use super::{
    Bounds3f, Float, Light, LightBounds, LightSampleContext, LightTrait, Normal3f, PI, Point3f,
    SampledSpectrum, SampledWavelengths, Vector3f, VectorLike, safe_sqrt, square,
};
use crate::geometry::primitives::OctahedralVector;
use crate::geometry::{DirectionCone, Normal};
use crate::utils::math::sample_discrete;
use std::collections::HashMap;
use std::sync::Arc;

use crate::core::pbrt::{ONE_MINUS_EPSILON, clamp_t, lerp};
use crate::utils::sampling::AliasTable;
use enum_dispatch::enum_dispatch;

#[derive(Clone, Copy, Debug, Default)]
#[repr(C)]
pub struct CompactLightBounds {
    pub w: OctahedralVector,
    pub phi: Float,

    // [0..15]  = qCosTheta_o
    // [15..30] = qCosTheta_e
    // [30..31] = twoSided
    // [31..32] = Unused/Padding
    packed_info: u32,

    pub qb: [[u16; 3]; 2],
}

const _: () = assert!(std::mem::size_of::<CompactLightBounds>() == 24);

impl CompactLightBounds {
    pub fn new(lb: &LightBounds, all_b: &Bounds3f) -> Self {
        let q_cos_o = Self::quantize_cos(lb.cos_theta_o);
        let q_cos_e = Self::quantize_cos(lb.cos_theta_e);
        let two_sided = if lb.two_sided { 1 } else { 0 };

        // | - twoSided (1) - | - qCosTheta_e (15) - | - qCosTheta_o (15) - |
        let packed_info = (q_cos_o & 0x7FFF) | ((q_cos_e & 0x7FFF) << 15) | (two_sided << 30);

        let mut qb = [[0u16; 3]; 2];
        for i in 0..3 {
            qb[0][i] = Self::quantize_bounds(lb.bounds.p_min[i], all_b.p_min[i], all_b.p_max[i])
                .floor() as u16;
            qb[1][i] = Self::quantize_bounds(lb.bounds.p_max[i], all_b.p_min[i], all_b.p_max[i])
                .ceil() as u16;
        }

        Self {
            w: OctahedralVector::new(lb.w.normalize()),
            phi: lb.phi,
            packed_info,
            qb,
        }
    }

    #[inline(always)]
    pub fn two_sided(&self) -> bool {
        (self.packed_info >> 30) & 1 == 1
    }

    #[inline(always)]
    fn q_cos_theta_o(&self) -> u32 {
        self.packed_info & 0x7FFF
    }

    #[inline(always)]
    fn q_cos_theta_e(&self) -> u32 {
        (self.packed_info >> 15) & 0x7FFF
    }

    #[inline]
    pub fn cos_theta_o(&self) -> Float {
        2.0 * (self.q_cos_theta_o() as Float / 32767.0) - 1.0
    }

    #[inline]
    pub fn cos_theta_e(&self) -> Float {
        2.0 * (self.q_cos_theta_e() as Float / 32767.0) - 1.0
    }

    pub fn importance(&self, p: Point3f, n: Normal3f, all_b: &Bounds3f) -> Float {
        let bounds = self.bounds(all_b);
        let cos_o = self.cos_theta_o();
        let cos_e = self.cos_theta_e();

        let pc = bounds.centroid();
        let d2 = p.distance_squared(pc).max(bounds.diagonal().norm() * 0.5);

        let cos_sub_clamped = |sin_a: Float, cos_a: Float, sin_b: Float, cos_b: Float| {
            if cos_a > cos_b {
                1.0
            } else {
                cos_a * cos_b + sin_a * sin_b
            }
        };
        let sin_sub_clamped = |sin_a: Float, cos_a: Float, sin_b: Float, cos_b: Float| {
            if cos_a > cos_b {
                0.0
            } else {
                sin_a * cos_b - cos_a * sin_b
            }
        };

        let wi = (p - pc).normalize();
        let w_vec = self.w.to_vector();
        let mut cos_w = w_vec.dot(wi);
        if self.two_sided() {
            cos_w = cos_w.abs();
        }
        let sin_w = safe_sqrt(1.0 - square(cos_w));
        let cos_b = DirectionCone::bound_subtended_directions(&bounds, p).cos_theta;
        let sin_b = safe_sqrt(1. - square(cos_b));
        let sin_o = safe_sqrt(1. - square(cos_o));
        let cos_x = cos_sub_clamped(sin_w, cos_w, sin_o, cos_o);
        let sin_x = sin_sub_clamped(sin_w, cos_w, sin_o, cos_o);
        let cos_p = cos_sub_clamped(sin_x, cos_x, sin_b, cos_b);
        if cos_p <= cos_e {
            return 0.;
        }
        let mut importance = self.phi * cos_p / d2;
        if n != Normal3f::zero() {
            let cos_i = wi.abs_dot(n.into());
            let sin_i = safe_sqrt(1. - square(cos_i));
            let cos_pi = cos_sub_clamped(sin_i, cos_i, sin_b, cos_b);
            importance *= cos_pi;
        }
        importance
    }

