use crate::core::bsdf::BSDFSample;
use crate::core::bxdf::{BxDFFlags, BxDFReflTransFlags, BxDFTrait, FArgs, TransportMode};
use crate::core::geometry::{Point2f, Vector3f, abs_cos_theta, same_hemisphere};
use crate::spectra::SampledSpectrum;
use crate::utils::sampling::{cosine_hemisphere_pdf, sample_cosine_hemisphere};
use crate::{Float, INV_PI};
use core::any::Any;

#[repr(C)]
#[derive(Debug, Copy, Clone, Default)]
pub struct DiffuseBxDF {
    pub r: SampledSpectrum,
}

impl DiffuseBxDF {
    pub fn new(r: SampledSpectrum) -> Self {
        Self { r }
    }
}

impl BxDFTrait for DiffuseBxDF {
    fn flags(&self) -> BxDFFlags {
        if !self.r.is_black() {
            BxDFFlags::DIFFUSE_REFLECTION
        } else {
            BxDFFlags::UNSET
        }
    }

    fn f(&self, wo: Vector3f, wi: Vector3f, _mode: TransportMode) -> SampledSpectrum {
        if !same_hemisphere(wo, wi) {
            return SampledSpectrum::new(0.);
        }
        self.r * INV_PI
    }

    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;
        }
        let mut wi = sample_cosine_hemisphere(u);
        if wo.z() < 0. {
            wi[2] *= -1.;
        }
        let pdf = cosine_hemisphere_pdf(abs_cos_theta(wi));
        let bsdf = BSDFSample {
            f: self.r * INV_PI,
            wi,
            pdf,
            flags: BxDFFlags::DIFFUSE_REFLECTION,
            ..Default::default()
        };
        Some(bsdf)
    }

    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) || !same_hemisphere(wo, wi) {
            return 0.;
        }
        cosine_hemisphere_pdf(abs_cos_theta(wi))
    }

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

    fn regularize(&mut self) {}
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct DiffuseTransmissionBxDF;