use ray_tracing_core::{aabb::AABB, prelude::*};

pub type Index = u32;

pub trait BvhIntersection {
    fn get_t(&self) -> Float;
}

pub trait BvhData {
    type Intersect<'a>: BvhIntersection
    where
        Self: 'a;

    // fn get_slice(&mut self) -> &mut [Self::Sortable];

    fn len(&self) -> Index;

    fn sort_range_dim(&mut self, id: std::ops::Range<Index>, dim: u8);

    fn get_aabb(&self, id: Index) -> AABB;

    fn get_aabb_range(&self, id: std::ops::Range<Index>) -> AABB {
        id.map(|i| self.get_aabb(i))
            .reduce(|a, b| a.extend(b.min()).extend(b.max()))
            .unwrap()
    }

    fn intersect(&self, id: Index, ray: Ray, min: Float, max: Float)
    -> Option<Self::Intersect<'_>>;
}

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

#[derive(Debug)]
pub struct Bvh<T> {
    data: T,
    bvh: Vec<Node>,
}

impl<T: BvhData> Bvh<T> {
    pub fn new(mut data: T, min: Index) -> Self {
        let len = data.len();
        let mut bvh = vec![Node::Leaf {
            start: 0,
            count: len,
        }];

        let aabb = data.get_aabb_range(0..len);

        Self::build_bvh(&mut data, &mut bvh, 0, aabb, min);

        Self { data, bvh }
    }

    pub fn inner(&self) -> &T {
        &self.data
    }

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

        if count <= Index::max(1, min) {
            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
        };

        // sort
        data.sort_range_dim(start..start + count, dim);

        let mid = count / 2;

        let left_aabb = data.get_aabb_range(start..(start + mid));
        let right_aabb = data.get_aabb_range((start + mid)..(start + count));

        assert!(aabb.contains_aabb(left_aabb));
        assert!(aabb.contains_aabb(right_aabb));

        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: mid });
        bvh.push(Node::Leaf {
            start: start + mid,
            count: count - mid,
        });

        Self::build_bvh(data, bvh, i as usize, left_aabb, min);
        Self::build_bvh(data, bvh, i as usize + 1, right_aabb, min);
    }

    fn intersect_inner(
        &self,
        node: Index,
        ray: Ray,
        min: Float,
        max: Float,
    ) -> Option<T::Intersect<'_>> {
        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_inner(right, ray, min, max),
                    (Some(_), None) => self.intersect_inner(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_inner(close, ray, min, max) {
                            if let Some(far_intersect) = self
                                .intersect_inner(
                                    far,
                                    ray,
                                    min,
                                    Float::min(max, close_intersect.get_t()),
                                )
                                .filter(|far_intersect| {
                                    far_intersect.get_t() < close_intersect.get_t()
                                })
                            {
                                Some(far_intersect)
                            } else {
                                Some(close_intersect)
                            }
                        } else {
                            self.intersect_inner(far, ray, min, max)
                        }
                    }
                }
            }
            Node::Leaf { start, count } => {
                let mut intersection = None;
                for i in start..(start + count) {
                    if let Some(t) = self.data.intersect(i, ray, min, max)
                        && min <= t.get_t()
                        && t.get_t() <= max
                        && !intersection.as_ref().is_some_and(
                            |old_t: &<T as BvhData>::Intersect<'_>| t.get_t() > old_t.get_t(),
                        )
                    {
                        intersection.replace(t);
                    }
                }

                intersection
            }
        }
    }

    pub fn intersect(&self, ray: Ray, min: Float, max: Float) -> Option<T::Intersect<'_>> {
        self.intersect_inner(0, ray, min, max)
    }
}
