use super::RayIntegratorTrait;
use super::base::IntegratorBase;
use super::constants::*;
use super::state::PathState;
use crate::Arena;
use shared::core::bsdf::{BSDF, BSDFSample};
use shared::core::bxdf::{BxDFFlags, FArgs, TransportMode};
use shared::core::camera::Camera;
use shared::core::film::VisibleSurface;
use shared::core::geometry::{Point2i, Ray, Vector3f, VectorLike};
use shared::core::interaction::{Interaction, InteractionTrait, SurfaceInteraction};
use shared::core::light::LightTrait;
use shared::core::light::{Light, LightSampleContext};
use shared::core::primitive::Primitive;
use shared::core::sampler::{Sampler, SamplerTrait};
use shared::lights::sampler::LightSampler;
use shared::lights::sampler::LightSamplerTrait;
use shared::spectra::{SampledSpectrum, SampledWavelengths};
use shared::utils::math::square;
use shared::utils::sampling::{
    power_heuristic, sample_uniform_hemisphere, sample_uniform_sphere, uniform_hemisphere_pdf,
    uniform_sphere_pdf,
};
use std::sync::Arc;

#[derive(Clone, Copy)]
pub struct PathConfig {
    pub max_depth: usize,
    pub regularize: bool,
    pub sample_lights: bool,
    pub sample_bsdf: bool,
    pub handle_volumes: bool,
    pub use_mis: bool,
}

impl PathConfig {
    pub const SIMPLE: Self = Self {
        max_depth: 5,
        regularize: false,
        sample_lights: true,
        sample_bsdf: true,
        handle_volumes: false,
        use_mis: false,
    };

    pub const FULL: Self = Self {
        max_depth: 8,
        regularize: true,
        sample_lights: true,
        sample_bsdf: true,
        handle_volumes: false,
        use_mis: true,
    };

    pub const VOLUMETRIC: Self = Self {
        max_depth: 8,
        regularize: true,
        sample_lights: true,
        sample_bsdf: true,
        handle_volumes: true,
        use_mis: true,
    };
}

pub struct PathIntegrator {
    base: IntegratorBase,
    camera: Arc<Camera>,
    sampler: LightSampler,
    config: PathConfig,
}

unsafe impl Send for PathIntegrator {}
unsafe impl Sync for PathIntegrator {}

impl PathIntegrator {
    pub fn new(
        aggregate: Arc<Primitive>,
        lights: Vec<Arc<Light>>,
        camera: Arc<Camera>,
        sampler: LightSampler,
        config: PathConfig,
    ) -> Self {
        let base = IntegratorBase::new(aggregate, lights.clone());
        Self {
            base,
            camera,
            sampler,
            config,
        }
    }

    fn sample_direct(
        &self,
        intr: &SurfaceInteraction,
        bsdf: &BSDF,
        _state: &PathState,
        lambda: &SampledWavelengths,
        sampler: &mut Sampler,
    ) -> SampledSpectrum {
        let ctx = LightSampleContext::from(intr);

        let Some(sampled) = self.sampler.sample_with_context(&ctx, sampler.get1d()) else {
            return SampledSpectrum::zero();
        };

        let Some(ls) = sampled.light.sample_li(&ctx, sampler.get2d(), lambda, true) else {
            return SampledSpectrum::zero();
        };

        if ls.l.is_black() || ls.pdf == 0.0 {
            return SampledSpectrum::zero();
        }

        let wo = intr.wo();
        let wi = ls.wi;

        let Some(f) = bsdf.f(wo, wi, TransportMode::Radiance) else {
            return SampledSpectrum::zero();
        };

        let f = f * wi.abs_dot(intr.shading.n.into());
        if f.is_black() {
            return SampledSpectrum::zero();
        }

        if !self
            .base
            .unoccluded(&Interaction::Surface(intr.clone()), &ls.p_light)
        {
            return SampledSpectrum::zero();
        }

        let p_l = sampled.p * ls.pdf;

        if !self.config.use_mis || sampled.light.light_type().is_delta_light() {
            ls.l * f / p_l
        } else {
            let p_b = bsdf.pdf(wo, wi, FArgs::default());
            let w_l = power_heuristic(1, p_l, 1, p_b);
            w_l * ls.l * f / p_l
        }
    }

    fn sample_direction(
        &self,
        wo: Vector3f,
        bsdf: &BSDF,
        isect: &SurfaceInteraction,
        sampler: &mut Sampler,
    ) -> Option<BSDFSample> {
        if self.config.sample_bsdf {
            bsdf.sample_f(wo, sampler.get1d(), sampler.get2d(), FArgs::default())
        } else {
            self.sample_uniform(wo, bsdf, isect, sampler)
        }
    }

