use ray_tracing_core::{
    color::Color,
    light::Light,
    prelude::{Float, Material, Rng},
    scene::{Intersection, Scene},
};

use crate::{
    bvh::{Bvh, BvhData, BvhIntersection},
    material::PbrtMaterial,
    shape::Shape,
};

#[derive(Debug)]
pub struct PbrtScene<R: Rng> {
    pub(crate) shapes: Bvh<Vec<Shape<R>>>,
    pub(crate) infinite_light: Option<PbrtInfiniteLight>,
}

pub struct SceneIntersectionWrapper<'a, R: Rng>(
    ray_tracing_core::scene::Intersection<R, UVMaterial<'a, R>, &'a dyn Light<R>>,
);

impl<R: Rng> BvhIntersection for SceneIntersectionWrapper<'_, R> {
    fn get_t(&self) -> Float {
        self.0.t()
    }
}

impl<R: Rng + std::fmt::Debug> BvhData for Vec<Shape<R>> {
    type Intersect<'a>
        = SceneIntersectionWrapper<'a, R>
    where
        R: 'a;

    fn len(&self) -> crate::bvh::Index {
        self.len() as crate::bvh::Index
    }

    fn sort_range_dim(&mut self, id: std::ops::Range<crate::bvh::Index>, dim: u8) {
        let get_key = |t: &Shape<R>| t.aabb.max()[dim as usize] - t.aabb.min()[dim as usize];
        self[id.start as usize..id.end as usize]
            .sort_by(|a, b| get_key(a).partial_cmp(&get_key(b)).unwrap());
    }

    fn get_aabb(&self, id: crate::bvh::Index) -> ray_tracing_core::prelude::AABB {
        self[id as usize].aabb
    }

    fn intersect(
        &self,
        id: crate::bvh::Index,
        ray: ray_tracing_core::prelude::Ray,
        min: Float,
        max: Float,
    ) -> Option<Self::Intersect<'_>> {
        self[id as usize]
            .intersect(ray, min, max)
            .map(|i| SceneIntersectionWrapper(i))
    }
}

#[derive(Debug)]
pub(crate) struct PbrtInfiniteLight {
    pub(crate) color: Color,
}

impl<R: Rng> Light<R> for PbrtInfiniteLight {
    fn emit(&self, w_in: ray_tracing_core::prelude::Dir3, rng: &mut R) -> Color {
        let _ = rng;
        let _ = w_in;
        self.color
    }
}

pub struct UVMaterial<'a, R: Rng> {
    pub(crate) u: Float,
    pub(crate) v: Float,
    pub(crate) material: &'a dyn PbrtMaterial<R>,
}

impl<'a, R: Rng> std::fmt::Debug for UVMaterial<'a, R> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("UVMaterial")
            .field("u", &self.u)
            .field("v", &self.v)
            .field("material", &self.material)
            .finish()
    }
}

impl<R: Rng> Material<R> for UVMaterial<'_, R> {
    fn eval(
        &self,
        w_in: ray_tracing_core::prelude::Dir3,
        w_out: ray_tracing_core::prelude::Dir3,
        rng: &mut R,
    ) -> Color {
        self.material.eval(self.u, self.v, w_in, w_out, rng)
    }

    fn sample(
        &self,
        w_in: ray_tracing_core::prelude::Dir3,
        rng: &mut R,
    ) -> ray_tracing_core::material::SampleResult {
        self.material.sample(self.u, self.v, w_in, rng)
    }

    fn pdf(
        &self,
        w_in: ray_tracing_core::prelude::Dir3,
        w_out: ray_tracing_core::prelude::Dir3,
    ) -> Float {
        self.material.pdf(self.u, self.v, w_in, w_out)
    }
}

impl<R: Rng + std::fmt::Debug> Scene<R> for PbrtScene<R> {
    type Mat<'a>
        = UVMaterial<'a, R>
    where
        R: 'a;

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

    fn intersect(
        &'_ self,
        ray: ray_tracing_core::prelude::Ray,
        min: ray_tracing_core::prelude::Float,
        max: ray_tracing_core::prelude::Float,
    ) -> Option<ray_tracing_core::scene::Intersection<R, Self::Mat<'_>, Self::Light<'_>>> {
        let i = self.shapes.intersect(ray, min, max).map(|i| i.0);

        i.or_else(|| {
            self.infinite_light.as_ref().map(|l| {
                Intersection::new(
                    Float::INFINITY,
                    -ray.dir(),
                    None,
                    Some(l as &dyn Light<R>),
                    0.0,
                )
            })
        })
    }

    fn sample_light<'b>(
        &self,
        _w_in: ray_tracing_core::prelude::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,
    {
        None
    }
}
