scena 1.0.0

use std::sync::Arc;

use crate::assets::EnvironmentDesc;
use crate::scene::Vec3;

use super::environment_prefilter::{build_brdf_lut, prefilter_specular_cubemap_mips};
use super::pbr_contract::{PbrMaterial, environment_split_sum_contribution, reflect_vec3};

/// Number of GGX-prefiltered specular mip levels emitted for the
/// environment cubemap. Mip 0 carries the source radiance; mips 1+
/// integrate the GGX kernel at increasing roughness so the WGSL
/// specular path can sample roughness via `prefilter_mip = roughness *
/// (mip_count - 1)`.
pub(in crate::render) const PREFILTER_MIP_COUNT: u32 = 5;
/// 2D BRDF LUT resolution. The split-sum approximation indexes the LUT
/// by `(N·V, roughness)`; 64×64 is enough resolution for visually
/// smooth specular without blowing the GPU upload budget.
pub(in crate::render) const BRDF_LUT_SIZE: u32 = 64;

#[derive(Debug, Clone, PartialEq)]
pub(in crate::render) struct PreparedEnvironmentLighting {
    diffuse_rgb: Vec3,
    specular_rgb: Vec3,
    intensity: f32,
    /// Phase 1C step 1: real cubemap radiance, decoded at prepare time from
    /// the active environment asset's six face-radiance values. The `Arc`
    /// keeps `PreparedEnvironmentLighting::clone` allocation-free in the hot
    /// CPU shading loops while still letting the GPU upload consume the same
    /// pixel data without copying. The pipeline keeps a 1×1 placeholder bind
    /// when this is `None` so the GPU bind group is always well-formed.
    cubemap: Option<Arc<PreparedEnvironmentCubemap>>,
}

#[derive(Debug, Clone, PartialEq)]
pub(in crate::render) struct PreparedEnvironmentCubemap {
    pub(in crate::render) resolution: u32,
    /// Phase 1C step 2: full GGX-prefiltered specular mip chain
    /// (PREFILTER_MIP_COUNT levels). Mip 0 is the source radiance, mips
    /// 1+ are convolved with a GGX kernel at increasing roughness. Each
    /// element is six face buffers laid out RGBA32F at that mip's
    /// resolution. The CPU rasterizer reads `mips[0]` as a six-face
    /// cube; the GPU upload streams every mip per face into the
    /// `texture_cube<f32>` mip chain.
    pub(in crate::render) mips: Vec<[Vec<f32>; 6]>,
    pub(in crate::render) mip_count: u32,
    /// 2D BRDF LUT (BRDF_LUT_SIZE × BRDF_LUT_SIZE) of `(scale, bias)`
    /// pairs that drive the split-sum specular composition
    /// `prefiltered * (F0 * scale + bias)` in the WGSL fragment shader.
    pub(in crate::render) brdf_lut: Vec<f32>,
    pub(in crate::render) brdf_lut_size: u32,
}

// Visibility note: both PreparedEnvironmentLighting and
// PreparedEnvironmentCubemap declare `pub(in crate::render)` to allow the
// GPU upload path in `crate::render::gpu` to consume the prepared cubemap
// while keeping these types out of the public crate surface.

impl Default for PreparedEnvironmentLighting {
    fn default() -> Self {
        Self {
            diffuse_rgb: Vec3::ZERO,
            specular_rgb: Vec3::ZERO,
            intensity: 0.0,
            cubemap: None,
        }
    }
}