    fn sample_uniform(
        &self,
        wo: Vector3f,
        bsdf: &BSDF,
        isect: &SurfaceInteraction,
        sampler: &mut Sampler,
    ) -> Option<BSDFSample> {
        let flags = bsdf.flags();

        let (wi, pdf) = if flags.is_reflective() && flags.is_transmissive() {
            (sample_uniform_sphere(sampler.get2d()), uniform_sphere_pdf())
        } else {
            let mut wi = sample_uniform_hemisphere(sampler.get2d());
            let same_hemi = wo.dot(isect.n().into()) * wi.dot(isect.n().into()) > 0.0;

            if (flags.is_reflective() && !same_hemi) || (flags.is_transmissive() && same_hemi) {
                wi = -wi;
            }
            (wi, uniform_hemisphere_pdf())
        };

        let f = bsdf.f(wo, wi, TransportMode::Radiance)?;

        Some(BSDFSample {
            f,
            wi,
            pdf,
            flags: BxDFFlags::empty(), // or appropriate flags
            eta: 1.0,
            pdf_is_proportional: false,
        })
    }

    fn compute_visible_surface(
        &self,
        isect: &SurfaceInteraction,
        bsdf: &BSDF,
        lambda: &SampledWavelengths,
    ) -> VisibleSurface {
        let albedo = bsdf.rho_wo(isect.wo(), &UC_RHO, &U_RHO);
        VisibleSurface::new(isect, &albedo, lambda)
    }
}

impl RayIntegratorTrait for PathIntegrator {
    fn li(
        &self,
        mut ray: Ray,
        lambda: &SampledWavelengths,
        sampler: &mut Sampler,
        want_visible: bool,
        _arena: &Arena,
    ) -> (SampledSpectrum, Option<VisibleSurface>) {
        let mut state = PathState::new();
        let mut visible = None;

        loop {
            let Some(mut si) = self.base.intersect(&ray, None) else {
                self.base.add_infinite_light_contribution(
                    &mut state,
                    &ray,
                    lambda,
                    Some(&self.sampler),
                    self.config.use_mis,
                );
                break;
            };

            let t_hit = si.t_hit();
            let isect = &mut si.intr;

            // Emission from hit surface
            let le = isect.le(-ray.d, lambda);
            if !le.is_black() {
                if state.depth == 0 || state.specular_bounce {
                    state.l += state.beta * le;
                } else if self.config.use_mis {
                    if !isect.area_light.is_null() {
                        let light = &isect.area_light;
                        let p_l = self.sampler.pmf_with_context(&state.prev_ctx, light)
                            * light.pdf_li(&state.prev_ctx, ray.d, true);
                        let w_b = power_heuristic(1, state.prev_pdf, 1, p_l);
                        state.l += state.beta * w_b * le;
                    }
                }
            }

            // Get BSDF
            let Some(mut bsdf) = isect.get_bsdf(&ray, lambda, &self.camera, sampler) else {
                state.specular_bounce = true;
                isect.skip_intersection(&mut ray, t_hit);
                continue;
            };

            // Visible surface at primary hit
            if state.depth == 0 && want_visible {
                visible = Some(self.compute_visible_surface(isect, &bsdf, lambda));
            }

            // Depth check
            if state.depth >= self.config.max_depth {
                break;
            }
            state.depth += 1;

            // Regularization
            if self.config.regularize && state.any_non_specular_bounces {
                bsdf.regularize();
            }

            // Direct lighting
            if self.config.sample_lights && bsdf.flags().is_non_specular() {
                state.l += state.beta * self.sample_direct(isect, &bsdf, &state, lambda, sampler);
            }

            // Sample BSDF for next direction
            let wo = -ray.d;
            let Some(bs) = self.sample_direction(wo, &bsdf, isect, sampler) else {
                break;
            };

            state.beta *= bs.f * bs.wi.abs_dot(isect.shading.n.into()) / bs.pdf;
            state.prev_pdf = if bs.pdf_is_proportional {
                bsdf.pdf(wo, bs.wi, FArgs::default())
            } else {
                bs.pdf
            };
            state.specular_bounce = bs.is_specular();
            state.any_non_specular_bounces |= !bs.is_specular();

            if bs.is_transmissive() {
                state.eta_scale *= square(bs.eta);
            }

            state.prev_ctx = LightSampleContext::from(&*isect);
            ray = isect.spawn_ray_with_differentials(&ray, bs.wi, bs.flags, bs.eta);

            if state.russian_roulette(sampler, 1) {
                break;
            }
        }

        (state.l, visible)
    }

    fn evaluate_pixel_sample(
        &self,
        p_pixel: Point2i,
        sample_ind: usize,
        sampler: &mut Sampler,
        arena: &Arena,
    ) {
        crate::integrators::pipeline::evaluate_pixel_sample(
            self,
            self.camera.as_ref(),
            sampler,
            p_pixel,
            sample_ind,
            arena,
        );
    }
}