/// Implementation based on blog post @
/// https://jacco.ompf2.com/2022/04/13/how-to-build-a-bvh-part-1-basics/
use std::f32::EPSILON;
use std::fmt;

use stl::STL;

use crate::{
    aabb::AABB,
    hitable::{Hit, HitRecord},
    material::Material,
    ray::Ray,
    vec3::{cross, dot, Vec3},
};

#[derive(Debug, PartialEq)]
struct BVHNode {
    aabb: AABB,
    // When prim_count==0, left_first holds the left child's index in bvh_nodes. When >0 left_first
    // holds the index for the first triangle in triangles.
    left_first: u32,
    prim_count: u32,
}

impl BVHNode {
    fn is_leaf(&self) -> bool {
        self.prim_count > 0
    }
}

#[derive(PartialEq)]
pub struct Triangle {
    centroid: Vec3,
    verts: [Vec3; 3],
}
impl fmt::Debug for Triangle {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Tri: <{}, {}, {}> @ {}",
            self.verts[0], self.verts[1], self.verts[2], self.centroid
        )
    }
}

pub struct BVHTriangles<M>
where
    M: Material,
{
    pub triangles: Vec<Triangle>,
    material: M,
    bvh_nodes: Vec<BVHNode>,
}

impl<M> fmt::Debug for BVHTriangles<M>
where
    M: Material,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} triangles", self.triangles.len())?;
        if f.alternate() {
            writeln!(f)?;
        }
        for (i, t) in self.triangles.iter().enumerate() {
            if f.alternate() {
                write!(f, "\t")?;
            }
            write!(f, "{i:>3} {:?} ", t)?;
            if f.alternate() {
                writeln!(f)?;
            }
        }
        if f.alternate() {
            writeln!(f)?;
        }
        for (i, n) in self.bvh_nodes.iter().enumerate() {
            write!(f, "N[{i}] {n:?}")?;
            if f.alternate() {
                writeln!(f)?;
            }
            let n = &self.bvh_nodes[i];
            if n.is_leaf() {
                for t_idx in n.left_first..(n.left_first + n.prim_count) {
                    if f.alternate() {
                        write!(f, "\t")?;
                    }
                    write!(f, "{:?} ", self.triangles[t_idx as usize])?;
                    if f.alternate() {
                        writeln!(f)?;
                    }
                }
            }
        }
        Ok(())
    }
}

