#![allow(dead_code)]

use kioku::Arena;

use crate::{
    accel::BVH4,
    bbox::BBox,
    boundable::Boundable,
    lerp::lerp_slice,
    math::{cross, dot, Normal, Point, XformFull},
    ray::{LocalRay, Ray},
    shading::SurfaceShader,
};

use super::{triangle, Surface, SurfaceIntersection, SurfaceIntersectionData};

const MAX_LEAF_TRIANGLE_COUNT: usize = 3;

#[derive(Copy, Clone, Debug)]
pub struct TriangleMesh<'a> {
    time_sample_count: usize,
    vertices: &'a [Point], // Vertices, with the time samples for each vertex stored contiguously
    normals: Option<&'a [Normal]>, // Vertex normals, organized the same as `vertices`
    indices: &'a [(u32, u32, u32, u32)], // (v0_idx, v1_idx, v2_idx, original_tri_idx)
    accel: BVH4<'a>,
}

impl<'a> TriangleMesh<'a> {
    pub fn from_verts_and_indices<'b>(
        arena: &'b Arena,
        verts: &[Vec<Point>],
        vert_normals: &Option<Vec<Vec<Normal>>>,
        tri_indices: &[(usize, usize, usize)],
    ) -> TriangleMesh<'b> {
        let vert_count = verts[0].len();
        let time_sample_count = verts.len();

        // Copy verts over to a contiguous area of memory, reorganizing them
        // so that each vertices' time samples are contiguous in memory.
        let vertices = {
            let vertices = arena.alloc_array_uninit(vert_count * time_sample_count);

            for vi in 0..vert_count {
                for ti in 0..time_sample_count {
                    unsafe {
                        *vertices[(vi * time_sample_count) + ti].as_mut_ptr() = verts[ti][vi];
                    }
                }
            }

            unsafe { std::mem::transmute(vertices) }
        };

        // Copy vertex normals, if any, organizing them the same as vertices
        // above.
        let normals = match vert_normals {
            Some(ref vnors) => {
                let normals = arena.alloc_array_uninit(vert_count * time_sample_count);

                for vi in 0..vert_count {
                    for ti in 0..time_sample_count {
                        unsafe {
                            *normals[(vi * time_sample_count) + ti].as_mut_ptr() = vnors[ti][vi];
                        }
                    }
                }

                unsafe { Some(std::mem::transmute(&normals[..])) }
            }

            None => None,
        };

        // Copy triangle vertex indices over, appending the triangle index itself to the tuple
        let indices: &mut [(u32, u32, u32, u32)] = {
            let indices = arena.alloc_array_uninit(tri_indices.len());
            for (i, tri_i) in tri_indices.iter().enumerate() {
                unsafe {
                    *indices[i].as_mut_ptr() =
                        (tri_i.0 as u32, tri_i.2 as u32, tri_i.1 as u32, i as u32);
                }
            }
            unsafe { std::mem::transmute(indices) }
        };

        // Create bounds array for use during BVH construction
        let bounds = {
            let mut bounds = Vec::with_capacity(indices.len() * time_sample_count);
            for tri in tri_indices {
                for ti in 0..time_sample_count {
                    let p0 = verts[ti][tri.0];
                    let p1 = verts[ti][tri.1];
                    let p2 = verts[ti][tri.2];
                    let minimum = p0.min(p1.min(p2));
                    let maximum = p0.max(p1.max(p2));
                    bounds.push(BBox::from_points(minimum, maximum));
                }
            }
            bounds
        };

        // Build BVH
        let accel = BVH4::from_objects(arena, &mut indices[..], MAX_LEAF_TRIANGLE_COUNT, |tri| {
            &bounds
                [(tri.3 as usize * time_sample_count)..((tri.3 as usize + 1) * time_sample_count)]
        });

        TriangleMesh {
            time_sample_count: time_sample_count,
            vertices: vertices,
            normals: normals,
            indices: indices,
            accel: accel,
        }
    }
}

impl<'a> Boundable for TriangleMesh<'a> {
    fn bounds(&self) -> &[BBox] {
        self.accel.bounds()
    }
}

