use std::sync::Arc;

use rand::seq::SliceRandom;
use ray_tracing_core::{
    prelude::*,
    scene::{Intersection, LightSample, Scene},
};
use sampling::sample_triangle;

use crate::util::triangle_normal;

pub mod triangle_bvh;

pub trait AccelerationStructure<T> {
    fn intersect(&self, ray: Ray, min: Float, max: Float) -> Option<(Float, Dir3, T)>;
}

#[derive(Debug)]
pub struct ASMaterial<R: Rng> {
    pub material: Option<Box<dyn Material<R>>>,
    pub light: Option<Box<dyn Light<R>>>,
}

impl<R: Rng> ASMaterial<R> {
    pub fn new_material<M: Material<R> + 'static>(material: M) -> Self {
        Self {
            material: Some(Box::new(material)),
            light: None,
        }
    }
    pub fn new_light<L: Light<R> + 'static>(light: L) -> Self {
        Self {
            material: None,
            light: Some(Box::new(light)),
        }
    }
    pub fn new_material_light<M: Material<R> + 'static, L: Light<R> + 'static>(
        material: M,
        light: L,
    ) -> Self {
        Self {
            material: Some(Box::new(material)),
            light: Some(Box::new(light)),
        }
    }
}

#[derive(Debug)]
pub struct AccelerationStructureScene<A, R: Rng> {
    acceleration_structure: A,
    materials: Arc<[ASMaterial<R>]>,
    lights: Vec<(usize, [Pos3; 3])>,
}

impl<A: Clone, R: Rng> Clone for AccelerationStructureScene<A, R> {
    fn clone(&self) -> Self {
        Self {
            acceleration_structure: self.acceleration_structure.clone(),
            materials: Arc::clone(&self.materials),
            lights: self.lights.clone(),
        }
    }
}

impl<A, R: Rng> AccelerationStructureScene<A, R> {
    pub fn new(a: A, materials: Arc<[ASMaterial<R>]>, lights: Vec<(usize, [Pos3; 3])>) -> Self {
        Self {
            acceleration_structure: a,
            materials,
            lights,
        }
    }
}

impl<A: AccelerationStructure<u32>, R: Rng> Scene<R> for AccelerationStructureScene<A, R> {
    type Mat<'a>
        = &'a dyn Material<R>
    where
        A: 'a,
        R: 'a;

    type Light<'b>
        = &'b dyn Light<R>
    where
        A: 'b,
        R: 'b;

    fn intersect(
        &self,
        ray: Ray,
        min: Float,
        max: Float,
    ) -> Option<ray_tracing_core::scene::Intersection<R, Self::Mat<'_>, Self::Light<'_>>> {
        let (t, n, i) = self.acceleration_structure.intersect(ray, min, max)?;

        let material = &self.materials[i as usize];

        let light_pdf = if let Some(_l) = &material.light {
            0.0
        } else {
            0.0
        };

        Some(Intersection::new(
            t,
            n,
            material.material.as_deref(),
            material.light.as_deref(),
            light_pdf,
        ))
    }

    fn sample_light<'b>(
        &self,
        _w_in: Dir3,
        _intersection: &ray_tracing_core::scene::Intersection<R, Self::Mat<'b>, Self::Light<'b>>,
        rng: &mut R,
    ) -> Option<ray_tracing_core::scene::LightSample<'_, R>>
    where
        Self: 'b,
    {
        let t = self.lights.choose(rng);

        if let Some(&(t, v)) = t {
            let light = self.materials[t].light.as_ref().unwrap().as_ref();
            let b = sample_triangle(rng.gen());

            let n = triangle_normal(v);
            let area = n.length() * 0.5;

            Some(LightSample::new(
                Pos3::from_barycentric(v, b),
                1.0 / ((self.lights.len() as Float) * area),
                n.normalize(),
                light,
            ))
        } else {
            None
        }
    }
}
