use crate::core::geometry::{Bounds3f, Ray, Vector3f};
use crate::core::primitive::{Primitive, PrimitiveTrait};
use crate::core::shape::ShapeIntersection;
use crate::{Float, Ptr, GVec, gvec};

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct LinearBVHNode {
    pub bounds: Bounds3f,
    pub primitives_offset: usize,
    pub n_primitives: u16,
    pub axis: u8,
    pub pad: u8,
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct BVHAggregate {
    pub node_count: u32,
    pub max_prims_in_node: u32,
    pub primitive_count: u32,
    pub primitives: GVec<Primitive>,
    pub nodes: GVec<LinearBVHNode>,
}

impl BVHAggregate {
    pub fn empty() -> Self {
        Self {
            max_prims_in_node: 0,
            primitives: gvec(),
            primitive_count: 0,
            nodes: gvec(),
            node_count: 0,
        }
    }

    #[inline(always)]
    fn node(&self, i: usize) -> &LinearBVHNode {
        unsafe { self.nodes.get_unchecked(i) }
    }

    #[inline(always)]
    fn primitive(&self, i: usize) -> &Primitive {
        unsafe { self.primitives.get_unchecked(i) }
    }
}

impl PrimitiveTrait for BVHAggregate {
    fn bounds(&self) -> Bounds3f {
        if self.nodes.is_empty() || self.node_count == 0 {
            Bounds3f::default()
        } else {
            self.node(0).bounds
        }
    }

    fn intersect(&self, r: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {
        if self.nodes.is_empty() {
            return None;
        }

        let mut hit_t = t_max.unwrap_or(Float::INFINITY);
        let mut best_si: Option<ShapeIntersection> = None;

        let inv_dir = Vector3f::new(1.0 / r.d.x(), 1.0 / r.d.y(), 1.0 / r.d.z());
        let dir_is_neg = [
            if inv_dir.x() < 0.0 { 1 } else { 0 },
            if inv_dir.y() < 0.0 { 1 } else { 0 },
            if inv_dir.z() < 0.0 { 1 } else { 0 },
        ];

        let mut stack = [0usize; 64];
        let mut stack_ptr = 0;
        let mut node_idx = 0usize;

        loop {
            let node = self.node(node_idx);

            if node
                .bounds
                .intersect_p(r.o, hit_t, inv_dir, &dir_is_neg)
                .is_none()
            {
                if stack_ptr == 0 {
                    break;
                }
                stack_ptr -= 1;
                node_idx = stack[stack_ptr];
                continue;
            }

            if node.n_primitives > 0 {
                // Leaf: test all primitives
                for i in 0..node.n_primitives {
                    let prim_idx = node.primitives_offset + i as usize;
                    let prim = self.primitive(prim_idx);
                    if let Some(si) = prim.intersect(r, Some(hit_t)) {
                        hit_t = si.t_hit();
                        best_si = Some(si);
                    }
                }

                if stack_ptr == 0 {
                    break;
                }
                stack_ptr -= 1;
                node_idx = stack[stack_ptr];
            } else {
                // Interior: push far, visit near
                if dir_is_neg[node.axis as usize] == 1 {
                    stack[stack_ptr] = node_idx + 1;
                    stack_ptr += 1;
                    node_idx = node.primitives_offset; // second child
                } else {
                    stack[stack_ptr] = node.primitives_offset;
                    stack_ptr += 1;
                    node_idx += 1; // first child
                }
            }
        }

        best_si
    }

    fn intersect_p(&self, r: &Ray, t_max: Option<Float>) -> bool {
        if self.nodes.is_empty() || self.node_count == 0 {
            return false;
        }

        let t_max = t_max.unwrap_or(Float::INFINITY);

        let inv_dir = Vector3f::new(1.0 / r.d.x(), 1.0 / r.d.y(), 1.0 / r.d.z());
        let dir_is_neg = [
            if inv_dir.x() < 0.0 { 1 } else { 0 },
            if inv_dir.y() < 0.0 { 1 } else { 0 },
            if inv_dir.z() < 0.0 { 1 } else { 0 },
        ];

        let mut stack = [0usize; 64];
        let mut stack_ptr = 0;
        let mut node_idx = 0usize;

        loop {
            let node = self.node(node_idx);

            if node
                .bounds
                .intersect_p(r.o, t_max, inv_dir, &dir_is_neg)
                .is_none()
            {
                if stack_ptr == 0 {
                    break;
                }
                stack_ptr -= 1;
                node_idx = stack[stack_ptr];
                continue;
            }

            if node.n_primitives > 0 {
                for i in 0..node.n_primitives {
                    let prim_idx = node.primitives_offset + i as usize;
                    let prim = self.primitive(prim_idx);
                    if prim.intersect_p(r, Some(t_max)) {
                        return true;
                    }
                }

                if stack_ptr == 0 {
                    break;
                }
                stack_ptr -= 1;
                node_idx = stack[stack_ptr];
            } else {
                if dir_is_neg[node.axis as usize] == 1 {
                    stack[stack_ptr] = node_idx + 1;
                    stack_ptr += 1;
                    node_idx = node.primitives_offset;
                } else {
                    stack[stack_ptr] = node.primitives_offset;
                    stack_ptr += 1;
                    node_idx += 1;
                }
            }
        }

        false
    }
}
