// SPDX-License-Identifier: MIT OR Apache-2.0
// Copyright (c) 2025 lacklustr@protonmail.com https://github.com/eadf

#[cfg(test)]
mod bvh_tests;
#[cfg(test)]
mod projection_tests;
use crate::common::VertexIndex;
use crate::common::sealed::ScalarType;
use crate::corner_table::TriangleIndex;
use vector_traits::num_traits::Float;
use vector_traits::prelude::*;

#[derive(Default)]
pub struct StaticBVH<S: ScalarType> {
    nodes: Vec<BVHNode<S>>,
    triangles: Vec<Triangle<S>>,
}

struct BVHNode<S: ScalarType> {
    bbox: Aabb<S>,
    data: NodeData,
}

enum NodeData {
    Leaf { start: u32, count: u32 },
    Internal { left: u32, right: u32 },
}

#[derive(Clone, Copy)]
struct Aabb<S: ScalarType> {
    min: S::Vec3,
    max: S::Vec3,
}

#[derive(Clone)]
pub struct Triangle<S: ScalarType> {
    v0: S::Vec3,
    v1: S::Vec3,
    v2: S::Vec3,
    normal: S::Vec3,
}

struct BestPoint<S: ScalarType> {
    point: S::Vec3,
    dist: S,
    triangle_index: u32,
}

impl<S: ScalarType> Aabb<S> {
    fn new(min: S::Vec3, max: S::Vec3) -> Self {
        Self { min, max }
    }

    fn from_triangle(tri: &Triangle<S>) -> Self {
        let min = tri.v0.min(tri.v1).min(tri.v2);
        let max = tri.v0.max(tri.v1).max(tri.v2);
        Self::new(min, max)
    }

    fn union(&self, other: &Aabb<S>) -> Aabb<S> {
        Aabb {
            min: self.min.min(other.min),
            max: self.max.max(other.max),
        }
    }

    fn distance_sq(&self, point: S::Vec3) -> S {
        let closest = point.clamp(self.min, self.max);
        point.distance_sq(closest)
    }

    #[allow(dead_code)]
    fn center(&self) -> S::Vec3 {
        (self.min + self.max) * 0.5.into()
    }
}

impl<S: ScalarType> StaticBVH<S> {
    pub fn new(indices: &[VertexIndex], vertices: &[S::Vec3]) -> Self {
        let triangles = indices
            .chunks_exact(3)
            .filter_map(|i| {
                // TODO: maybe convert to simd first?
                let v0 = vertices[i[0].0 as usize];
                let v1 = vertices[i[1].0 as usize];
                let v2 = vertices[i[2].0 as usize];
                let edge0 = v1 - v0;
                let edge1 = v2 - v0;
                let raw_normal = edge0.cross(edge1);

                let magnitude_sq = raw_normal.magnitude_sq();
                if magnitude_sq < 1e-10.into() {
                    return None; // Filter out degenerate triangles
                }

                Some(Triangle {
                    v0,
                    v1,
                    v2,
                    normal: raw_normal / magnitude_sq.sqrt(),
                })
            })
            .collect();
        Self::new_from_triangles(triangles)
    }

    pub fn new_from_triangles(triangles: Vec<Triangle<S>>) -> Self {
        let mut bvh = StaticBVH {
            nodes: Vec::with_capacity(2 * triangles.len()),
            triangles,
        };

        if !bvh.triangles.is_empty() {
            let _ = bvh.build_recursive(0, bvh.triangles.len() as u32);
        }

        bvh
    }

    fn build_recursive(&mut self, start: u32, end: u32) -> u32 {
        let count = end - start;
        let node_idx = self.nodes.len() as u32;

        // Compute bbox for this range
        let mut bbox = Aabb::from_triangle(&self.triangles[start as usize]);
        for i in (start + 1)..end {
            bbox = bbox.union(&Aabb::from_triangle(&self.triangles[i as usize]));
        }

        // Leaf node
        if count <= 4 {
            self.nodes.push(BVHNode {
                bbox,
                data: NodeData::Leaf { start, count },
            });
            return node_idx;
        }

        // Split along longest axis
        let extent = bbox.max - bbox.min;
        let axis: usize = if extent.x() > extent.y() && extent.x() > extent.z() {
            0
        } else if extent.y() > extent.z() {
            1
        } else {
            2
        };

        // Sort by centroid along axis
        let mid = (start + end) / 2;
        let _ = self.triangles[(start as usize)..(end as usize)].select_nth_unstable_by(
            (mid - start) as usize,
            |a, b| {
                let ca = (a.v0 + a.v1 + a.v2) / S::THREE;
                let cb = (b.v0 + b.v1 + b.v2) / S::THREE;
                ca[axis].partial_cmp(&cb[axis]).unwrap()
            },
        );

        // Placeholder node
        self.nodes.push(BVHNode {
            bbox,
            data: NodeData::Internal { left: 0, right: 0 },
        });

        // Build children
        let left = self.build_recursive(start, mid);
        let right = self.build_recursive(mid, end);

        // Update node
        self.nodes[node_idx as usize].data = NodeData::Internal { left, right };

        node_idx
    }

