use crate::core::bsdf::BSDFSample;
use crate::core::bxdf::{BxDFFlags, BxDFReflTransFlags, BxDFTrait, FArgs, TransportMode};
use crate::core::geometry::{
    Normal3f, Point2f, Vector3f, VectorLike, abs_cos_theta, same_hemisphere,
};
use crate::core::scattering::{TrowbridgeReitzDistribution, fr_complex_from_spectrum, reflect};
use crate::spectra::SampledSpectrum;
use crate::utils::sampling::{cosine_hemisphere_pdf, sample_cosine_hemisphere};
use crate::{Float, INV_PI};
use core::any::Any;
use num_traits::Float as NumFloat;

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct ConductorBxDF {
    pub mf_distrib: TrowbridgeReitzDistribution,
    pub eta: SampledSpectrum,
    pub k: SampledSpectrum,
}

unsafe impl Send for ConductorBxDF {}
unsafe impl Sync for ConductorBxDF {}

impl ConductorBxDF {
    pub fn new(
        mf_distrib: &TrowbridgeReitzDistribution,
        eta: SampledSpectrum,
        k: SampledSpectrum,
    ) -> Self {
        Self {
            mf_distrib: *mf_distrib,
            eta,
            k,
        }
    }
}

impl BxDFTrait for ConductorBxDF {
    fn flags(&self) -> BxDFFlags {
        if self.mf_distrib.effectively_smooth() {
            BxDFFlags::SPECULAR_REFLECTION
        } else {
            BxDFFlags::GLOSSY_REFLECTION
        }
    }

    fn sample_f(&self, wo: Vector3f, _uc: Float, u: Point2f, f_args: FArgs) -> Option<BSDFSample> {
        let reflection_flags =
            BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::REFLECTION.bits());
        if !f_args.sample_flags.contains(reflection_flags) {
            return None;
        }

        if self.mf_distrib.effectively_smooth() {
            let wi = Vector3f::new(-wo.x(), -wo.y(), wo.z());
            let f =
                fr_complex_from_spectrum(abs_cos_theta(wi), self.eta, self.k) / abs_cos_theta(wi);

            let bsdf = BSDFSample {
                f,
                wi,
                pdf: 1.,
                flags: BxDFFlags::SPECULAR_REFLECTION,
                ..Default::default()
            };

            return Some(bsdf);
        }

        if wo.z() == 0. {
            return None;
        }
        let wm = self.mf_distrib.sample_wm(wo, u);
        let wi = reflect(wo, wm.into());
        if !same_hemisphere(wo, wi) {
            return None;
        }

        let pdf = self.mf_distrib.pdf(wo, wm) / (4. * wo.dot(wm).abs());

        let cos_theta_o = abs_cos_theta(wo);
        let cos_theta_i = abs_cos_theta(wi);
        if cos_theta_i == 0. || cos_theta_o == 0. {
            return None;
        }

        let f_spectrum = fr_complex_from_spectrum(wo.dot(wi).abs(), self.eta, self.k);
        let f = self.mf_distrib.d(wm) * f_spectrum * self.mf_distrib.g(wo, wi)
            / (4. * cos_theta_i * cos_theta_o);

        let bsdf = BSDFSample {
            f,
            wi,
            pdf,
            flags: BxDFFlags::GLOSSY_REFLECTION,
            ..Default::default()
        };

        Some(bsdf)
    }

    fn f(&self, wo: Vector3f, wi: Vector3f, _mode: TransportMode) -> SampledSpectrum {
        if !same_hemisphere(wo, wi) {
            return SampledSpectrum::default();
        }
        if self.mf_distrib.effectively_smooth() {
            return SampledSpectrum::default();
        }

        let cos_theta_o = abs_cos_theta(wo);
        let cos_theta_i = abs_cos_theta(wi);
        if cos_theta_i == 0. || cos_theta_o == 0. {
            return SampledSpectrum::new(0.);
        }

        let wm = wi + wo;
        if wm.norm_squared() == 0. {
            return SampledSpectrum::new(0.);
        }
        let wm_norm = wm.normalize();

        let f_spectrum = fr_complex_from_spectrum(wo.dot(wm).abs(), self.eta, self.k);
        self.mf_distrib.d(wm_norm) * f_spectrum * self.mf_distrib.g(wo, wi)
            / (4. * cos_theta_i * cos_theta_o)
    }

    fn pdf(&self, wo: Vector3f, wi: Vector3f, f_args: FArgs) -> Float {
        let reflection_flags =
            BxDFReflTransFlags::from_bits_truncate(BxDFReflTransFlags::REFLECTION.bits());
        if !f_args.sample_flags.contains(reflection_flags) {
            return 0.;
        }
        if !same_hemisphere(wo, wi) {
            return 0.;
        }
        if self.mf_distrib.effectively_smooth() {
            return 0.;
        }
        let wm = wo + wi;
        if wm.norm_squared() == 0. {
            return 0.;
        }
        let wm_corr = Normal3f::new(0., 0., 1.).face_forward(wm);
        self.mf_distrib.pdf(wo, wm_corr.into()) / (4. * wo.dot(wm).abs())
    }

    fn regularize(&mut self) {
        self.mf_distrib.regularize();
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}