#![recursion_limit = "256"]

use std::sync::OnceLock;

use iris_core::{SHAPE_SHADER, rounded_rect_coverage, util::Vec2};
use pollster::FutureExt;
use wgpu::util::DeviceExt;

const TOP_LEFT: Vec2 = Vec2::new(10.5, 20.25);
const BOT_RIGHT: Vec2 = Vec2::new(170.75, 90.0);

const RADII: [f32; 5] = [0.0, 0.75, 8.0, 20.0, 34.0];

const GRID_W: u32 = 384;
const GRID_H: u32 = 192;
const STEP: f32 = 0.5;
const ORIGIN: Vec2 = Vec2::new(TOP_LEFT.x - 4.0, TOP_LEFT.y - 4.0);

const TOLERANCE: f32 = 1e-5;

#[test]
fn mask_sdf_matches_the_shader() {
    let gpu = Gpu::open();
    let mut worst = 0.0f32;
    let mut worst_at = (Vec2::new(0.0, 0.0), 0.0f32, 0.0f32, 0.0f32);
    let (mut inside, mut feather, mut outside) = (0u32, 0u32, 0u32);

    for radius in RADII {
        let coverage = run_shader(&gpu, radius);
        for y in 0..GRID_H {
            for x in 0..GRID_W {
                let pos = probe_at(x, y);
                let got = coverage[(y * GRID_W + x) as usize];
                let want = rounded_rect_coverage(pos, TOP_LEFT, BOT_RIGHT, radius);
                let diff = (got - want).abs();
                if diff > worst {
                    worst = diff;
                    worst_at = (pos, radius, want, got);
                }
                if got > 0.999 {
                    inside += 1;
                } else if got > 0.001 {
                    feather += 1;
                } else {
                    outside += 1;
                }
            }
        }
    }

    let (pos, radius, want, got) = worst_at;
    assert!(
        worst <= TOLERANCE,
        "shader.wgsl's rounded_rect_coverage and iris_core's disagree by {worst} at {pos:?} \
         (radius {radius}): the CPU says {want}, the GPU {got}. One of the two was edited \
         without the other -- they are transliterations and have to stay so, or a masked \
         corner stops being tappable where it is drawn.",
    );

    assert!(
        inside > 0 && feather > 0 && outside > 0,
        "the grid never crossed an edge ({inside} in, {feather} on the feather, {outside} out), \
         so agreeing proved nothing",
    );
}

fn probe_at(x: u32, y: u32) -> Vec2 {
    Vec2::new(ORIGIN.x + x as f32 * STEP, ORIGIN.y + y as f32 * STEP)
}

fn run_shader(gpu: &Gpu, radius: f32) -> Vec<f32> {
    let Gpu { device, queue, .. } = gpu;
    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("mask sdf probe"),
        source: wgpu::ShaderSource::Wgsl(probe_source().into()),
    });

    let texture = device.create_texture(&wgpu::TextureDescriptor {
        label: Some("mask sdf coverage"),
        size: wgpu::Extent3d {
            width: GRID_W,
            height: GRID_H,
            depth_or_array_layers: 1,
        },
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D2,
        format: wgpu::TextureFormat::R32Float,
        usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
        view_formats: &[],
    });
    let view = texture.create_view(&Default::default());

    let probe = Probe {
        top_left: [TOP_LEFT.x, TOP_LEFT.y],
        bot_right: [BOT_RIGHT.x, BOT_RIGHT.y],
        origin: [ORIGIN.x, ORIGIN.y],
        step: [STEP, STEP],
        radius,
        _pad: [0.0; 3],
    };
    let uniform = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: Some("mask sdf probe"),
        contents: bytemuck::bytes_of(&probe),
        usage: wgpu::BufferUsages::UNIFORM,
    });

    let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
        label: Some("mask sdf probe"),
        layout: None,
        vertex: wgpu::VertexState {
            module: &module,
            entry_point: Some("probe_vs"),
            compilation_options: Default::default(),
            buffers: &[],
        },
        fragment: Some(wgpu::FragmentState {
            module: &module,
            entry_point: Some("probe_fs"),
            compilation_options: Default::default(),
            targets: &[Some(wgpu::TextureFormat::R32Float.into())],
        }),
        primitive: Default::default(),
        depth_stencil: None,
        multisample: Default::default(),
        multiview_mask: None,
        cache: None,
    });
    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("mask sdf probe"),
        layout: &pipeline.get_bind_group_layout(0),
        entries: &[wgpu::BindGroupEntry {
            binding: 0,
            resource: uniform.as_entire_binding(),
        }],
    });

    let row_bytes = GRID_W * 4;
    assert_eq!(row_bytes % 256, 0, "GRID_W must keep rows 256-byte aligned");
    let out_size = u64::from(row_bytes) * u64::from(GRID_H);
    let read_buf = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("mask sdf readback"),
        size: out_size,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });

    let mut enc = device.create_command_encoder(&Default::default());
    {
        let mut pass = enc.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("mask sdf probe"),
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: &view,
                depth_slice: None,
                resolve_target: None,
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                    store: wgpu::StoreOp::Store,
                },
            })],
            depth_stencil_attachment: None,
            multiview_mask: None,
            timestamp_writes: None,
            occlusion_query_set: None,
        });
        pass.set_pipeline(&pipeline);
        pass.set_bind_group(0, &bind_group, &[]);
        pass.draw(0..3, 0..1);
    }
    enc.copy_texture_to_buffer(
        texture.as_image_copy(),
        wgpu::TexelCopyBufferInfo {
            buffer: &read_buf,
            layout: wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(row_bytes),
                rows_per_image: Some(GRID_H),
            },
        },
        wgpu::Extent3d {
            width: GRID_W,
            height: GRID_H,
            depth_or_array_layers: 1,
        },
    );
    queue.submit([enc.finish()]);

    let slice = read_buf.slice(..);
    slice.map_async(wgpu::MapMode::Read, |r| r.expect("mapping the readback"));
    device
        .poll(wgpu::PollType::wait_indefinitely())
        .expect("waiting for the probe");
    let mapped = slice
        .get_mapped_range()
        .expect("reading back the mapped probe buffer");
    let coverage = bytemuck::cast_slice::<u8, f32>(&mapped).to_vec();
    drop(mapped);
    read_buf.unmap();
    coverage
}