const ROOT_NODE_IDX: usize = 0;
impl<M> BVHTriangles<M>
where
    M: Material,
{
    pub fn new(stl: &STL, material: M) -> BVHTriangles<M> {
        assert_eq!(std::mem::size_of::<BVHNode>(), 32);
        let div3 = 1. / 3.;
        let triangles: Vec<_> = stl
            .triangles
            .iter()
            .map(|t| {
                let v0 = t.verts[0];
                let v1 = t.verts[1];
                let v2 = t.verts[2];
                let centroid = (v0 + v1 + v2) * div3;
                Triangle {
                    centroid,
                    verts: [v0, v1, v2],
                }
            })
            .collect();

        let n = 2 * triangles.len() - 2;
        let bvh_nodes = Vec::with_capacity(n);
        let mut bvh = BVHTriangles {
            triangles,
            bvh_nodes,
            material,
        };
        bvh.build_bvh();
        bvh
    }

    fn build_bvh(&mut self) {
        // assign all triangles to root node
        let root = BVHNode {
            aabb: AABB::default(),
            left_first: 0,
            prim_count: self.triangles.len() as u32,
        };
        self.bvh_nodes.push(root);
        self.update_node_bounds(ROOT_NODE_IDX);
        // subdivide recursively
        self.subdivide(ROOT_NODE_IDX);
    }
    fn update_node_bounds(&mut self, node_idx: usize) {
        let node = &mut self.bvh_nodes[node_idx];
        let mut aabb_min: Vec3 = f32::MAX.into();
        let mut aabb_max: Vec3 = f32::MIN.into();
        for i in node.left_first..(node.left_first + node.prim_count) {
            let leaf_tri = &self.triangles[i as usize];
            aabb_min = vec3::min(aabb_min, leaf_tri.verts[0]);
            aabb_min = vec3::min(aabb_min, leaf_tri.verts[1]);
            aabb_min = vec3::min(aabb_min, leaf_tri.verts[2]);
            aabb_max = vec3::max(aabb_max, leaf_tri.verts[0]);
            aabb_max = vec3::max(aabb_max, leaf_tri.verts[1]);
            aabb_max = vec3::max(aabb_max, leaf_tri.verts[2]);
        }
        node.aabb = AABB::new(aabb_min, aabb_max);
    }
    fn subdivide(&mut self, idx: usize) {
        // Early out if we're down to just 2 or less triangles
        if self.bvh_nodes[idx].prim_count <= 2 {
            return;
        }
        let (left_first, prim_count, left_count, i) = {
            let node = &self.bvh_nodes[idx];

            // Compute split plane and position.
            let extent = node.aabb.max() - node.aabb.min();
            let axis = node.aabb.longest_axis();
            let split_pos = node.aabb.min()[axis] + extent[axis] * 0.5;

            // Split the group in two halves.
            let mut i = node.left_first as isize;
            let mut j = i + node.prim_count as isize - 1;
            while i <= j {
                if self.triangles[i as usize].centroid[axis] < split_pos {
                    i += 1;
                } else {
                    self.triangles.swap(i as usize, j as usize);
                    j -= 1;
                }
            }

            // Create child nodes for each half.
            let left_count = i as u32 - node.left_first;
            if left_count == 0 || left_count == node.prim_count {
                return;
            }
            (node.left_first, node.prim_count, left_count, i)
        };

        // create child nodes
        let left_child_idx = self.bvh_nodes.len() as u32;
        let right_child_idx = left_child_idx + 1 as u32;
        let left = BVHNode {
            aabb: AABB::default(),
            left_first,
            prim_count: left_count as u32,
        };
        let right = BVHNode {
            aabb: AABB::default(),
            left_first: i as u32,
            prim_count: (prim_count - left_count) as u32,
        };
        self.bvh_nodes.push(left);
        self.bvh_nodes.push(right);
        let node = &mut self.bvh_nodes[idx];
        node.left_first = left_child_idx;
        node.prim_count = 0;

        // Recurse
        self.update_node_bounds(left_child_idx as usize);
        self.update_node_bounds(right_child_idx as usize);
        self.subdivide(left_child_idx as usize);
        self.subdivide(right_child_idx as usize);
    }

    fn intersect_bvh(&self, r: Ray, node_idx: u32, t_min: f32, t_max: f32) -> Option<HitRecord> {
        let node = &self.bvh_nodes[node_idx as usize];
        if !node.aabb.hit(r, t_min, t_max) {
            return None;
        }
        //dbg!(&self);

        if node.is_leaf() {
            return self
                .triangles
                .iter()
                .map(|tri| {
                    if let Some(RayTriangleResult { t, p }) = intersect_tri(r, tri) {
                        // We don't support UV (yet?).
                        let uv = (0.5, 0.5);
                        let v0 = tri.verts[0];
                        let v1 = tri.verts[1];
                        let v2 = tri.verts[2];

                        let v0v1 = v1 - v0;
                        let v0v2 = v2 - v0;
                        let normal = cross(v0v1, v0v2).unit_vector();
                        //println!("hit triangle {tri:?}");
                        Some(HitRecord {
                            t,
                            uv,
                            p,
                            normal,
                            material: &self.material,
                        })
                    } else {
                        None
                    }
                })
                .filter_map(|hr| hr)
                .min_by(|a, b| a.t.partial_cmp(&b.t).unwrap());
        } else {
            let r1 = self.intersect_bvh(r, node.left_first, t_min, t_max);
            let r2 = self.intersect_bvh(r, node.left_first + 1, t_min, t_max);
            // Merge results, if both hit, take the one closest to the ray origin (smallest t
            // value).
            match (&r1, &r2) {
                (Some(_), None) => return r1,
                (None, Some(_)) => return r2,
                (None, None) => (),
                (Some(rp1), Some(rp2)) => return if rp1.t < rp2.t { r1 } else { r2 },
            }
        }
        None
    }
}