impl PreparedEnvironmentLighting {
    pub(in crate::render) fn from_environment(environment: Option<&EnvironmentDesc>) -> Self {
        let Some(environment) = environment else {
            return Self::default();
        };
        // Phase 1C step 1: parse the cubemap regardless of whether the CPU
        // shading path is going to consume scalar irradiance, so the GPU
        // pipeline can sample real per-fragment radiance. The scalar
        // diffuse/specular still come from `preview_irradiance_rgb` to keep
        // CPU rasterizer parity with the pre-Phase-1C fixtures.
        let cubemap_faces = environment.cubemap_faces();
        let cubemap = cubemap_faces.map(|faces| {
            let resolution = faces.resolution();
            let source_pixels = faces.build_face_pixels_rgba32f();
            let mips =
                prefilter_specular_cubemap_mips(&source_pixels, resolution, PREFILTER_MIP_COUNT);
            Arc::new(PreparedEnvironmentCubemap {
                resolution,
                mips,
                mip_count: PREFILTER_MIP_COUNT,
                brdf_lut: build_brdf_lut(BRDF_LUT_SIZE),
                brdf_lut_size: BRDF_LUT_SIZE,
            })
        });
        // glTF/PBR color-contract fallback: when the environment records no scalar
        // `preview_irradiance_rgb` but does carry a real cubemap (the common
        // case for bundled HDR environments), derive an average radiance from
        // the cubemap mip-0 pixels so the CPU rasterizer's PBR path can still
        // light metallic surfaces. This is a generic environment fallback, not
        // an asset-specific color calibration path.
        let irradiance = match environment.preview_irradiance_rgb() {
            Some(stored) => stored,
            None => match cubemap.as_ref() {
                Some(prepared) => average_cubemap_radiance(prepared),
                None => {
                    return Self {
                        diffuse_rgb: Vec3::ZERO,
                        specular_rgb: Vec3::ZERO,
                        intensity: 0.0,
                        cubemap,
                    };
                }
            },
        };
        let diffuse_rgb = Vec3::new(
            sanitize_environment_channel(irradiance[0]),
            sanitize_environment_channel(irradiance[1]),
            sanitize_environment_channel(irradiance[2]),
        );
        if diffuse_rgb.x <= f32::EPSILON
            && diffuse_rgb.y <= f32::EPSILON
            && diffuse_rgb.z <= f32::EPSILON
        {
            return Self {
                diffuse_rgb: Vec3::ZERO,
                specular_rgb: Vec3::ZERO,
                intensity: 0.0,
                cubemap,
            };
        }
        Self {
            diffuse_rgb,
            specular_rgb: diffuse_rgb,
            intensity: 1.0,
            cubemap,
        }
    }

    pub(in crate::render) fn cubemap(&self) -> Option<&PreparedEnvironmentCubemap> {
        self.cubemap.as_deref()
    }

    pub(in crate::render::prepare) fn is_active(&self) -> bool {
        self.intensity > 0.0
            && (self.diffuse_rgb.x > f32::EPSILON
                || self.diffuse_rgb.y > f32::EPSILON
                || self.diffuse_rgb.z > f32::EPSILON
                || self.specular_rgb.x > f32::EPSILON
                || self.specular_rgb.y > f32::EPSILON
                || self.specular_rgb.z > f32::EPSILON)
    }

    pub(in crate::render::prepare) fn gpu_diffuse_intensity(&self) -> [f32; 4] {
        [
            self.diffuse_rgb.x,
            self.diffuse_rgb.y,
            self.diffuse_rgb.z,
            self.intensity,
        ]
    }

    pub(in crate::render::prepare) fn gpu_specular_intensity(&self) -> [f32; 4] {
        [
            self.specular_rgb.x,
            self.specular_rgb.y,
            self.specular_rgb.z,
            self.intensity,
        ]
    }

    pub(in crate::render::prepare) fn pbr_contribution(
        &self,
        material: PbrMaterial,
        normal: Vec3,
        view: Vec3,
    ) -> Vec3 {
        if !self.is_active() {
            return Vec3::ZERO;
        }
        let diffuse = self
            .cubemap
            .as_deref()
            .map(|cubemap| sample_cubemap_mip(cubemap, 0, normal))
            .unwrap_or(self.diffuse_rgb);
        let reflection = reflect_vec3(Vec3::new(-view.x, -view.y, -view.z), normal);
        let prefiltered = self
            .cubemap
            .as_deref()
            .map(|cubemap| sample_prefiltered_specular(cubemap, reflection, material.roughness))
            .unwrap_or(self.specular_rgb);
        let brdf = self
            .cubemap
            .as_deref()
            .map(|cubemap| sample_brdf_lut(cubemap, dot_vec3(normal, view), material.roughness))
            .unwrap_or((1.0, 0.0));
        scale_vec3(
            environment_split_sum_contribution(material, normal, view, diffuse, prefiltered, brdf),
            self.intensity,
        )
    }
}

pub(in crate::render) fn collect_environment_lighting(
    environment: Option<&EnvironmentDesc>,
) -> PreparedEnvironmentLighting {
    PreparedEnvironmentLighting::from_environment(environment)
}

