//! GPU integration tests for the Brush render path.
//!
//! These run the real renderer headlessly: Burn's WGPU device picks the
//! best adapter available, which includes software Vulkan (lavapipe) in
//! containers and CI. Frames come from the synthetic source, and the
//! assertions check image *content* — where the subject lands, what
//! colour it is, what the backdrop looks like — rather than exact
//! pixels, so they hold on any conforming driver.

use std::sync::LazyLock;

use glam::{Vec2, uvec2, vec3};

use render::{RenderStats, SplatRenderer};
use replay::synth::{self, NoiseParams, SUBJECT};
use rig::{Camera, FrameSource, SplatBatch};
use splatgen::{SplatGen, SplatSource};

/// Small enough to stay fast on software Vulkan, big enough that the
/// subject spans several pixels.
const CENTER: (u32, u32) = (64, 53);
const SIDE: (u32, u32) = (80, 60);

fn renderer() -> &'static SplatRenderer {
    static RENDERER: LazyLock<SplatRenderer> = LazyLock::new(SplatRenderer::default);
    &RENDERER
}

/// The synthetic subject's head centre (the scene is deterministic at
/// frame 0, where the sway offset is zero).
fn head_center() -> glam::Vec3 {
    glam::vec3(0.0, 0.18, 2.0)
}

/// Synthetic skin colour, RGB in [0, 1].
const SKIN: [f32; 3] = [0.85, 0.62, 0.52];

fn first_frame_with(noise: NoiseParams) -> SplatBatch {
    let rig = synth::dev_rig_with_resolutions(uvec2(CENTER.0, CENTER.1), uvec2(SIDE.0, SIDE.1));
    let raw = synth::dev_source_with(rig.center, SUBJECT, 30.0).with_noise(noise);
    let mut source = SplatSource::new(raw, SplatGen::default());
    source.next_frame().expect("synthetic frame").splats
}

fn first_frame() -> SplatBatch {
    first_frame_with(NoiseParams::OFF)
}

/// The mean of pixels within `tolerance` of a colour, in image
/// coordinates; `None` if too few match.
fn color_centroid(image: &render::RenderedImage, color: [f32; 3], tolerance: f32) -> Option<Vec2> {
    let mut sum = Vec2::ZERO;
    let mut count = 0.0;
    for y in 0..image.height {
        for x in 0..image.width {
            let [r, g, b, _a] = image.pixel(x, y);
            let px = [r as f32 / 255.0, g as f32 / 255.0, b as f32 / 255.0];
            let dist =
                (px[0] - color[0]).abs() + (px[1] - color[1]).abs() + (px[2] - color[2]).abs();
            if dist < tolerance {
                sum += Vec2::new(x as f32, y as f32);
                count += 1.0;
            }
        }
    }
    (count > 5.0).then(|| sum / count)
}

/// A checkerboard cell colour, wherever the given pixel landed.
#[allow(dead_code)]
fn is_checker_grey(px: [u8; 4]) -> bool {
    let [r, g, b, a] = px;
    a > 0 && r == g && g == b && ((r as i32 - 184).abs() < 24 || (r as i32 - 71).abs() < 24)
}

#[tokio::test]
async fn renders_subject_and_backdrop_from_the_center() {
    let rig = synth::dev_rig_with_resolutions(uvec2(CENTER.0, CENTER.1), uvec2(SIDE.0, SIDE.1));
    let batch = first_frame();
    assert!(!batch.is_empty());

    // The virtual camera at t = 0.5 sits at the rig's center camera.
    let camera = rig.virtual_camera(0.5, rig.left.intrinsics, rig.left.resolution);
    let (image, stats) = renderer()
        .render_async(&batch, &camera)
        .await
        .expect("render");

    assert_eq!((image.width, image.height), SIDE);
    assert_eq!(image.rgba.len(), (SIDE.0 * SIDE.1 * 4) as usize);
    assert!(stats.visible_splats > 0);
    assert!(stats.visible_splats as usize <= stats.total_splats);

    // The subject's head lands where the camera math says it should.
    let projected = camera.project(head_center()).expect("in front of camera");
    let rendered = color_centroid(&image, SKIN, 0.2).expect("skin visible");
    assert!(
        (projected - rendered).length() < 3.0,
        "head projected to {projected:?} but rendered at {rendered:?}"
    );

    // The image corners show the checkerboard backdrop, not the subject:
    // grey, opaque, and neither skin nor the renderer background. (Exact
    // cell colours are not asserted — a corner can legitimately blend a
    // checker boundary.)
    let corners = [
        image.pixel(2, 2),
        image.pixel(image.width - 3, 2),
        image.pixel(2, image.height - 3),
        image.pixel(image.width - 3, image.height - 3),
    ];
    for corner in corners {
        let [r, g, b, a] = corner;
        assert!(a >= 250, "backdrop is opaque, got {corner:?}");
        assert!(
            (r as i32 - g as i32).abs() <= 1 && (g as i32 - b as i32).abs() <= 1,
            "backdrop is grey, got {corner:?}"
        );
        assert!(
            (40..215).contains(&r),
            "backdrop is a checker grey, got {corner:?}"
        );
    }

    // The head is opaque: stacked ~0.95-opacity surfels saturate alpha
    // (within the rasterizer's 1e-4 transmittance floor).
    let head_px = image.pixel(projected.x as u32, projected.y as u32);
    assert!(head_px[3] >= 250, "head pixel must be opaque: {head_px:?}");
}