    pub fn bounds(&self, all_b: &Bounds3f) -> Bounds3f {
        let mut p_min = Point3f::default();
        let mut p_max = Point3f::default();

        for i in 0..3 {
            let t_min = self.qb[0][i] as Float / 65535.0;
            let t_max = self.qb[1][i] as Float / 65535.0;
            p_min[i] = lerp(t_min, all_b.p_min[i], all_b.p_max[i]);
            p_max[i] = lerp(t_max, all_b.p_min[i], all_b.p_max[i]);
        }
        Bounds3f::from_points(p_min, p_max)
    }

    fn quantize_cos(c: Float) -> u32 {
        (32767.0 * ((c + 1.0) * 0.5)).floor() as u32
    }

    fn quantize_bounds(c: Float, min: Float, max: Float) -> Float {
        if min == max {
            return 0.0;
        }
        65535.0 * clamp_t((c - min) / (max - min), 0.0, 1.0)
    }
}

#[derive(Debug, Clone)]
pub struct SampledLight {
    pub light: Arc<Light>,
    pub p: Float,
}

impl SampledLight {
    pub fn new(light: Arc<Light>, p: Float) -> Self {
        Self { light, p }
    }
}

#[enum_dispatch]
pub trait LightSamplerTrait: Send + Sync + std::fmt::Debug {
    fn sample_with_context(&self, ctx: &LightSampleContext, u: Float) -> Option<SampledLight>;
    fn pmf_with_context(&self, ctx: &LightSampleContext, light: &Arc<Light>) -> Float;
    fn sample(&self, u: Float) -> Option<SampledLight>;
    fn pmf(&self, light: &Arc<Light>) -> Float;
}

#[derive(Clone, Debug)]
#[enum_dispatch(LightSamplerTrait)]
pub enum LightSampler {
    Uniform(UniformLightSampler),
    Power(PowerLightSampler),
    BVH(BVHLightSampler),
}

#[derive(Clone, Debug)]
pub struct UniformLightSampler {
    lights: Vec<Arc<Light>>,
}

impl UniformLightSampler {
    pub fn new(lights: &[Arc<Light>]) -> Self {
        Self {
            lights: lights.to_vec(),
        }
    }
}

impl LightSamplerTrait for UniformLightSampler {
    fn sample_with_context(&self, _ctx: &LightSampleContext, u: Float) -> Option<SampledLight> {
        self.sample(u)
    }
    fn pmf_with_context(&self, _ctx: &LightSampleContext, light: &Arc<Light>) -> Float {
        self.pmf(light)
    }
    fn sample(&self, u: Float) -> Option<SampledLight> {
        if self.lights.is_empty() {
            return None;
        }

        let light_index = (u as usize * self.lights.len()).min(self.lights.len() - 1);
        Some(SampledLight {
            light: self.lights[light_index].clone(),
            p: 1. / self.lights.len() as Float,
        })
    }
    fn pmf(&self, _light: &Arc<Light>) -> Float {
        if self.lights.is_empty() {
            return 0.;
        }
        1. / self.lights.len() as Float
    }
}

#[derive(Clone, Debug)]
pub struct PowerLightSampler {
    lights: Vec<Arc<Light>>,
    light_to_index: HashMap<usize, usize>,
    alias_table: AliasTable,
}

impl PowerLightSampler {
    pub fn new(lights: &[Arc<Light>]) -> Self {
        if lights.is_empty() {
            return Self {
                lights: Vec::new(),
                light_to_index: HashMap::new(),
                alias_table: AliasTable::new(&[]),
            };
        }

        let mut lights_vec = Vec::with_capacity(lights.len());
        let mut light_to_index = HashMap::with_capacity(lights.len());
        let mut light_power = Vec::with_capacity(lights.len());

        let lambda = SampledWavelengths::sample_visible(0.5);

        for (i, light) in lights.iter().enumerate() {
            lights_vec.push(light.clone());

            let ptr = Arc::as_ptr(light) as usize;
            light_to_index.insert(ptr, i);

            let phi = SampledSpectrum::safe_div(&light.phi(lambda), &lambda.pdf());
            light_power.push(phi.average());
        }

        let alias_table = AliasTable::new(&light_power);

        Self {
            lights: lights_vec,
            light_to_index,
            alias_table,
        }
    }
}

