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 angle-weighted pseudonormal voting across all
    /// triangles within a tolerance of the nearest distance. This handles
    /// edge and vertex proximity correctly, where a single triangle's face
    /// normal would give ambiguous or wrong sign.
    pub fn signed_distance(&self, mesh: &TriangleMesh, p: Vec3) -> f64 {
        let mut best_dist_sq = f64::INFINITY;
        self.query_nearest_dist(mesh, p, 0, &mut best_dist_sq);

        if best_dist_sq == f64::INFINITY {
            return f64::INFINITY;
        }

        // Collect all triangles within a relative tolerance of the best distance.
        // This catches co-nearest triangles sharing an edge or vertex.
        let eps_factor = 1.0 + 1e-4;
        let threshold_sq = best_dist_sq * eps_factor * eps_factor + 1e-20;
        let mut neighbors: Vec<(usize, Vec3)> = Vec::new(); // (tri_idx, closest_point)
        self.collect_near(mesh, p, 0, threshold_sq, &mut neighbors);

        let sign = if neighbors.len() <= 1 {
            // Single closest triangle: use its face normal directly
            if let Some(&(tri_idx, closest)) = neighbors.first() {
                let to_point = p - closest;
                let normal = mesh.triangles[tri_idx].normal();
                if to_point.dot(normal) >= 0.0 { 1.0 } else { -1.0 }
            } else {
                1.0
            }
        } else {
            // Multiple co-nearest triangles: angle-weighted normal vote.
            // Weight each triangle's contribution by the solid angle it subtends
            // at p, approximated by the triangle area divided by distance squared.
            let mut weighted_dot = 0.0f64;
            for &(tri_idx, closest) in &neighbors {
                let tri = &mesh.triangles[tri_idx];
                let to_point = p - closest;
                let normal = tri.normal();
                let area = tri.area();
                let d_sq = to_point.dot(to_point);
                // Weight: area / (distance^2 + epsilon) — larger, closer faces dominate
                let weight = area / (d_sq + 1e-20);
                weighted_dot += to_point.dot(normal) * weight;
            }
            if weighted_dot >= 0.0 { 1.0 } else { -1.0 }
        };

        best_dist_sq.sqrt() * 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())
    }

    /// Find the minimum squared distance from p to any triangle.
    fn query_nearest_dist(
        &self,
        mesh: &TriangleMesh,
        p: Vec3,
        node_idx: usize,
        best_dist_sq: &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 (_, dist) = tri.closest_point(p);
                    let dist_sq = dist * dist;
                    if dist_sq < *best_dist_sq {
                        *best_dist_sq = dist_sq;
                    }
                }
            }
            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_dist(mesh, p, *left, best_dist_sq);
                    self.query_nearest_dist(mesh, p, *right, best_dist_sq);
                } else {
                    self.query_nearest_dist(mesh, p, *right, best_dist_sq);
                    self.query_nearest_dist(mesh, p, *left, best_dist_sq);
                }
            }
        }
    }

    /// Collect all (triangle_index, closest_point) pairs within threshold_sq of p.
    fn collect_near(
        &self,
        mesh: &TriangleMesh,
        p: Vec3,
        node_idx: usize,
        threshold_sq: f64,
        out: &mut Vec<(usize, Vec3)>,
    ) {
        match &self.nodes[node_idx] {
            BVHNode::Leaf { bounds, first, count } => {
                if bounds.distance_to_point(p).powi(2) > threshold_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);
                    if dist * dist <= threshold_sq {
                        out.push((tri_idx, closest));
                    }
                }
            }
            BVHNode::Internal { bounds, left, right } => {
                if bounds.distance_to_point(p).powi(2) > threshold_sq {
                    return;
                }
                self.collect_near(mesh, p, *left, threshold_sq, out);
                self.collect_near(mesh, p, *right, threshold_sq, out);
            }
        }
    }

    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;
const SAH_NUM_BINS: usize = 16;
const SAH_TRAVERSAL_COST: f64 = 1.0;
const SAH_INTERSECT_COST: f64 = 1.5;

/// Surface area of an AABB (used as probability proxy in SAH).
fn aabb_surface_area(b: &AABB) -> f64 {
    let e = b.extent();
    2.0 * (e.x * e.y + e.y * e.z + e.z * e.x)
}

/// SAH bin for accumulating triangle bounds during split evaluation.
struct SAHBin {
    bounds: AABB,
    count: usize,
}

impl SAHBin {
    fn empty() -> Self {
        Self { bounds: AABB::empty(), count: 0 }
    }
}

