//! Brush renderer integration and the virtual camera.
//!
//! Pipeline role (PLAN.md, "Pipeline per frame" — render): take the
//! freshly constructed [`SplatBatch`] each frame, hand it to Brush's
//! renderer (WebGPU via Burn) as a brand-new Gaussian list — no training,
//! no state carried between frames — and render it from the virtual
//! camera pose supplied by the UI.
//!
//! [`SplatRenderer`] owns the Burn device; [`SplatBatch`] and [`rig::Camera`]
//! types are converted to Brush's representation at the boundary:
//!
//! * scales → log scales, opacity → logit, colour → SH degree 0
//!   (Brush's conventions, see `brush_render::sh::rgb_to_sh`),
//! * intrinsics → Brush's fov + normalized center.
//!
//! The output is Brush's packed RGBA8 frame: premultiplied over
//! `background`, `[H, W, 4]` bytes, row-major.

use anyhow::Result;
use burn::tensor::Device;
use burn_wgpu::WgpuDevice;
use glam::{UVec2, Vec2, Vec3};

use brush_render::camera::Camera as BrushCamera;
use brush_render::gaussian_splats::{SplatRenderMode, Splats, TextureMode, inverse_sigmoid};
use brush_render::kernels::camera_model::CameraModel;
use brush_render::render_splats;
use brush_render::sh::rgb_to_sh;

use rig::{Camera, SplatBatch};

/// A rendered virtual-camera frame: Brush's packed output, premultiplied
/// over the renderer's background.
#[derive(Debug, Clone, PartialEq)]
pub struct RenderedImage {
    pub width: u32,
    pub height: u32,
    /// `[height][width][4]` RGBA8, premultiplied over the background.
    pub rgba: Vec<u8>,
}

impl RenderedImage {
    /// The pixel at `(x, y)`: `(r, g, b, a)`.
    pub fn pixel(&self, x: u32, y: u32) -> [u8; 4] {
        let idx = ((y * self.width + x) * 4) as usize;
        [
            self.rgba[idx],
            self.rgba[idx + 1],
            self.rgba[idx + 2],
            self.rgba[idx + 3],
        ]
    }

    /// A solid image of the background colour (alpha 0, colour
    /// premultiplied): what Brush produces for an empty Gaussian list.
    pub fn background(width: u32, height: u32, background: Vec3) -> Self {
        let px = [
            (background.x * 255.0) as u8,
            (background.y * 255.0) as u8,
            (background.z * 255.0) as u8,
            0,
        ];
        let rgba = px.repeat((width * height) as usize);
        Self {
            width,
            height,
            rgba,
        }
    }
}

/// Counting stats from a render pass, for the UI and tests.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RenderStats {
    /// Splats in the batch.
    pub total_splats: usize,
    /// Splats that survived culling for this view.
    pub visible_splats: u32,
}

/// Renders splat batches from arbitrary pinhole cameras using Brush.
pub struct SplatRenderer {
    device: Device,
    background: Vec3,
}

impl Default for SplatRenderer {
    fn default() -> Self {
        Self::new(Vec3::new(0.10, 0.10, 0.12))
    }
}

impl SplatRenderer {
    /// Create a renderer on the default Burn WGPU device.
    ///
    /// The device initializes lazily on first render; it needs a
    /// wgpu-capable GPU adapter (Vulkan; software Vulkan such as lavapipe
    /// works for tests and CI).
    pub fn new(background: Vec3) -> Self {
        Self {
            device: Device::from(WgpuDevice::default()),
            background,
        }
    }

    /// The background colour splats are composited over.
    pub fn background(&self) -> Vec3 {
        self.background
    }

    /// Render `batch` from `camera` (blocking on the GPU).
    pub fn render(
        &self,
        batch: &SplatBatch,
        camera: &Camera,
    ) -> Result<(RenderedImage, RenderStats)> {
        pollster::block_on(self.render_async(batch, camera))
    }

    /// Render `batch` from `camera`, async.
    pub async fn render_async(
        &self,
        batch: &SplatBatch,
        camera: &Camera,
    ) -> Result<(RenderedImage, RenderStats)> {
        let resolution = camera.resolution;
        if batch.is_empty() {
            return Ok((
                RenderedImage::background(resolution.x, resolution.y, self.background),
                RenderStats {
                    total_splats: 0,
                    visible_splats: 0,
                },
            ));
        }

        let splats = splats_for_brush(batch, &self.device);
        let brush_camera = brush_camera(camera);

        let (image, aux) = render_splats(
            splats,
            &brush_camera,
            UVec2::new(resolution.x, resolution.y),
            self.background,
            None,
            TextureMode::Packed,
        )
        .await;

        let stats = RenderStats {
            total_splats: batch.len(),
            visible_splats: aux.num_visible,
        };

        let data = image
            .into_data_async()
            .await
            .map_err(|e| anyhow::anyhow!("frame readback failed: {e}"))?;
        let rgba = data.into_bytes().to_vec();

        Ok((
            RenderedImage {
                width: resolution.x,
                height: resolution.y,
                rgba,
            },
            stats,
        ))
    }
}

/// Convert a rig camera to Brush's representation: fov from focal length,
/// normalized principal point, ideal pinhole.
pub fn brush_camera(camera: &Camera) -> BrushCamera {
    let res = camera.resolution;
    BrushCamera::new(
        camera.pose.position,
        camera.pose.rotation,
        camera.fov_x() as f64,
        camera.fov_y() as f64,
        Vec2::new(
            camera.intrinsics.cx / res.x as f32,
            camera.intrinsics.cy / res.y as f32,
        ),
        CameraModel::Pinhole,
    )
}