    /// Returns the closest point + the normal of the closest triangle
    ///
    /// # Parameters
    /// - `normal_threshold`: Cosine of maximum angle between query_normal and triangle normal
    ///   - `-1.0`: Accept all triangles (geometry preservation, backward faces OK)
    ///   - `0.0`: Accept triangles within 90° (prevent wrapping to opposite side)
    ///   - `0.5`: Accept triangles within 60° (smooth organic surfaces)
    ///   - `0.7`: Accept triangles within 45° (stricter smoothing)
    pub fn closest_point(
        &self,
        point: S::Vec3,
        query_normal: S::Vec3Simd,
        normal_threshold: S,
    ) -> Option<(S::Vec3, S::Vec3)> {
        if self.nodes.is_empty() {
            return None;
        }

        let mut best_point = BestPoint {
            point: S::Vec3::ZERO,
            dist: S::MAX,
            triangle_index: TriangleIndex::INVALID.0,
        };

        self.query_recursive(0, point, query_normal, normal_threshold, &mut best_point);

        if best_point.triangle_index != TriangleIndex::INVALID.0 {
            Some((
                best_point.point,
                self.triangles[best_point.triangle_index as usize].normal,
            ))
        } else {
            None
        }
    }

    fn query_recursive(
        &self,
        node_idx: u32,
        point: S::Vec3,
        query_normal: S::Vec3Simd,
        normal_threshold: S,
        best_point: &mut BestPoint<S>,
    ) {
        let node = &self.nodes[node_idx as usize];

        if node.bbox.distance_sq(point) >= best_point.dist {
            return;
        }

        match node.data {
            NodeData::Leaf { start, count } => {
                for i in start..(start + count) {
                    let tri = &self.triangles[i as usize];

                    // Filter based on normal orientation
                    if tri.normal.to_simd().dot(query_normal) < normal_threshold {
                        continue;
                    }

                    let proj = project_point_to_triangle(point, tri);
                    let dist = point.distance_sq(proj);

                    if dist < best_point.dist {
                        best_point.dist = dist;
                        best_point.point = proj;
                        best_point.triangle_index = i;
                    }
                }
            }
            NodeData::Internal { left, right } => {
                let dist_left = self.nodes[left as usize].bbox.distance_sq(point);
                let dist_right = self.nodes[right as usize].bbox.distance_sq(point);

                if dist_left < dist_right {
                    self.query_recursive(left, point, query_normal, normal_threshold, best_point);
                    self.query_recursive(right, point, query_normal, normal_threshold, best_point);
                } else {
                    self.query_recursive(right, point, query_normal, normal_threshold, best_point);
                    self.query_recursive(left, point, query_normal, normal_threshold, best_point);
                }
            }
        }
    }
}

/// Returns the closest point to `p` on the triangle
fn project_point_to_triangle<S: ScalarType>(p: S::Vec3, tri: &Triangle<S>) -> S::Vec3 {
    let p = p.to_simd();
    let v0 = tri.v0.to_simd();
    let v1 = tri.v1.to_simd();
    let v2 = tri.v2.to_simd();

    // Compute edges
    let edge0 = v1 - v0;
    let edge1 = v2 - v0;

    let normal = tri.normal.to_simd();

    // Project point onto triangle plane
    let v0_to_p = p - v0;
    let dist_to_plane = v0_to_p.dot(normal);
    let p_on_plane = p - normal * dist_to_plane;

    // Check if projection is inside triangle using barycentric coordinates
    let v0_to_proj = p_on_plane - v0;

    let d00 = edge0.dot(edge0);
    let d01 = edge0.dot(edge1);
    let d11 = edge1.dot(edge1);
    let d20 = v0_to_proj.dot(edge0);
    let d21 = v0_to_proj.dot(edge1);

    let denom = d00 * d11 - d01 * d01;

    if Float::abs(denom) < 1e-10.into() {
        unreachable!()
    }

    let v = (d11 * d20 - d01 * d21) / denom;
    let w = (d00 * d21 - d01 * d20) / denom;
    let u = S::ONE - v - w;

    // If inside triangle, return projection
    if u >= S::ZERO && v >= S::ZERO && w >= S::ZERO {
        return S::Vec3::from_simd(p_on_plane);
    }

    // Point is outside triangle, find the closest point on edges/vertices

    #[inline(always)]
    // Helper to clamp point to edge
    fn closest_point_on_edge<S: ScalarType>(
        p: S::Vec3Simd,
        a: S::Vec3Simd,
        b: S::Vec3Simd,
    ) -> S::Vec3Simd {
        let ab = b - a;
        let ap = p - a;
        let t = Float::clamp(ap.dot(ab) / ab.magnitude_sq(), S::ZERO, S::ONE);
        a + ab * t
    }

    // Check all three edges
    let c0 = closest_point_on_edge::<S>(p, v0, v1);
    let c1 = closest_point_on_edge::<S>(p, v1, v2);
    let c2 = closest_point_on_edge::<S>(p, v2, v0);

    let d0 = p.distance_sq(c0);
    let d1 = p.distance_sq(c1);
    let d2 = p.distance_sq(c2);

    if d0 <= d1 && d0 <= d2 {
        S::Vec3::from_simd(c0)
    } else if d1 <= d2 {
        S::Vec3::from_simd(c1)
    } else {
        S::Vec3::from_simd(c2)
    }
}