/// Evaluate SAH cost for splitting at each bin boundary along the given axis.
/// Returns (best_split_bin, best_cost). Split puts bins [0..split] left, [split..N] right.
fn sah_best_split(
    triangles: &[Triangle],
    indices: &[usize],
    centroids: &[Vec3],
    start: usize,
    end: usize,
    parent_bounds: &AABB,
    axis: usize,
) -> (usize, f64) {
    let get_axis = |v: Vec3| match axis {
        0 => v.x,
        1 => v.y,
        _ => v.z,
    };

    // Centroid bounds (tighter than triangle bounds) for bin mapping
    let mut centroid_min = f64::INFINITY;
    let mut centroid_max = f64::NEG_INFINITY;
    for &idx in &indices[start..end] {
        let c = get_axis(centroids[idx]);
        centroid_min = centroid_min.min(c);
        centroid_max = centroid_max.max(c);
    }

    let span = centroid_max - centroid_min;
    if span < 1e-12 {
        // All centroids coincide on this axis; can't split here
        return (SAH_NUM_BINS / 2, f64::INFINITY);
    }
    let inv_span = SAH_NUM_BINS as f64 / span;

    // Bin triangles
    let mut bins: Vec<SAHBin> = (0..SAH_NUM_BINS).map(|_| SAHBin::empty()).collect();

    for &idx in &indices[start..end] {
        let c = get_axis(centroids[idx]);
        let b = ((c - centroid_min) * inv_span) as usize;
        let b = b.min(SAH_NUM_BINS - 1);
        bins[b].bounds = bins[b].bounds.union(&AABB::from_triangle(&triangles[idx]));
        bins[b].count += 1;
    }

    // Sweep from left: accumulate bounds and counts for splits [0..i+1] | [i+1..N]
    let mut left_area = vec![0.0f64; SAH_NUM_BINS - 1];
    let mut left_count = vec![0usize; SAH_NUM_BINS - 1];
    let mut running_bounds = AABB::empty();
    let mut running_count = 0usize;

    for i in 0..(SAH_NUM_BINS - 1) {
        running_bounds = running_bounds.union(&bins[i].bounds);
        running_count += bins[i].count;
        left_area[i] = aabb_surface_area(&running_bounds);
        left_count[i] = running_count;
    }

    // Sweep from right
    let mut right_area = vec![0.0f64; SAH_NUM_BINS - 1];
    let mut right_count = vec![0usize; SAH_NUM_BINS - 1];
    running_bounds = AABB::empty();
    running_count = 0;

    for i in (0..(SAH_NUM_BINS - 1)).rev() {
        running_bounds = running_bounds.union(&bins[i + 1].bounds);
        running_count += bins[i + 1].count;
        right_area[i] = aabb_surface_area(&running_bounds);
        right_count[i] = running_count;
    }

    let parent_area = aabb_surface_area(parent_bounds);
    let inv_parent = if parent_area > 1e-12 { 1.0 / parent_area } else { 1.0 };

    let mut best_cost = f64::INFINITY;
    let mut best_bin = SAH_NUM_BINS / 2;

    for i in 0..(SAH_NUM_BINS - 1) {
        if left_count[i] == 0 || right_count[i] == 0 {
            continue;
        }
        let cost = SAH_TRAVERSAL_COST
            + SAH_INTERSECT_COST * inv_parent
                * (left_area[i] * left_count[i] as f64
                    + right_area[i] * right_count[i] as f64);
        if cost < best_cost {
            best_cost = cost;
            best_bin = i;
        }
    }

    (best_bin, best_cost)
}

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

    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;
    }

    // SAH split: evaluate all three axes, pick lowest cost
    let leaf_cost = SAH_INTERSECT_COST * count as f64;
    let mut best_axis = 0usize;
    let mut best_bin = SAH_NUM_BINS / 2;
    let mut best_cost = f64::INFINITY;

    for axis in 0..3 {
        let (bin, cost) = sah_best_split(
            triangles, indices, centroids, start, end, &bounds, axis,
        );
        if cost < best_cost {
            best_cost = cost;
            best_bin = bin;
            best_axis = axis;
        }
    }

    // If SAH says a leaf is cheaper, or no valid split found, fall back to median
    let use_median = best_cost >= leaf_cost || best_cost == f64::INFINITY;

    let mid = if use_median {
        // Median split along longest axis
        let axis = bounds.longest_axis();
        let get_axis = |v: Vec3| match axis {
            0 => v.x,
            1 => v.y,
            _ => v.z,
        };
        indices[start..end].sort_unstable_by(|&a, &b| {
            get_axis(centroids[a])
                .partial_cmp(&get_axis(centroids[b]))
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        start + count / 2
    } else {
        // Partition around the SAH split bin
        let get_axis = |v: Vec3| match best_axis {
            0 => v.x,
            1 => v.y,
            _ => v.z,
        };

        let mut centroid_min = f64::INFINITY;
        let mut centroid_max = f64::NEG_INFINITY;
        for &idx in &indices[start..end] {
            let c = get_axis(centroids[idx]);
            centroid_min = centroid_min.min(c);
            centroid_max = centroid_max.max(c);
        }
        let span = centroid_max - centroid_min;
        let inv_span = SAH_NUM_BINS as f64 / span;

        // Partition: triangles in bins [0..=best_bin] go left
        let slice = &mut indices[start..end];
        let mut lo = 0usize;
        let mut hi = slice.len();
        while lo < hi {
            let c = get_axis(centroids[slice[lo]]);
            let b = ((c - centroid_min) * inv_span) as usize;
            let b = b.min(SAH_NUM_BINS - 1);
            if b <= best_bin {
                lo += 1;
            } else {
                hi -= 1;
                slice.swap(lo, hi);
            }
        }

        let mid = start + lo;
        // Guard: if partition degenerates, fall back to midpoint
        if mid == start || mid == end {
            start + count / 2
        } else {
            mid
        }
    };

    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
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_unit_cube_mesh() -> TriangleMesh {
        // 12 triangles forming a unit cube centered at origin
        let h = 0.5;
        let corners = [
            Vec3::new(-h, -h, -h), Vec3::new( h, -h, -h),
            Vec3::new( h,  h, -h), Vec3::new(-h,  h, -h),
            Vec3::new(-h, -h,  h), Vec3::new( h, -h,  h),
            Vec3::new( h,  h,  h), Vec3::new(-h,  h,  h),
        ];
        // CCW winding viewed from outside each face
        let faces: [(usize, usize, usize); 12] = [
            (0,2,1), (0,3,2), // -Z face: normal points -Z
            (4,5,6), (4,6,7), // +Z face: normal points +Z
            (0,1,5), (0,5,4), // -Y face: normal points -Y
            (2,3,7), (2,7,6), // +Y face: normal points +Y
            (0,4,7), (0,7,3), // -X face: normal points -X
            (1,2,6), (1,6,5), // +X face: normal points +X
        ];
        let mut triangles = Vec::new();
        let mut bounds = AABB::empty();
        for &(a, b, c) in &faces {
            let tri = Triangle { v: [corners[a], corners[b], corners[c]] };
            for v in &tri.v {
                bounds.min = bounds.min.component_min(*v);
                bounds.max = bounds.max.component_max(*v);
            }
            triangles.push(tri);
        }
        TriangleMesh { triangles, bounds }
    }

    #[test]
    fn bvh_builds_for_cube() {
        let mesh = make_unit_cube_mesh();
        let bvh = BVH::build(&mesh);
        assert!(!bvh.nodes.is_empty());
    }

    #[test]
    fn signed_distance_inside_cube() {
        let mesh = make_unit_cube_mesh();
        let bvh = BVH::build(&mesh);
        let d = bvh.signed_distance(&mesh, Vec3::new(0.0, 0.0, 0.0));
        assert!(d < 0.0, "origin should be inside cube, got {d}");
    }

    #[test]
    fn signed_distance_outside_cube() {
        let mesh = make_unit_cube_mesh();
        let bvh = BVH::build(&mesh);
        let d = bvh.signed_distance(&mesh, Vec3::new(2.0, 0.0, 0.0));
        assert!(d > 0.0, "point far from cube should be outside, got {d}");
    }

    #[test]
    fn signed_distance_on_edge_correct_sign() {
        // Point near a cube edge: the multi-triangle voting should give
        // consistent sign even though two face normals are involved.
        let mesh = make_unit_cube_mesh();
        let bvh = BVH::build(&mesh);
        // Slightly outside the cube near the +X/+Y edge
        let d = bvh.signed_distance(&mesh, Vec3::new(0.6, 0.6, 0.0));
        assert!(d > 0.0, "point outside near edge should be positive, got {d}");
        // Slightly inside near the same edge
        let d2 = bvh.signed_distance(&mesh, Vec3::new(0.4, 0.4, 0.0));
        assert!(d2 < 0.0, "point inside near edge should be negative, got {d2}");
    }

    #[test]
    fn signed_distance_on_vertex_correct_sign() {
        // Point near a cube vertex: three face normals are involved.
        let mesh = make_unit_cube_mesh();
        let bvh = BVH::build(&mesh);
        let d = bvh.signed_distance(&mesh, Vec3::new(0.7, 0.7, 0.7));
        assert!(d > 0.0, "point outside near vertex should be positive, got {d}");
    }

    #[test]
    fn nearest_triangle_finds_closest() {
        let mesh = make_unit_cube_mesh();
        let bvh = BVH::build(&mesh);
        let (_, dist) = bvh.nearest_triangle(&mesh, Vec3::new(1.5, 0.0, 0.0));
        assert!((dist - 1.0).abs() < 0.01, "distance to +X face should be ~1.0, got {dist}");
    }
}