use crate::mesh::{AABB, Triangle, TriangleMesh, Vec3};

/// BVH over a triangle mesh for fast nearest-point queries.
pub struct BVH {
    nodes: Vec<BVHNode>,
    tri_indices: Vec<usize>,
}

enum BVHNode {
    Leaf {
        bounds: AABB,
        first: usize,
        count: usize,
    },
    Internal {
        bounds: AABB,
        left: usize,
        right: usize,
    },
}

impl BVH {
    pub fn build(mesh: &TriangleMesh) -> Self {
        let n = mesh.triangles.len();
        let mut tri_indices: Vec<usize> = (0..n).collect();
        let mut centroids: Vec<Vec3> = mesh.triangles.iter().map(|t| t.centroid()).collect();
        let mut nodes = Vec::with_capacity(2 * n);

        build_recursive(
            &mesh.triangles,
            &mut tri_indices,
            &mut centroids,
            &mut nodes,
            0,
            n,
        );

        Self { nodes, tri_indices }
    }

    /// Find the signed distance from point p to the mesh.
    /// Sign is determined by the normal of the nearest triangle.
    pub fn signed_distance(&self, mesh: &TriangleMesh, p: Vec3) -> f64 {
        let mut best_dist_sq = f64::INFINITY;
        let mut best_sign = 1.0f64;
        self.query_nearest(mesh, p, 0, &mut best_dist_sq, &mut best_sign);
        best_dist_sq.sqrt() * best_sign
    }

    /// Find the nearest triangle index and unsigned distance.
    pub fn nearest_triangle(&self, mesh: &TriangleMesh, p: Vec3) -> (usize, f64) {
        let mut best_dist_sq = f64::INFINITY;
        let mut best_idx = 0;
        self.query_nearest_idx(mesh, p, 0, &mut best_dist_sq, &mut best_idx);
        (best_idx, best_dist_sq.sqrt())
    }

    fn query_nearest(
        &self,
        mesh: &TriangleMesh,
        p: Vec3,
        node_idx: usize,
        best_dist_sq: &mut f64,
        best_sign: &mut f64,
    ) {
        match &self.nodes[node_idx] {
            BVHNode::Leaf { bounds, first, count } => {
                if bounds.distance_to_point(p).powi(2) > *best_dist_sq {
                    return;
                }
                for i in *first..(*first + *count) {
                    let tri_idx = self.tri_indices[i];
                    let tri = &mesh.triangles[tri_idx];
                    let (closest, dist) = tri.closest_point(p);
                    let dist_sq = dist * dist;
                    if dist_sq < *best_dist_sq {
                        *best_dist_sq = dist_sq;
                        let to_point = p - closest;
                        let normal = tri.normal();
                        *best_sign = if to_point.dot(normal) >= 0.0 { 1.0 } else { -1.0 };
                    }
                }
            }
            BVHNode::Internal { bounds, left, right } => {
                if bounds.distance_to_point(p).powi(2) > *best_dist_sq {
                    return;
                }
                let left_bounds = node_bounds(&self.nodes[*left]);
                let right_bounds = node_bounds(&self.nodes[*right]);
                let dl = left_bounds.distance_to_point(p);
                let dr = right_bounds.distance_to_point(p);

                if dl < dr {
                    self.query_nearest(mesh, p, *left, best_dist_sq, best_sign);
                    self.query_nearest(mesh, p, *right, best_dist_sq, best_sign);
                } else {
                    self.query_nearest(mesh, p, *right, best_dist_sq, best_sign);
                    self.query_nearest(mesh, p, *left, best_dist_sq, best_sign);
                }
            }
        }
    }

