#ifndef ARE_INCLUDE_CORE_BVH_H
#define ARE_INCLUDE_CORE_BVH_H

#include "basic/types.h"
#include "resource/buffer.h"
#include "scene/mesh.h"
#include <memory>
#include <vector>

namespace are {

// Axis-aligned bounding box
struct AABB {
	Vec3 min_;
	Vec3 max_;

	// Construct AABB from min and max points
	AABB(const Vec3 &min = Vec3(0.0f), const Vec3 &max = Vec3(0.0f))
		: min_(min), max_(max) {
	}

	// Expand AABB to include point
	void expand(const Vec3 &point);

	// Expand AABB to include another AABB
	void expand(const AABB &other);

	// Get center of AABB
	Vec3 center() const {
		return (min_ + max_) * 0.5f;
	}

	// Get surface area of AABB
	float surface_area() const;

	// Check if AABB is valid
	bool is_valid() const;
};

// Triangle primitive for BVH
struct Triangle {
	Vec3 v0_, v1_, v2_;
	Vec3 n0_, n1_, n2_;
	Vec2 uv0_, uv1_, uv2_;
	uint material_id_;

	// Get bounding box of triangle
	AABB get_bounds() const;

	// Get centroid of triangle
	Vec3 get_centroid() const;
};

// BVH node for GPU
struct BVHNode {
	Vec3 aabb_min_;
	uint left_first_; // Left child index or first primitive index
	Vec3 aabb_max_;
	uint count_; // 0 for interior node, >0 for leaf node
};

// GPU-friendly BVH node layout (std430 aligned)
struct BVHNodeGpu {
	Vec4 aabb_min_left_first_; ///< xyz = aabb min, w = left_first (uint)
	Vec4 aabb_max_count_; ///< xyz = aabb max, w = count (uint, 0 for interior)
};

// GPU-friendly triangle layout (std430 aligned)
struct TriangleGpu {
	Vec4 v0_material_; ///< xyz = v0, w = material_id (uint)
	Vec4 v1_; ///< xyz = v1, w = reserved
	Vec4 v2_; ///< xyz = v2, w = reserved
	Vec4 n0_; ///< xyz = n0, w = reserved
	Vec4 n1_; ///< xyz = n1, w = reserved
	Vec4 n2_; ///< xyz = n2, w = reserved
	Vec4 uv0_uv1_; ///< xy = uv0, zw = uv1
	Vec4 uv2_; ///< xy = uv2, zw = reserved
};

// Bounding Volume Hierarchy for ray tracing acceleration
class BVH {
public:
	// Constructor
	BVH();

	// Destructor
	~BVH();

	/*
	 * @brief Build BVH from meshes
	 * @param meshes Mesh list
	 * @return True if build succeeded
	 */
	bool build(const std::vector<std::shared_ptr<Mesh>> &meshes);

	/*
	 * @brief Upload BVH to GPU
	 * @param node_buffer Buffer for BVH nodes
	 * @param triangle_buffer Buffer for triangles
	 * @return True if upload succeeded
	 */
	bool upload_to_gpu(Buffer &node_buffer, Buffer &triangle_buffer);

	/*
	 * @brief Get total node count
	 * @return Node count
	 */
	uint get_node_count() const {
		return static_cast<uint>(nodes_.size());
	}

	/*
	 * @brief Get total triangle count
	 * @return Triangle count
	 */
	uint get_triangle_count() const {
		return static_cast<uint>(triangles_.size());
	}

	// Clear BVH data
	void clear();

private:
	std::vector<BVHNode> nodes_;
	std::vector<Triangle> triangles_;
	std::vector<uint> triangle_indices_;

	/*
	 * @brief Recursively build BVH
	 * @param node_idx Current node index
	 * @param first_prim First primitive index
	 * @param prim_count Primitive count
	 */
	void build_recursive_(uint node_idx, uint first_prim, uint prim_count);

	/*
	 * @brief Find best split using SAH
	 * @param first_prim First primitive index
	 * @param prim_count Primitive count
	 * @param axis Split axis (output)
	 * @param split_pos Split position (output)
	 * @return Split cost
	 */
	float find_best_split_(uint first_prim, uint prim_count, int &axis, float &split_pos);

	/*
	 * @brief Calculate node bounds
	 * @param first_prim First primitive index
	 * @param prim_count Primitive count
	 * @return Bounding box
	 */
	AABB calculate_bounds_(uint first_prim, uint prim_count);

	/*
	 * @brief Calculate centroid bounds
	 * @param first_prim First primitive index
	 * @param prim_count Primitive count
	 * @return Centroid bounding box
	 */
	AABB calculate_centroid_bounds_(uint first_prim, uint prim_count);
};

} // namespace are

#endif // ARE_INCLUDE_CORE_BVH_H