fn vulkan_instance() -> &'static wgpu::Instance {
    static INSTANCE: OnceLock<wgpu::Instance> = OnceLock::new();
    INSTANCE.get_or_init(wgpu::Instance::default)
}

struct Gpu {
    device: wgpu::Device,
    queue: wgpu::Queue,
}

impl Gpu {
    fn open() -> Self {
        let instance = vulkan_instance();
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions::default())
            .block_on()
            .expect(
                "no wgpu adapter on this machine, so the CPU/shader SDF agreement went \
                 unchecked. This VM has a virtio-gpu render node (the `this-machine-graphics` \
                 skill says what it is and how it fails); if that is gone, fix it rather \
                 than deleting this test.",
            );
        let info = adapter.get_info();
        eprintln!(
            "mask_sdf: {} ({:?}, {})",
            info.name, info.backend, info.driver
        );
        let (device, queue) = adapter
            .request_device(&wgpu::DeviceDescriptor {
                required_limits: iris_core::device_limits(),
                ..Default::default()
            })
            .block_on()
            .expect("could not get a device from the adapter");
        Self { device, queue }
    }
}

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct Probe {
    top_left: [f32; 2],
    bot_right: [f32; 2],
    origin: [f32; 2],
    step: [f32; 2],
    radius: f32,
    _pad: [f32; 3],
}

fn probe_source() -> String {
    format!(
        "{}\n{}\n\
         struct Probe {{\n\
         top_left: vec2<f32>,\n\
         bot_right: vec2<f32>,\n\
         origin: vec2<f32>,\n\
         step: vec2<f32>,\n\
         radius: f32,\n\
         }}\n\
         @group(0) @binding(0) var<uniform> probe: Probe;\n\
         @vertex\n\
         fn probe_vs(@builtin(vertex_index) vi: u32) -> @builtin(position) vec4<f32> {{\n\
         var xy = array(vec2(-1.0, -1.0), vec2(3.0, -1.0), vec2(-1.0, 3.0));\n\
         return vec4(xy[vi], 0.0, 1.0);\n\
         }}\n\
         @fragment\n\
         fn probe_fs(@builtin(position) pos: vec4<f32>) -> @location(0) f32 {{\n\
         let at = probe.origin + floor(pos.xy) * probe.step;\n\
         return rounded_rect_coverage(at, probe.top_left, probe.bot_right, probe.radius);\n\
         }}\n",
        wgsl_fn("distance_from_rect"),
        wgsl_fn("rounded_rect_coverage"),
    )
}

fn wgsl_fn(name: &str) -> &'static str {
    let start = SHAPE_SHADER
        .find(&format!("fn {name}("))
        .unwrap_or_else(|| panic!("shader.wgsl has no `fn {name}(` -- renamed, or gone"));
    let body = SHAPE_SHADER[start..]
        .find('{')
        .expect("a wgsl fn signature is followed by its body");
    let mut depth = 0usize;
    for (i, c) in SHAPE_SHADER[start + body..].char_indices() {
        match c {
            '{' => depth += 1,
            '}' => {
                depth -= 1;
                if depth == 0 {
                    return &SHAPE_SHADER[start..start + body + i + 1];
                }
            }
            _ => {}
        }
    }
    panic!("`fn {name}`'s body in shader.wgsl is never closed");
}