/// Flatten a batch into Brush's raw arrays (pure CPU math, no device):
/// positions, xyzw quaternions, log scales, SH degree-0 coefficients,
/// logit opacities.
pub fn brush_raw_arrays(batch: &SplatBatch) -> BrushRawArrays {
    let n = batch.len();
    let mut positions = Vec::with_capacity(n * 3);
    let mut rotations = Vec::with_capacity(n * 4);
    let mut log_scales = Vec::with_capacity(n * 3);
    let mut sh_coeffs = Vec::with_capacity(n * 3);
    let mut raw_opacities = Vec::with_capacity(n);

    for splat in batch.iter() {
        positions.extend_from_slice(&splat.mean.to_array());
        rotations.extend_from_slice(&splat.rotation.to_array());
        log_scales.extend_from_slice(&splat.scale.to_array().map(f32::ln));
        let sh = rgb_to_sh(Vec3::from(splat.color));
        sh_coeffs.extend_from_slice(&sh.to_array());
        raw_opacities.push(inverse_sigmoid(splat.opacity));
    }

    BrushRawArrays {
        positions,
        rotations,
        log_scales,
        sh_coeffs,
        raw_opacities,
    }
}

/// The flattened per-splat arrays [`brush_raw_arrays`] produces.
#[derive(Debug, Clone, PartialEq)]
pub struct BrushRawArrays {
    pub positions: Vec<f32>,
    pub rotations: Vec<f32>,
    pub log_scales: Vec<f32>,
    pub sh_coeffs: Vec<f32>,
    pub raw_opacities: Vec<f32>,
}

/// Upload a batch to the device as Brush `Splats` (forward rendering,
/// SH degree 0, default render mode).
pub fn splats_for_brush(batch: &SplatBatch, device: &Device) -> Splats {
    let arrays = brush_raw_arrays(batch);
    Splats::from_raw(
        arrays.positions,
        arrays.rotations,
        arrays.log_scales,
        arrays.sh_coeffs,
        arrays.raw_opacities,
        SplatRenderMode::Default,
        device,
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use glam::{uvec2, vec3};
    use rig::{CameraIntrinsics, CameraPose, Splat};

    fn camera() -> Camera {
        Camera {
            id: "test".into(),
            pose: CameraPose::looking_at(Vec3::ZERO, vec3(0.0, 0.0, 2.0)),
            intrinsics: CameraIntrinsics {
                fx: 500.0,
                fy: 480.0,
                cx: 320.0,
                cy: 240.0,
            },
            resolution: uvec2(640, 480),
        }
    }

    #[test]
    fn brush_camera_conversion() {
        let cam = camera();
        let bc = brush_camera(&cam);
        assert_eq!(bc.position, Vec3::ZERO);
        // fov matches Brush's focal formula (f32 round-trips, so a
        // milli-pixel tolerance is plenty tight).
        let focal_x = (640.0 / 2.0) / (bc.fov_x / 2.0).tan();
        assert!((focal_x - 500.0).abs() < 1e-3);
        let focal_y = (480.0 / 2.0) / (bc.fov_y / 2.0).tan();
        assert!((focal_y - 480.0).abs() < 1e-3);
        // Normalized principal point.
        assert!((bc.center_uv.x - 0.5).abs() < 1e-6);
        assert!((bc.center_uv.y - 0.5).abs() < 1e-6);
        // Rotation preserved (camera-to-world).
        let fwd = cam.pose.rotation * Vec3::Z;
        let bfwd = bc.rotation * Vec3::Z;
        assert!((fwd - bfwd).length() < 1e-6);
        assert!(bc.is_valid());
    }

    #[test]
    fn raw_arrays_follow_brush_conventions() {
        let mut batch = SplatBatch::new();
        batch.push(Splat {
            mean: vec3(0.1, 0.2, 2.0),
            rotation: glam::Quat::from_rotation_y(0.5),
            scale: vec3(0.004, 0.005, 0.003),
            color: [0.85, 0.62, 0.52],
            opacity: 0.95,
        });

        let a = brush_raw_arrays(&batch);
        // Position.
        assert_eq!(a.positions, vec![0.1, 0.2, 2.0]);
        // Quaternion is glam's (x, y, z, w) layout.
        assert_eq!(a.rotations, glam::Quat::from_rotation_y(0.5).to_array());
        // Log scales.
        for (got, want) in a.log_scales.iter().zip([0.004_f32, 0.005, 0.003]) {
            assert!((got - want.ln()).abs() < 1e-6);
        }
        // Logit opacity: sigmoid(raw) == 0.95.
        let raw = a.raw_opacities[0];
        let sigmoid = 1.0 / (1.0 + (-raw).exp());
        assert!((sigmoid - 0.95).abs() < 1e-6);
        // SH degree 0: colour = 0.5 + SH_C0 * coeff must reconstruct the
        // splat colour (the constant matches Brush's, cast like Brush
        // does).
        let sh_c0 = 0.28209479177387814_f64 as f32;
        for (coeff, want) in a.sh_coeffs.iter().zip([0.85_f32, 0.62, 0.52]) {
            let color = 0.5 + sh_c0 * coeff;
            assert!((color - want).abs() < 1e-5);
        }
    }

    #[test]
    fn background_image_is_transparent() {
        let img = RenderedImage::background(4, 3, Vec3::new(1.0, 0.0, 0.0));
        assert_eq!((img.width, img.height), (4, 3));
        assert_eq!(img.rgba.len(), 4 * 3 * 4);
        // Premultiplied background: colour at full value, alpha 0.
        assert_eq!(img.pixel(0, 0), [255, 0, 0, 0]);
        assert_eq!(img.pixel(3, 2), [255, 0, 0, 0]);
    }
}