    fn query_nearest_idx(
        &self,
        mesh: &TriangleMesh,
        p: Vec3,
        node_idx: usize,
        best_dist_sq: &mut f64,
        best_idx: &mut usize,
    ) {
        match &self.nodes[node_idx] {
            BVHNode::Leaf { bounds, first, count } => {
                if bounds.distance_to_point(p).powi(2) > *best_dist_sq {
                    return;
                }
                for i in *first..(*first + *count) {
                    let tri_idx = self.tri_indices[i];
                    let tri = &mesh.triangles[tri_idx];
                    let (_, dist) = tri.closest_point(p);
                    let dist_sq = dist * dist;
                    if dist_sq < *best_dist_sq {
                        *best_dist_sq = dist_sq;
                        *best_idx = tri_idx;
                    }
                }
            }
            BVHNode::Internal { bounds, left, right } => {
                if bounds.distance_to_point(p).powi(2) > *best_dist_sq {
                    return;
                }
                let left_bounds = node_bounds(&self.nodes[*left]);
                let right_bounds = node_bounds(&self.nodes[*right]);
                let dl = left_bounds.distance_to_point(p);
                let dr = right_bounds.distance_to_point(p);

                if dl < dr {
                    self.query_nearest_idx(mesh, p, *left, best_dist_sq, best_idx);
                    self.query_nearest_idx(mesh, p, *right, best_dist_sq, best_idx);
                } else {
                    self.query_nearest_idx(mesh, p, *right, best_dist_sq, best_idx);
                    self.query_nearest_idx(mesh, p, *left, best_dist_sq, best_idx);
                }
            }
        }
    }

    /// Count triangles whose centroid falls within a given AABB.
    pub fn count_in_region(&self, mesh: &TriangleMesh, region: &AABB) -> usize {
        self.count_recursive(mesh, region, 0)
    }

    fn count_recursive(&self, mesh: &TriangleMesh, region: &AABB, node_idx: usize) -> usize {
        match &self.nodes[node_idx] {
            BVHNode::Leaf { bounds, first, count } => {
                if !aabb_overlaps(bounds, region) {
                    return 0;
                }
                let mut n = 0;
                for i in *first..(*first + *count) {
                    let c = mesh.triangles[self.tri_indices[i]].centroid();
                    if point_in_aabb(c, region) {
                        n += 1;
                    }
                }
                n
            }
            BVHNode::Internal { bounds, left, right } => {
                if !aabb_overlaps(bounds, region) {
                    return 0;
                }
                self.count_recursive(mesh, region, *left)
                    + self.count_recursive(mesh, region, *right)
            }
        }
    }
}

fn node_bounds(node: &BVHNode) -> &AABB {
    match node {
        BVHNode::Leaf { bounds, .. } | BVHNode::Internal { bounds, .. } => bounds,
    }
}

fn aabb_overlaps(a: &AABB, b: &AABB) -> bool {
    a.min.x <= b.max.x && a.max.x >= b.min.x
        && a.min.y <= b.max.y && a.max.y >= b.min.y
        && a.min.z <= b.max.z && a.max.z >= b.min.z
}

fn point_in_aabb(p: Vec3, b: &AABB) -> bool {
    p.x >= b.min.x && p.x <= b.max.x
        && p.y >= b.min.y && p.y <= b.max.y
        && p.z >= b.min.z && p.z <= b.max.z
}

const MAX_LEAF_SIZE: usize = 8;

fn build_recursive(
    triangles: &[Triangle],
    indices: &mut [usize],
    centroids: &mut [Vec3],
    nodes: &mut Vec<BVHNode>,
    start: usize,
    end: usize,
) -> usize {
    let count = end - start;

    // Compute bounds
    let mut bounds = AABB::empty();
    for &idx in &indices[start..end] {
        bounds = bounds.union(&AABB::from_triangle(&triangles[idx]));
    }

    if count <= MAX_LEAF_SIZE {
        let node_idx = nodes.len();
        nodes.push(BVHNode::Leaf { bounds, first: start, count });
        return node_idx;
    }

    // Split along longest axis at centroid median
    let axis = bounds.longest_axis();
    let mid = start + count / 2;

    // Partial sort: partition around the median centroid on the chosen axis
    let get_axis = |v: Vec3| match axis {
        0 => v.x,
        1 => v.y,
        _ => v.z,
    };

    // Simple partition: sort the range by centroid along axis
    let index_slice = &mut indices[start..end];
    index_slice.sort_unstable_by(|&a, &b| {
        let ca = get_axis(centroids[a]);
        let cb = get_axis(centroids[b]);
        ca.partial_cmp(&cb).unwrap_or(std::cmp::Ordering::Equal)
    });

    let node_idx = nodes.len();
    nodes.push(BVHNode::Leaf { bounds: AABB::empty(), first: 0, count: 0 }); // placeholder

    let left = build_recursive(triangles, indices, centroids, nodes, start, mid);
    let right = build_recursive(triangles, indices, centroids, nodes, mid, end);

    nodes[node_idx] = BVHNode::Internal { bounds, left, right };
    node_idx
}