/// Average mip-0 radiance across the six cubemap faces. Used as a fallback
/// scalar irradiance for the CPU rasterizer when the asset does not record a
/// pre-baked `preview_irradiance_rgb` value. Without this, metallic surfaces
/// (where `1 − metallic = 0` cancels the diffuse term) get zero light from
/// the environment on the CPU path and render as pitch-black silhouettes.
fn average_cubemap_radiance(cubemap: &PreparedEnvironmentCubemap) -> [f32; 3] {
    let Some(faces) = cubemap.mips.first() else {
        return [0.0; 3];
    };
    let mut total = [0.0_f64; 3];
    let mut count = 0u64;
    for face in faces {
        for pixel in face.chunks_exact(4) {
            total[0] += f64::from(pixel[0]);
            total[1] += f64::from(pixel[1]);
            total[2] += f64::from(pixel[2]);
            count += 1;
        }
    }
    if count == 0 {
        return [0.0; 3];
    }
    let count = count as f64;
    [
        (total[0] / count) as f32,
        (total[1] / count) as f32,
        (total[2] / count) as f32,
    ]
}

fn sanitize_environment_channel(value: f32) -> f32 {
    if value.is_finite() {
        value.clamp(0.0, 64.0)
    } else {
        0.0
    }
}

fn sample_prefiltered_specular(
    cubemap: &PreparedEnvironmentCubemap,
    direction: Vec3,
    roughness: f32,
) -> Vec3 {
    let max_mip = cubemap.mip_count.saturating_sub(1);
    let mip = (roughness.clamp(0.0, 1.0) * max_mip as f32).round() as u32;
    sample_cubemap_mip(cubemap, mip, direction)
}

fn sample_cubemap_mip(cubemap: &PreparedEnvironmentCubemap, mip: u32, direction: Vec3) -> Vec3 {
    let Some(faces) = cubemap.mips.get(mip as usize) else {
        return Vec3::ZERO;
    };
    let resolution = (cubemap.resolution >> mip).max(1);
    let (face_index, u, v) = cubemap_face_uv(direction);
    let x = (u.clamp(0.0, 1.0) * (resolution - 1) as f32).round() as u32;
    let y = (v.clamp(0.0, 1.0) * (resolution - 1) as f32).round() as u32;
    let pixel = ((y * resolution + x) * 4) as usize;
    let face = &faces[face_index];
    if pixel + 2 >= face.len() {
        return Vec3::ZERO;
    }
    Vec3::new(face[pixel], face[pixel + 1], face[pixel + 2])
}

fn cubemap_face_uv(direction: Vec3) -> (usize, f32, f32) {
    let ax = direction.x.abs();
    let ay = direction.y.abs();
    let az = direction.z.abs();
    let (face, sc, tc, major) = if ax >= ay && ax >= az {
        if direction.x >= 0.0 {
            (0, -direction.z, -direction.y, ax)
        } else {
            (1, direction.z, -direction.y, ax)
        }
    } else if ay >= ax && ay >= az {
        if direction.y >= 0.0 {
            (2, direction.x, direction.z, ay)
        } else {
            (3, direction.x, -direction.z, ay)
        }
    } else if direction.z >= 0.0 {
        (4, direction.x, -direction.y, az)
    } else {
        (5, -direction.x, -direction.y, az)
    };
    if major <= f32::EPSILON || !major.is_finite() {
        return (4, 0.5, 0.5);
    }
    (face, 0.5 * (sc / major + 1.0), 0.5 * (tc / major + 1.0))
}

fn sample_brdf_lut(
    cubemap: &PreparedEnvironmentCubemap,
    n_dot_v: f32,
    roughness: f32,
) -> (f32, f32) {
    let size = cubemap.brdf_lut_size.max(1);
    let x = (n_dot_v.clamp(0.0, 1.0) * (size - 1) as f32).round() as u32;
    let y = (roughness.clamp(0.0, 1.0) * (size - 1) as f32).round() as u32;
    let index = ((y * size + x) * 2) as usize;
    if index + 1 >= cubemap.brdf_lut.len() {
        return (1.0, 0.0);
    }
    (cubemap.brdf_lut[index], cubemap.brdf_lut[index + 1])
}

fn dot_vec3(left: Vec3, right: Vec3) -> f32 {
    left.x * right.x + left.y * right.y + left.z * right.z
}

fn scale_vec3(value: Vec3, scale: f32) -> Vec3 {
    Vec3::new(value.x * scale, value.y * scale, value.z * scale)
}