/// Based on https://www.scratchapixel.com/lessons/3d-basic-rendering/ray-tracing-rendering-a-triangle/moller-trumbore-ray-triangle-intersection.html
fn intersect_tri(r: Ray, tri: &Triangle) -> Option<RayTriangleResult> {
    // #ifdef MOLLER_TRUMBORE
    //     Vec3f v0v1 = v1 - v0;
    //     Vec3f v0v2 = v2 - v0;
    //     Vec3f pvec = dir.crossProduct(v0v2);
    //     float det = v0v1.dotProduct(pvec);
    // #ifdef CULLING
    //     // if the determinant is negative, the triangle is 'back facing'
    //     // if the determinant is close to 0, the ray misses the triangle
    //     if (det < kEpsilon) return false;
    // #else
    //     // ray and triangle are parallel if det is close to 0
    //     if (fabs(det) < kEpsilon) return false;
    // #endif
    //     float invDet = 1 / det;
    //
    //     Vec3f tvec = orig - v0;
    //     u = tvec.dotProduct(pvec) * invDet;
    //     if (u < 0 || u > 1) return false;
    //
    //     Vec3f qvec = tvec.crossProduct(v0v1);
    //     v = dir.dotProduct(qvec) * invDet;
    //     if (v < 0 || u + v > 1) return false;
    //
    //     t = v0v2.dotProduct(qvec) * invDet;
    //
    let v0 = tri.verts[0];
    let v1 = tri.verts[1];
    let v2 = tri.verts[2];

    let v0v1 = v1 - v0;
    let v0v2 = v2 - v0;
    let p = cross(r.direction, v0v2);
    let det = dot(v0v1, p);
    if det.abs() < EPSILON {
        return None;
    }

    let inv_det = 1. / det;

    let t = r.origin - v0;
    let u = dot(t, p) * inv_det;
    if u < 0. || u > 1. {
        return None;
    }

    let q = cross(t, v0v1);
    let v = dot(r.direction, q) * inv_det;
    if v < 0. || u + v > 1. {
        return None;
    }
    let t = dot(v0v2, q) * inv_det;

    if t > EPSILON {
        return Some(RayTriangleResult {
            t,
            p: r.point_at_parameter(t),
        });
    }
    None
}

impl<M> Hit for BVHTriangles<M>
where
    M: Material,
{
    fn hit(&self, r: Ray, t_min: f32, t_max: f32) -> Option<HitRecord> {
        self.intersect_bvh(r, 0, t_min, t_max)
    }

    fn bounding_box(&self, _t_min: f32, _t_max: f32) -> Option<AABB> {
        Some(self.bvh_nodes[0].aabb)
    }
}