impl LightSamplerTrait for PowerLightSampler {
    fn sample_with_context(&self, _ctx: &LightSampleContext, u: Float) -> Option<SampledLight> {
        self.sample(u)
    }

    fn pmf_with_context(&self, _ctx: &LightSampleContext, light: &Arc<Light>) -> Float {
        self.pmf(light)
    }

    fn sample(&self, u: Float) -> Option<SampledLight> {
        if self.alias_table.size() == 0 {
            return None;
        }

        let (light_index, pmf, _) = self.alias_table.sample(u);

        Some(SampledLight {
            light: self.lights[light_index].clone(),
            p: pmf,
        })
    }

    fn pmf(&self, light: &Arc<Light>) -> Float {
        if self.alias_table.size() == 0 {
            return 0.;
        }

        let ptr = Arc::as_ptr(light) as usize;

        if let Some(&index) = self.light_to_index.get(&ptr) {
            self.alias_table.pmf(index)
        } else {
            0.0
        }
    }
}

#[derive(Clone, Copy, Debug, Default)]
#[repr(C, align(32))]
pub struct LightBVHNode {
    pub light_bounds: CompactLightBounds,

    // Bit 31 (MSB) : isLeaf (1 bit)
    // Bits 0..31   : childOrLightIndex (31 bits)
    packed_data: u32,
}

const _: () = assert!(std::mem::size_of::<LightBVHNode>() == 32);

impl LightBVHNode {
    /// Mask to isolate the Leaf Flag (Bit 31)
    const LEAF_MASK: u32 = 0x8000_0000;
    /// Mask to isolate the Index (Bits 0-30)
    const INDEX_MASK: u32 = 0x7FFF_FFFF;

    pub fn make_leaf(light_index: u32, cb: CompactLightBounds) -> Self {
        debug_assert!(
            (light_index & Self::LEAF_MASK) == 0,
            "Light index too large"
        );

        Self {
            light_bounds: cb,
            // Set index and flip the MSB to 1
            packed_data: light_index | Self::LEAF_MASK,
        }
    }

    pub fn make_interior(child_index: u32, cb: CompactLightBounds) -> Self {
        debug_assert!(
            (child_index & Self::LEAF_MASK) == 0,
            "Child index too large"
        );

        Self {
            light_bounds: cb,
            // Set index, MSB remains 0
            packed_data: child_index,
        }
    }

    #[inline(always)]
    pub fn is_leaf(&self) -> bool {
        (self.packed_data & Self::LEAF_MASK) != 0
    }

    #[inline(always)]
    pub fn light_index(&self) -> u32 {
        debug_assert!(self.is_leaf());
        self.packed_data & Self::INDEX_MASK
    }

    #[inline(always)]
    pub fn child_index(&self) -> u32 {
        debug_assert!(!self.is_leaf());
        self.packed_data & Self::INDEX_MASK
    }

    #[inline(always)]
    pub fn child_or_light_index(&self) -> u32 {
        self.packed_data & Self::INDEX_MASK
    }

    pub fn sample(&self, _ctx: &LightSampleContext, _u: Float) -> Option<SampledLight> {
        todo!("Implement LightBVHNode::Sample logic")
    }
}

#[derive(Clone, Debug)]
pub struct BVHLightSampler {
    lights: Vec<Arc<Light>>,
    infinite_lights: Vec<Arc<Light>>,
    all_light_bounds: Bounds3f,
    nodes: Vec<LightBVHNode>,
    light_to_bit_trail: HashMap<usize, usize>,
}

impl BVHLightSampler {
    fn evaluate_cost(&self, b: &LightBounds, bounds: &Bounds3f, dim: usize) -> Float {
        let theta_o = b.cos_theta_o.acos();
        let theta_e = b.cos_theta_e.acos();
        let theta_w = (theta_o + theta_e).min(PI);
        let sin_o = safe_sqrt(1. - square(b.cos_theta_o));
        let m_omega = 2. * PI * (1. - b.cos_theta_o)
            + PI / 2.
                * (2. * theta_w - (theta_o - 2. * theta_w).cos() - 2. * theta_o * sin_o
                    + b.cos_theta_o);
        let kr = bounds.diagonal().max_component_value() / bounds.diagonal()[dim];
        b.phi * m_omega * kr * b.bounds.surface_area()
    }
}