impl<'a> Surface for TriangleMesh<'a> {
    fn intersect_ray(
        &self,
        ray: &mut Ray,
        local_ray: &LocalRay,
        space: &XformFull,
        isect: &mut SurfaceIntersection,
        shader: &dyn SurfaceShader,
    ) {
        self.accel.traverse(ray, local_ray, |idx_range, ray| {
            // Iterate through the triangles and test the ray against them.
            let mut non_shadow_hit = false;
            let mut hit_tri = std::mem::MaybeUninit::uninit();
            let mut hit_tri_indices = std::mem::MaybeUninit::uninit();
            let mut hit_tri_data = std::mem::MaybeUninit::uninit();
            let ray_pre = triangle::RayTriPrecompute::new(ray.dir);
            for tri_idx in idx_range.clone() {
                let tri_indices = self.indices[tri_idx];

                // Get triangle.
                let mut tri = if self.time_sample_count == 1 {
                    // No deformation motion blur, so fast-path it.
                    (
                        self.vertices[tri_indices.0 as usize],
                        self.vertices[tri_indices.1 as usize],
                        self.vertices[tri_indices.2 as usize],
                    )
                } else {
                    // Deformation motion blur, need to interpolate.
                    let p0_slice = &self.vertices[(tri_indices.0 as usize * self.time_sample_count)
                        ..((tri_indices.0 as usize + 1) * self.time_sample_count)];
                    let p1_slice = &self.vertices[(tri_indices.1 as usize * self.time_sample_count)
                        ..((tri_indices.1 as usize + 1) * self.time_sample_count)];
                    let p2_slice = &self.vertices[(tri_indices.2 as usize * self.time_sample_count)
                        ..((tri_indices.2 as usize + 1) * self.time_sample_count)];

                    let p0 = lerp_slice(p0_slice, ray.time);
                    let p1 = lerp_slice(p1_slice, ray.time);
                    let p2 = lerp_slice(p2_slice, ray.time);

                    (p0, p1, p2)
                };

                // Transform triangle into world space.
                tri.0 = tri.0.xform(space);
                tri.1 = tri.1.xform(space);
                tri.2 = tri.2.xform(space);

                // Test ray against triangle
                if let Some((t, b0, b1, b2)) =
                    triangle::intersect_ray(ray.orig, ray_pre, ray.max_t, tri)
                {
                    if ray.is_occlusion() {
                        *isect = SurfaceIntersection::Occlude;
                        ray.mark_done();
                        break;
                    } else {
                        non_shadow_hit = true;
                        ray.max_t = t;
                        unsafe {
                            *hit_tri.as_mut_ptr() = tri;
                            *hit_tri_indices.as_mut_ptr() = tri_indices;
                            *hit_tri_data.as_mut_ptr() = (t, b0, b1, b2);
                        }
                    }
                }
            }

            // Calculate intersection data if necessary.
            if non_shadow_hit {
                let hit_tri = unsafe { hit_tri.assume_init() };
                let (t, b0, b1, b2) = unsafe { hit_tri_data.assume_init() };

                // Calculate intersection point and error magnitudes
                let (pos, pos_err) = triangle::surface_point(hit_tri, (b0, b1, b2));

                // Calculate geometric surface normal
                let geo_normal = cross(hit_tri.0 - hit_tri.1, hit_tri.0 - hit_tri.2).into_normal();

                // Calculate interpolated surface normal, if any
                let shading_normal = if let Some(normals) = self.normals {
                    let hit_tri_indices = unsafe { hit_tri_indices.assume_init() };
                    let n0_slice = &normals[(hit_tri_indices.0 as usize * self.time_sample_count)
                        ..((hit_tri_indices.0 as usize + 1) * self.time_sample_count)];
                    let n1_slice = &normals[(hit_tri_indices.1 as usize * self.time_sample_count)
                        ..((hit_tri_indices.1 as usize + 1) * self.time_sample_count)];
                    let n2_slice = &normals[(hit_tri_indices.2 as usize * self.time_sample_count)
                        ..((hit_tri_indices.2 as usize + 1) * self.time_sample_count)];

                    let n0 = lerp_slice(n0_slice, ray.time).normalized();
                    let n1 = lerp_slice(n1_slice, ray.time).normalized();
                    let n2 = lerp_slice(n2_slice, ray.time).normalized();

                    let s_nor = ((n0 * b0) + (n1 * b1) + (n2 * b2)).xform_fast(&space);
                    if dot(s_nor, geo_normal) >= 0.0 {
                        s_nor
                    } else {
                        -s_nor
                    }
                } else {
                    geo_normal
                };

                let intersection_data = SurfaceIntersectionData {
                    incoming: ray.dir,
                    t: t,
                    pos: pos,
                    pos_err: pos_err,
                    nor: shading_normal,
                    nor_g: geo_normal,
                    local_space: *space,
                    sample_pdf: 0.0,
                };

                // Fill in intersection data
                *isect = SurfaceIntersection::Hit {
                    intersection_data: intersection_data,
                    closure: shader.shade(&intersection_data, ray.time),
                };
            }
        });
    }
}