use ray_tracing_core::prelude::*;

use crate::util::triangle_normal;

use super::AccelerationStructure;

type Index = u32;

#[derive(Debug, Clone)]
pub struct TriangleBVH<T> {
    vertices: Vec<Pos3>,
    triangles: Vec<Triangle<T>>,
    bvh: Vec<Node>,
}

#[derive(Debug, Clone, Copy)]
enum Node {
    Inner {
        left: Index,
        left_aabb: AABB,
        right: Index,
        right_aabb: AABB,
    },
    Leaf {
        start: Index,
        count: Index,
    },
}

#[derive(Debug, Clone, Copy)]
pub struct Triangle<T> {
    pub vertices: [Index; 3],
    pub data: T,
}

impl<T> Triangle<T> {
    pub fn new(vertices: [Index; 3], data: T) -> Self {
        Self { vertices, data }
    }
}

impl<T: Copy> TriangleBVH<T> {
    pub fn new(vertices: Vec<Pos3>, mut triangles: Vec<Triangle<T>>) -> Self {
        let mut bvh = vec![Node::Leaf {
            start: 0,
            count: triangles.len() as Index,
        }];

        let aabb = Self::calculate_aabb(&vertices, &triangles);

        Self::build_bvh(&vertices, &mut triangles, &mut bvh, 0, aabb);

        Self {
            vertices,
            triangles,
            bvh,
        }
    }

    fn calculate_aabb(vertices: &[Pos3], triangles: &[Triangle<T>]) -> AABB {
        let mut aabb = AABB::new(
            vertices[triangles[0].vertices[0] as usize],
            vertices[triangles[0].vertices[1] as usize],
        );
        aabb = aabb.extend(vertices[triangles[0].vertices[2] as usize]);

        for t in triangles.iter().skip(1) {
            for v in t.vertices {
                aabb = aabb.extend(vertices[v as usize]);
            }
        }

        aabb
    }

    fn build_bvh(
        vertices: &[Pos3],
        triangles: &mut [Triangle<T>],
        bvh: &mut Vec<Node>,
        node: usize,
        aabb: AABB,
    ) {
        let (start, count) = if let Node::Leaf { start, count } = bvh[node] {
            (start, count)
        } else {
            unreachable!()
        };

        if count < 8 {
            return;
        }

        let size = aabb.size();

        let dim = if size.x() > Float::max(size.y(), size.z()) {
            0
        } else if size.y() > Float::max(size.x(), size.z()) {
            1
        } else {
            2
        };

        let get_key = |t: &Triangle<T>| {
            t.vertices
                .iter()
                .map(|&v| vertices[v as usize][dim])
                .sum::<Float>()
                / 3.0
        };

        triangles.sort_by(|a, b| get_key(a).partial_cmp(&get_key(b)).unwrap());

        let (left_range, right_range) = triangles.split_at_mut((count / 2) as usize);

        let left_aabb = Self::calculate_aabb(vertices, left_range);
        let right_aabb = Self::calculate_aabb(vertices, right_range);

        let i = bvh.len() as Index;
        bvh[node] = Node::Inner {
            left: i,
            left_aabb,
            right: i + 1,
            right_aabb,
        };

        bvh.push(Node::Leaf {
            start,
            count: left_range.len() as Index,
        });
        bvh.push(Node::Leaf {
            start: start + left_range.len() as Index,
            count: right_range.len() as Index,
        });

        Self::build_bvh(vertices, left_range, bvh, i as usize, left_aabb);
        Self::build_bvh(vertices, right_range, bvh, i as usize + 1, right_aabb);
    }

    fn intersect_bvh(
        &self,
        node: Index,
        ray: Ray,
        min: Float,
        max: Float,
    ) -> Option<(Float, Dir3, T)> {
        match self.bvh[node as usize] {
            Node::Inner {
                left,
                left_aabb,
                right,
                right_aabb,
            } => {
                let left_intersect = left_aabb.intersect_ray(ray, min, max);
                let right_intersect = right_aabb.intersect_ray(ray, min, max);

                match (left_intersect, right_intersect) {
                    (None, None) => None,
                    (None, Some(_)) => self.intersect_bvh(right, ray, min, max),
                    (Some(_), None) => self.intersect_bvh(left, ray, min, max),
                    (Some(l), Some(r)) => {
                        let close;
                        let far;
                        if l < r {
                            close = left;
                            far = right;
                        } else {
                            close = right;
                            far = left;
                        }

                        if let Some(close_intersect) = self.intersect_bvh(close, ray, min, max) {
                            if let Some(far_intersect) = self
                                .intersect_bvh(far, ray, min, Float::min(max, close_intersect.0))
                                .filter(|far_intersect| far_intersect.0 < close_intersect.0)
                            {
                                Some(far_intersect)
                            } else {
                                Some(close_intersect)
                            }
                        } else {
                            self.intersect_bvh(far, ray, min, max)
                        }
                    }
                }
            }
            Node::Leaf { start, count } => {
                let mut intersection = None;
                for i in start..(start + count) {
                    if let Some(t) = crate::util::triangle_intersection(
                        ray,
                        self.get_vertices(&self.triangles[i as usize]),
                    ) {
                        if min <= t
                            && t <= max
                            && !intersection.is_some_and(|(old_t, _, _)| t > old_t)
                        {
                            let n = triangle_normal(self.get_vertices(&self.triangles[i as usize]))
                                .normalize();
                            intersection.replace((t, n, self.triangles[i as usize].data));
                        }
                    }
                }

                intersection
            }
        }
    }

    fn get_vertices(&self, triangle: &Triangle<T>) -> [Pos3; 3] {
        [
            self.vertices[triangle.vertices[0] as usize],
            self.vertices[triangle.vertices[1] as usize],
            self.vertices[triangle.vertices[2] as usize],
        ]
    }
}

impl<T: Copy> AccelerationStructure<T> for TriangleBVH<T> {
    fn intersect(&self, ray: Ray, min: Float, max: Float) -> Option<(Float, Dir3, T)> {
        self.intersect_bvh(0, ray, min, max)
    }
}