impl LightSamplerTrait for BVHLightSampler {
    fn sample_with_context(&self, ctx: &LightSampleContext, mut u: Float) -> Option<SampledLight> {
        let empty_nodes = if self.nodes.is_empty() { 0. } else { 1. };
        let inf_size = self.infinite_lights.len() as Float;
        let light_size = self.lights.len() as Float;

        let p_inf = inf_size / (inf_size + empty_nodes);

        if u < p_inf {
            u /= p_inf;
            let ind = (u * light_size).min(light_size - 1.) as usize;
            let pmf = p_inf / inf_size;
            Some(SampledLight::new(self.infinite_lights[ind].clone(), pmf))
        } else {
            if self.nodes.is_empty() {
                return None;
            }
            let p = ctx.p();
            let n = ctx.ns;
            u = ((u - p_inf) / (1. - p_inf)).min(ONE_MINUS_EPSILON);
            let mut node_ind = 0;
            let mut pmf = 1. - p_inf;

            loop {
                let node = self.nodes[node_ind];
                if !node.is_leaf() {
                    let children: [LightBVHNode; 2] = [
                        self.nodes[node_ind + 1],
                        self.nodes[node.child_or_light_index() as usize],
                    ];
                    let ci: [Float; 2] = [
                        children[0]
                            .light_bounds
                            .importance(p, n, &self.all_light_bounds),
                        children[1]
                            .light_bounds
                            .importance(p, n, &self.all_light_bounds),
                    ];

                    if ci[0] == 0. && ci[1] == 0. {
                        return None;
                    }

                    let mut node_pmf: Float = 0.;
                    let child = sample_discrete(&ci, u, Some(&mut node_pmf), Some(&mut u));
                    pmf *= node_pmf;
                    node_ind = if child == 0 {
                        node_ind + 1
                    } else {
                        node.child_or_light_index() as usize
                    };
                } else {
                    if node_ind > 0
                        || node.light_bounds.importance(p, n, &self.all_light_bounds) > 0.
                    {
                        return Some(SampledLight::new(
                            self.lights[node.child_or_light_index() as usize].clone(),
                            pmf,
                        ));
                    }
                    return None;
                }
            }
        }
    }

    fn pmf_with_context(&self, ctx: &LightSampleContext, light: &Arc<Light>) -> Float {
        let ptr = Arc::as_ptr(light) as usize;
        let empty_nodes = if self.nodes.is_empty() { 0. } else { 1. };
        if self.light_to_bit_trail.contains_key(&ptr) {
            return 1. / (self.infinite_lights.len() as Float + empty_nodes);
        }

        let mut bit_trail = self.light_to_bit_trail[&ptr];
        let p = ctx.p();
        let n = ctx.ns;
        let p_inf = self.infinite_lights.len() as Float
            / (self.infinite_lights.len() as Float + empty_nodes);
        let mut pmf = 1. - p_inf;
        let mut node_ind = 0;

        loop {
            let node = self.nodes[node_ind];
            if node.is_leaf() {
                return pmf;
            }
            let child0 = self.nodes[node_ind + 1];
            let child1 = self.nodes[node.child_or_light_index() as usize];
            let ci = [
                child0.light_bounds.importance(p, n, &self.all_light_bounds),
                child1.light_bounds.importance(p, n, &self.all_light_bounds),
            ];
            pmf *= ci[bit_trail & 1] / (ci[0] + ci[1]);
            node_ind = if (bit_trail & 1) != 0 {
                node.child_or_light_index() as usize
            } else {
                node_ind + 1
            };
            bit_trail >>= 1;
        }
    }

    fn sample(&self, u: Float) -> Option<SampledLight> {
        if self.lights.is_empty() {
            return None;
        }

        let light_ind =
            (u * self.lights.len() as Float).min(self.lights.len() as Float - 1.) as usize;

        Some(SampledLight::new(
            self.lights[light_ind].clone(),
            1. / self.lights.len() as Float,
        ))
    }

    fn pmf(&self, _light: &Arc<Light>) -> Float {
        if self.lights.is_empty() {
            return 0.;
        }

        1. / self.lights.len() as Float
    }
}