struct RayTriangleResult {
    t: f32,
    p: Vec3,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        bvh_triangles::BVHTriangles,
        cuboid::Cuboid,
        hitable::Hit,
        material::{Dielectric, Lambertian},
        ray::Ray,
        texture::ConstantTexture,
    };
    use pretty_assertions::assert_eq;
    use std::{
        io::{BufReader, Cursor},
        sync::Arc,
    };
    use stl::STL;

    /*
        #[test]
        fn build_bvh() {
            let stl_triangles: Vec<_> = (0..4)
                .flat_map(|y| {
                    (0..2).map(move |x| {
                        let x = x as f32;
                        let y = y as f32;
                        stl::Triangle {
                            normal: [1., 0., 0.].into(),
                            verts: [
                                [2. * x + 0., 2. * y + 0., 0.].into(),
                                [2. * x + 1., 2. * y + 0., 0.].into(),
                                [2. * x + 1., 2. * y + 1., 0.].into(),
                            ],
                            attr: 0,
                        }
                    })
                })
                .collect();
            let stl = STL {
                header: [0; 80],
                triangles: stl_triangles,
            };
            /*
            let mut bvh_triangles: Vec<_> = stl_triangles
            .iter()
            .map(|tri| {
            let div3 = 1. / 3.;
            let v0 = tri.verts[0];
            let v1 = tri.verts[1];
            let v2 = tri.verts[2];
            let centroid = (v0 + v1 + v2) * div3;
            Triangle {
            centroid,
            verts: tri.verts,
            }
            })
            .collect();
            bvh_triangles.sort_by(|a, b| a.centroid.y.partial_cmp(&b.centroid.y).unwrap());
            let material = Lambertian::new(ConstantTexture::new([0., 0., 0.]));
            let bvh_nodes = Default::default();
            let want = BVHTriangles {
            triangles: bvh_triangles,
            bvh_nodes,
            material,
            };
            */
            let material = Lambertian::new(ConstantTexture::new([0., 0., 0.]));
            let bvh = BVHTriangles::new(&stl, material);
            dbg!(&bvh);
            assert_eq!(bvh.bvh_nodes.len(), 2 * bvh.triangles.len() - 2);
        }
    */

    #[test]
    fn compare_cuboid() {
        let c = Cuboid::new(
            [0., 0., 0.].into(),
            [20., 20., 20.].into(),
            Arc::new(Dielectric::new(1.5)),
        );
        let stl_cube = STL::parse(
            BufReader::new(Cursor::new(include_bytes!("../stls/cube.stl"))),
            false,
        )
        .expect("failed to parse cube");

        let s = BVHTriangles::new(
            &stl_cube,
            Dielectric::new(1.5),
            //Metal::new(Vec3::new(0.8, 0.8, 0.8), 0.2),
            //Lambertian::new(ConstantTexture::new(Vec3::new(1.0, 0.2, 0.2))),
        );
        //dbg!(&s);
        let mut rays: Vec<_> = (1..20)
            .flat_map(|y| {
                (1..20).flat_map(move |x| {
                    let x = x as f32;
                    let y = y as f32;
                    vec![
                        // Outward in angle
                        Ray::new([-1., x, y], [1., 0., 0.], 0.),
                        Ray::new([21., x, y], [-1., 0., 0.], 0.),
                        Ray::new([x, -1., y], [0., 1., 0.], 0.),
                        Ray::new([x, 21., y], [0., -1., 0.], 0.),
                        Ray::new([x, y, -1.], [0., 0., 1.], 0.),
                        Ray::new([x, y, 21.], [0., 0., -1.], 0.),
                        // Inward out (
                        Ray::new([x, y, 10.], [1., 0., 0.], 0.),
                        Ray::new([x, y, 10.], [-1., 0., 0.], 0.),
                        Ray::new([x, 10., y], [1., 0., 0.], 0.),
                        Ray::new([x, 10., y], [-1., 0., 0.], 0.),
                        Ray::new([10., x, y], [1., 0., 0.], 0.),
                        Ray::new([10., x, y], [-1., 0., 0.], 0.),
                        Ray::new([x, y, 10.], [0., 1., 0.], 0.),
                        Ray::new([x, y, 10.], [0., -1., 0.], 0.),
                        Ray::new([x, 10., y], [0., 1., 0.], 0.),
                        Ray::new([x, 10., y], [0., -1., 0.], 0.),
                        Ray::new([10., x, y], [0., 1., 0.], 0.),
                        Ray::new([10., x, y], [0., -1., 0.], 0.),
                        Ray::new([x, y, 10.], [0., 0., 1.], 0.),
                        Ray::new([x, y, 10.], [0., 0., -1.], 0.),
                        Ray::new([x, 10., y], [0., 0., 1.], 0.),
                        Ray::new([x, 10., y], [0., 0., -1.], 0.),
                        Ray::new([10., x, y], [0., 0., 1.], 0.),
                        Ray::new([10., x, y], [0., 0., -1.], 0.),
                    ]
                    .into_iter()
                })
            })
            .collect();
        // These currently differ between STL and cuboid.
        if false {
            // Outward in at an angle.
            let sqrt2 = 2f32.sqrt();
            rays.extend(vec![
                Ray::new([-1., 10., 10.], [sqrt2, sqrt2, 0.], 0.),
                Ray::new([-1., 10., 10.], [sqrt2, -sqrt2, 0.], 0.),
                Ray::new([-1., 10., 10.], [sqrt2, 0., sqrt2], 0.),
                Ray::new([-1., 10., 10.], [sqrt2, 0., -sqrt2], 0.),
            ]);
        }

        // TODO(wathiede): proptest this, it's still not perfectly equal when rendering.

        for r in rays.into_iter() {
            let c_hit = c
                .hit(r, 0., f32::MAX)
                .expect(&format!("c_hit missed {r:#?}"));
            let s_hit = s
                .hit(r, 0., f32::MAX)
                .expect(&format!("s_hit missed {r:#?}"));
            assert!(
                (c_hit.t - s_hit.t).abs() < EPSILON,
                "{r:?} [t] c_hit: {c_hit:#?}, s_hit: {s_hit:#?}"
            );
            // uv isn't valid for BVHTriangles.
            // assert_eq!( c_hit.uv, s_hit.uv, "{i}: [uv] c_hit: {c_hit:?}, s_hit: {s_hit:?}");
            assert_eq!(
                c_hit.p, s_hit.p,
                "{r:?}: [p] c_hit: {c_hit:#?}, s_hit: {s_hit:#?}"
            );
            assert_eq!(
                c_hit.normal, s_hit.normal,
                "{r:?}: [normal] c_hit: {c_hit:?}, s_hit: {s_hit:?}"
            );
        }
    }
}