#[tokio::test]
async fn steering_shifts_the_backdrop_but_not_the_subject() {
    // The rig converges on the subject: at the convergence distance the
    // parallax is zero, so the subject stays framed as the virtual camera
    // moves between the side cameras — the point of the whole design —
    // while everything off the convergence plane (the backdrop) shifts.
    let rig = synth::dev_rig_with_resolutions(uvec2(CENTER.0, CENTER.1), uvec2(SIDE.0, SIDE.1));
    let batch = first_frame();

    let render_at = |t: f32| {
        let camera = rig.virtual_camera(t, rig.left.intrinsics, rig.left.resolution);
        let batch = &batch;
        async move { renderer().render_async(batch, &camera).await }
    };

    let (left, _) = render_at(0.0).await.expect("render left");
    let (right, _) = render_at(1.0).await.expect("render right");

    // The rendered head aligns with the camera math in both views…
    let left_head = color_centroid(&left, SKIN, 0.2).expect("skin from left");
    let right_head = color_centroid(&right, SKIN, 0.2).expect("skin from right");
    let cam_left = rig.virtual_camera(0.0, rig.left.intrinsics, rig.left.resolution);
    let cam_right = rig.virtual_camera(1.0, rig.left.intrinsics, rig.left.resolution);
    let proj_left = cam_left.project(head_center()).unwrap();
    let proj_right = cam_right.project(head_center()).unwrap();

    assert!((proj_left - left_head).length() < 3.0);
    assert!((proj_right - right_head).length() < 3.0);

    // …and the convergence geometry keeps the subject stable across the
    // full steering range…
    assert!(
        (left_head - right_head).length() < 3.0,
        "subject must stay framed: {left_head:?} vs {right_head:?}"
    );

    // …while the backdrop parallaxes: a good fraction of the frame must
    // change between the two ends of the baseline.
    let mut changed = 0;
    let mut total = 0;
    let (left_px, _) = left.rgba.as_chunks::<4>();
    let (right_px, _) = right.rgba.as_chunks::<4>();
    for (a, b) in left_px.iter().zip(right_px) {
        total += 1;
        let diff = (a[0] as i32 - b[0] as i32)
            .abs()
            .max((a[1] as i32 - b[1] as i32).abs());
        if diff > 25 {
            changed += 1;
        }
    }
    assert!(
        changed * 10 > total,
        "only {changed}/{total} pixels changed between t=0 and t=1 — no parallax"
    );
}

#[tokio::test]
async fn noisy_frames_render_like_clean_ones() {
    // The realistic path: sensor noise on, and the splatgen
    // preprocessing (flier removal, bilateral filter) has to clean it
    // up well enough that the same content assertions hold.
    let rig = synth::dev_rig_with_resolutions(uvec2(CENTER.0, CENTER.1), uvec2(SIDE.0, SIDE.1));
    let batch = first_frame_with(NoiseParams::default());
    assert!(!batch.is_empty());

    let camera = rig.virtual_camera(0.5, rig.left.intrinsics, rig.left.resolution);
    let (image, stats) = renderer()
        .render_async(&batch, &camera)
        .await
        .expect("render");
    assert!(stats.visible_splats > 0);

    // The subject still lands where the camera math says, within a
    // slightly looser tolerance than the clean-frame test.
    let projected = camera.project(head_center()).expect("in front of camera");
    let rendered = color_centroid(&image, SKIN, 0.2).expect("skin visible");
    assert!(
        (projected - rendered).length() < 4.0,
        "head projected to {projected:?} but rendered at {rendered:?}"
    );

    // And the backdrop corners are still grey and opaque.
    let corner = image.pixel(2, 2);
    let [r, g, b, a] = corner;
    assert!(a >= 250, "backdrop opaque, got {corner:?}");
    assert!((r as i32 - g as i32).abs() <= 1 && (g as i32 - b as i32).abs() <= 1);
}

#[tokio::test]
async fn empty_batch_renders_the_background() {
    let rig = synth::dev_rig_with_resolutions(uvec2(CENTER.0, CENTER.1), uvec2(SIDE.0, SIDE.1));
    let camera: Camera = rig.virtual_camera(0.5, rig.left.intrinsics, rig.left.resolution);
    let (image, stats) = renderer()
        .render_async(&SplatBatch::new(), &camera)
        .await
        .expect("render");
    assert_eq!(
        stats,
        RenderStats {
            total_splats: 0,
            visible_splats: 0
        }
    );
    assert_eq!((image.width, image.height), SIDE);
    // Every pixel is the (transparent) background.
    for y in 0..image.height {
        for x in 0..image.width {
            assert_eq!(image.pixel(x, y)[3], 0);
        }
    }
}

#[tokio::test]
async fn splats_behind_the_camera_are_culled() {
    let rig = synth::dev_rig_with_resolutions(uvec2(CENTER.0, CENTER.1), uvec2(SIDE.0, SIDE.1));
    let batch = first_frame();
    // A virtual camera looking away from the subject: everything ends up
    // behind the near plane.
    let mut camera = rig.virtual_camera(0.5, rig.left.intrinsics, rig.left.resolution);
    camera.pose = rig::CameraPose::looking_at(
        SUBJECT + vec3(0.0, 0.0, 2.0),
        SUBJECT + vec3(0.0, 0.0, 10.0),
    );
    let (image, stats) = renderer()
        .render_async(&batch, &camera)
        .await
        .expect("render");
    assert_eq!(stats.total_splats, batch.len());
    assert_eq!(stats.visible_splats, 0, "everything is behind the camera");
    // The image is all background.
    assert!(image.rgba.as_chunks::<4>().0.iter().all(|px| px[3] == 0));
}