use crate::{
    color::Color,
    lerp::lerp_slice,
    math::XformFull,
    ray::{LocalRay, Ray},
    scene::{Assembly, InstanceType, Object},
    shading::{SimpleSurfaceShader, SurfaceShader},
    surface::SurfaceIntersection,
};

pub struct Tracer<'a> {
    root: &'a Assembly<'a>,
    ray_trace_count: u64,
}

impl<'a> Tracer<'a> {
    pub fn from_assembly(assembly: &'a Assembly) -> Tracer<'a> {
        Tracer {
            root: assembly,
            ray_trace_count: 0,
        }
    }

    pub fn rays_traced(&self) -> u64 {
        self.ray_trace_count
    }

    pub fn trace(&mut self, mut ray: Ray) -> SurfaceIntersection {
        self.ray_trace_count += 1;

        let local_ray = ray.to_local();
        let space = XformFull::identity();
        let mut isect = SurfaceIntersection::Miss;

        self.trace_assembly(self.root, &mut ray, &local_ray, &space, &mut isect);

        isect
    }

    fn trace_assembly(
        &mut self,
        assembly: &Assembly,
        ray: &mut Ray,
        local_ray: &LocalRay,
        space: &XformFull,
        isect: &mut SurfaceIntersection,
    ) {
        assembly
            .object_accel
            .traverse(ray, local_ray, |idx_range, ray| {
                for inst_idx in idx_range {
                    let inst = &assembly.instances[inst_idx];

                    // Handle transforms if needed.
                    let (local_space, local_ray) = if let Some((xstart, xend)) =
                        inst.transform_indices
                    {
                        let instance_xform = lerp_slice(&assembly.xforms[xstart..xend], ray.time);
                        let combined_xform = instance_xform.compose(&space.fwd);

                        if let Some(xform) = combined_xform.to_full() {
                            (xform, ray.to_local_xform(&xform))
                        } else {
                            // Invalid transform, so skip traversing into this instance.
                            continue;
                        }
                    } else {
                        (*space, *local_ray)
                    };

                    // Trace ray.
                    match inst.instance_type {
                        InstanceType::Object => {
                            self.trace_object(
                                &assembly.objects[inst.data_index],
                                inst.surface_shader_index
                                    .map(|i| assembly.surface_shaders[i]),
                                ray,
                                &local_ray,
                                &local_space,
                                isect,
                            );
                        }

                        InstanceType::Assembly => {
                            self.trace_assembly(
                                &assembly.assemblies[inst.data_index],
                                ray,
                                &local_ray,
                                &local_space,
                                isect,
                            );
                        }
                    }

                    if ray.is_done() {
                        return;
                    }
                }
            });
    }

    fn trace_object<'b>(
        &mut self,
        obj: &Object,
        surface_shader: Option<&dyn SurfaceShader>,
        ray: &mut Ray,
        local_ray: &LocalRay,
        space: &XformFull,
        isect: &mut SurfaceIntersection,
    ) {
        match *obj {
            Object::Surface(surface) => {
                let unassigned_shader = SimpleSurfaceShader::Emit {
                    color: Color::new_xyz(color::rec709_to_xyz((1.0, 0.0, 1.0))),
                };
                let shader = surface_shader.unwrap_or(&unassigned_shader);

                surface.intersect_ray(ray, local_ray, space, isect, shader);
            }

            Object::SurfaceLight(surface) => {
                // Lights don't use shaders
                let bogus_shader = SimpleSurfaceShader::Emit {
                    color: Color::new_xyz(color::rec709_to_xyz((1.0, 0.0, 1.0))),
                };

                surface.intersect_ray(ray, local_ray, space, isect, &bogus_shader);
            }
        }
    }
}