//! Rasterising local-adjustment masks (ARCH §5.4).
//!
//! Turns a [`MaskStack`]'s rules into an r8unorm texture array, one slice per
//! active layer, which the composed adjust shader samples. Nothing here reads
//! back, and no mask ever exists in CPU memory.
//!
//! # What runs when
//!
//! Rasterising is **not** on the slider path. Dragging exposure on a masked
//! layer changes uniforms only; the mask array is reused untouched. This pass
//! runs when a mask's *shape* changes — a different selection, a moved
//! gradient, a resized output — which is what keeps a local adjustment as
//! responsive as a global one.
//!
//! # The label field
//!
//! Region masks index a compacted label field uploaded once per segmentation.
//! Compacted, rather than the watershed's raw basin roots, because a root is a
//! sparse index into pixel space: indexing a per-region array by one would
//! need a table the size of the image, where compacted ids index an array of
//! `region_count`. The compaction is CPU-side and once per image, which is the
//! same place and cadence the region adjacency graph is already built at.

use dr_pipeline::mask::{MaskSource, MaskStack, MAX_LAYERS};
use wgpu::util::DeviceExt;

use crate::{GpuContext, GpuError};

/// Modes understood by `mask.wgsl`. Kept beside the shader's `switch`.
const MODE_REGIONS: u32 = 0;
const MODE_LINEAR: u32 = 1;
const MODE_RADIAL: u32 = 2;

#[repr(C)]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct MaskParams {
    width: u32,
    height: u32,
    label_width: u32,
    label_height: u32,

    mode: u32,
    region_count: u32,
    feather: f32,
    _pad0: f32,

    centre: [f32; 2],
    axis: [f32; 2],
    softness: f32,
    angle: f32,
    _pad1: [f32; 2],
}

/// The segmentation a region mask indexes into, resident on the GPU.
///
/// Uploaded once per image. Holds the compacted label field and nothing else —
/// the hierarchy that produced the ids stays on the CPU, where the interactive
/// operations (walk up a level, add a region) are cheap graph work.
pub struct LabelField {
    buffer: wgpu::Buffer,
    width: u32,
    height: u32,
    region_count: u32,
}

impl LabelField {
    /// Upload a compacted label field.
    ///
    /// `labels` is one region id per pixel, every value below `region_count` —
    /// exactly [`dr_segment::RegionField::labels`].
    pub fn upload(
        ctx: &GpuContext,
        labels: &[u32],
        width: u32,
        height: u32,
        region_count: u32,
    ) -> Result<Self, GpuError> {
        if labels.len() != (width * height) as usize {
            return Err(GpuError::InvalidMask(format!(
                "label field is {} entries, expected {}x{}",
                labels.len(),
                width,
                height
            )));
        }

        let buffer = ctx
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("mask-labels"),
                contents: bytemuck::cast_slice(labels),
                usage: wgpu::BufferUsages::STORAGE,
            });

        Ok(Self {
            buffer,
            width,
            height,
            region_count,
        })
    }

    pub fn region_count(&self) -> u32 {
        self.region_count
    }

    pub fn size(&self) -> (u32, u32) {
        (self.width, self.height)
    }
}

/// The rasterised masks for one edit.
pub struct MaskArray {
    texture: wgpu::Texture,
    view: wgpu::TextureView,
    width: u32,
    height: u32,
    layers: u32,
}

impl MaskArray {
    pub const FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::R8Unorm;

    /// The view the adjust shader binds at `@binding(3)`.
    pub fn view(&self) -> &wgpu::TextureView {
        &self.view
    }

    pub fn layers(&self) -> u32 {
        self.layers
    }

    pub fn size(&self) -> (u32, u32) {
        (self.width, self.height)
    }

    fn matches(&self, width: u32, height: u32, layers: u32) -> bool {
        self.width == width && self.height == height && self.layers == layers
    }
}

/// Rasterises mask layers.
pub struct MaskPass {
    ctx: GpuContext,
    layout: wgpu::BindGroupLayout,
    pipeline: wgpu::RenderPipeline,
    array: Option<MaskArray>,
    /// How many times the array texture has been (re)allocated.
    ///
    /// Exists to be asserted on. Reallocating per frame instead of per resize
    /// is the kind of regression that costs a lot of bandwidth and shows up
    /// nowhere in the output, so the cheap reuse path is worth a test that
    /// can actually see it.
    allocations: usize,
    /// A one-region, always-unselected field, for a stack with no region mask.
    ///
    /// The shader's bindings are fixed, so *something* must be bound at the
    /// label slots even when rasterising a gradient. A placeholder is cheaper
    /// and far simpler than two pipelines differing only in what they ignore.
    placeholder: LabelField,
}

impl MaskPass {
    pub fn new(ctx: &GpuContext) -> Result<Self, GpuError> {
        let scope = ctx.device.push_error_scope(wgpu::ErrorFilter::Validation);

        let module = ctx
            .device
            .create_shader_module(wgpu::ShaderModuleDescriptor {
                label: Some("mask"),
                source: wgpu::ShaderSource::Wgsl(include_str!("shaders/mask.wgsl").into()),
            });

        let layout = ctx
            .device
            .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("mask-bgl"),
                entries: &[uniform_entry(0), storage_entry(1), storage_entry(2)],
            });

        let pipeline_layout = ctx
            .device
            .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("mask-layout"),
                bind_group_layouts: &[Some(&layout)],
                immediate_size: 0,
            });

        let pipeline = ctx
            .device
            .create_render_pipeline(&wgpu::RenderPipelineDescriptor {
                label: Some("mask-pipeline"),
                layout: Some(&pipeline_layout),
                vertex: wgpu::VertexState {
                    module: &module,
                    entry_point: Some("vs"),
                    compilation_options: Default::default(),
                    buffers: &[],
                },
                fragment: Some(wgpu::FragmentState {
                    module: &module,
                    entry_point: Some("fs"),
                    compilation_options: Default::default(),
                    targets: &[Some(MaskArray::FORMAT.into())],
                }),
                primitive: wgpu::PrimitiveState::default(),
                depth_stencil: None,
                multisample: wgpu::MultisampleState::default(),
                multiview_mask: None,
                cache: None,
            });

        if let Some(err) = pollster::block_on(scope.pop()) {
            return Err(GpuError::ShaderCompilation(err.to_string()));
        }

        let placeholder = LabelField::upload(ctx, &[0], 1, 1, 0)?;

        Ok(Self {
            ctx: ctx.clone(),
            layout,
            pipeline,
            array: None,
            allocations: 0,
            placeholder,
        })
    }

    /// Rasterise every active layer, returning the array to bind.
    ///
    /// `labels` may be `None` when no layer is a region mask; a region layer
    /// without one is skipped rather than drawn wrong, since a mask that
    /// silently covers the whole frame would apply an edit everywhere.
    pub fn render(
        &mut self,
        stack: &MaskStack,
        labels: Option<&LabelField>,
        width: u32,
        height: u32,
    ) -> Result<&MaskArray, GpuError> {
        // At least one layer, because a zero-layer texture array is invalid
        // and the shader binds this slot unconditionally.
        let active = stack.active_count().clamp(1, MAX_LAYERS) as u32;
        self.ensure_array(width, height, active)?;

        let mut encoder = self
            .ctx
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("mask-encoder"),
            });

        for (slot, layer) in stack.active().enumerate().take(MAX_LAYERS) {
            let field = match (&layer.source, labels) {
                (MaskSource::Regions { .. }, None) => {
                    log::warn!(
                        "mask layer {} is a region mask with no segmentation loaded; skipping",
                        layer.id
                    );
                    continue;
                }
                (MaskSource::Regions { .. }, Some(f)) => f,
                (_, _) => &self.placeholder,
            };

            let params = self.params(layer, field, width, height);
            let selected = self.selection_buffer(layer, field);
            self.draw(&mut encoder, slot as u32, &params, field, &selected);
        }

        self.ctx.queue.submit([encoder.finish()]);
        Ok(self.array.as_ref().expect("array was just ensured"))
    }

    /// The currently rasterised array, if any.
    pub fn array(&self) -> Option<&MaskArray> {
        self.array.as_ref()
    }

    /// How many times the array texture has been allocated. For tests.
    pub fn allocations(&self) -> usize {
        self.allocations
    }

    fn params(
        &self,
        layer: &dr_pipeline::mask::MaskLayer,
        field: &LabelField,
        width: u32,
        height: u32,
    ) -> MaskParams {
        let base = MaskParams {
            width,
            height,
            label_width: field.width,
            label_height: field.height,
            mode: MODE_REGIONS,
            region_count: field.region_count,
            feather: 0.0,
            _pad0: 0.0,
            centre: [0.5, 0.5],
            axis: [1.0, 0.0],
            softness: 0.0,
            angle: 0.0,
            _pad1: [0.0, 0.0],
        };

        match &layer.source {
            MaskSource::Regions { .. } => MaskParams {
                // A pixel of softening at the proxy-to-output ratio, so the
                // edge is equally soft whatever size the render is.
                feather: (width as f32 / field.width.max(1) as f32).clamp(0.0, 4.0),
                ..base
            },
            MaskSource::Linear {
                centre,
                angle,
                width: ramp,
            } => MaskParams {
                mode: MODE_LINEAR,
                centre: [centre.0, centre.1],
                axis: [angle.cos(), angle.sin()],
                softness: *ramp,
                ..base
            },
            MaskSource::Radial {
                centre,
                radii,
                angle,
                feather,
            } => MaskParams {
                mode: MODE_RADIAL,
                centre: [centre.0, centre.1],
                axis: [radii.0.max(1e-6), radii.1.max(1e-6)],
                softness: *feather,
                angle: *angle,
                ..base
            },
        }
    }

    /// One byte-flag per region, or a single zero for a non-region layer.
    fn selection_buffer(
        &self,
        layer: &dr_pipeline::mask::MaskLayer,
        field: &LabelField,
    ) -> wgpu::Buffer {
        let mut flags = vec![0u32; field.region_count.max(1) as usize];
        if let MaskSource::Regions { ids, .. } = &layer.source {
            for &id in ids {
                if let Some(slot) = flags.get_mut(id as usize) {
                    *slot = 1;
                }
            }
        }

        self.ctx
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("mask-selection"),
                contents: bytemuck::cast_slice(&flags),
                usage: wgpu::BufferUsages::STORAGE,
            })
    }

    #[allow(clippy::too_many_arguments)]
    fn draw(
        &self,
        encoder: &mut wgpu::CommandEncoder,
        slot: u32,
        params: &MaskParams,
        field: &LabelField,
        selected: &wgpu::Buffer,
    ) {
        let params_buf = self
            .ctx
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("mask-params"),
                contents: bytemuck::bytes_of(params),
                usage: wgpu::BufferUsages::UNIFORM,
            });

        let bind_group = self
            .ctx
            .device
            .create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("mask-bind"),
                layout: &self.layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: params_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: field.buffer.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: selected.as_entire_binding(),
                    },
                ],
            });

        // The array slice is selected by the attachment rather than by a
        // uniform the shader reads — one fewer value that can disagree with
        // where the pass actually writes.
        let array = self.array.as_ref().expect("array ensured by caller");
        let view = array.texture.create_view(&wgpu::TextureViewDescriptor {
            label: Some("mask-slice"),
            dimension: Some(wgpu::TextureViewDimension::D2),
            base_array_layer: slot,
            array_layer_count: Some(1),
            ..Default::default()
        });

        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("mask-pass"),
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: &view,
                depth_slice: None,
                resolve_target: None,
                ops: wgpu::Operations {
                    // Cleared rather than loaded: every pixel is written by the
                    // triangle below, and declaring that lets a tiler skip
                    // reading the previous contents in.
                    load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
                    store: wgpu::StoreOp::Store,
                },
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            occlusion_query_set: None,
            multiview_mask: None,
        });
        pass.set_pipeline(&self.pipeline);
        pass.set_bind_group(0, &bind_group, &[]);
        pass.draw(0..3, 0..1);
    }

    fn ensure_array(&mut self, width: u32, height: u32, layers: u32) -> Result<(), GpuError> {
        if self.array.as_ref().is_some_and(|a| a.matches(width, height, layers)) {
            return Ok(());
        }

        let texture = self.ctx.device.create_texture(&wgpu::TextureDescriptor {
            label: Some("mask-array"),
            size: wgpu::Extent3d {
                width,
                height,
                depth_or_array_layers: layers,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: MaskArray::FORMAT,
            usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::TEXTURE_BINDING,
            view_formats: &[],
        });

        let view = texture.create_view(&wgpu::TextureViewDescriptor {
            label: Some("mask-array-view"),
            dimension: Some(wgpu::TextureViewDimension::D2Array),
            ..Default::default()
        });

        self.allocations += 1;
        self.array = Some(MaskArray {
            texture,
            view,
            width,
            height,
            layers,
        });
        Ok(())
    }
}

fn uniform_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
    wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::FRAGMENT,
        ty: wgpu::BindingType::Buffer {
            ty: wgpu::BufferBindingType::Uniform,
            has_dynamic_offset: false,
            min_binding_size: None,
        },
        count: None,
    }
}

fn storage_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
    wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::FRAGMENT,
        ty: wgpu::BindingType::Buffer {
            ty: wgpu::BufferBindingType::Storage { read_only: true },
            has_dynamic_offset: false,
            min_binding_size: None,
        },
        count: None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use dr_pipeline::descriptor::ParamId;
    use dr_pipeline::mask::MaskLayer;

    fn ctx() -> Option<GpuContext> {
        pollster::block_on(GpuContext::new_headless()).ok()
    }

    /// A 4x2 label field: regions 0 and 1 left, 2 and 3 right.
    fn labels() -> (Vec<u32>, u32, u32, u32) {
        (vec![0, 0, 2, 2, 1, 1, 3, 3], 4, 2, 4)
    }

    fn lit(source: MaskSource) -> MaskLayer {
        let mut layer = MaskLayer::new("m1", source);
        layer.set_param("exposure", ParamId("exposure"), 1.0);
        layer
    }

    #[test]
    fn a_label_field_of_the_wrong_size_is_rejected() {
        let Some(ctx) = ctx() else {
            eprintln!("no adapter; skipping");
            return;
        };
        assert!(LabelField::upload(&ctx, &[0, 1, 2], 4, 2, 4).is_err());
    }

    #[test]
    fn region_masks_rasterise_to_the_selected_regions() {
        let Some(ctx) = ctx() else {
            eprintln!("no adapter; skipping");
            return;
        };
        let (data, w, h, n) = labels();
        let field = LabelField::upload(&ctx, &data, w, h, n).expect("upload");

        let mut stack = MaskStack::new();
        stack.push(lit(MaskSource::Regions {
            signature: 1,
            level: 4,
            ids: vec![0, 1],
        }));

        let mut pass = MaskPass::new(&ctx).expect("mask pass");
        let array = pass.render(&stack, Some(&field), w, h).expect("render");
        assert_eq!(array.size(), (w, h));
        assert_eq!(array.layers(), 1);
    }

    /// A region layer with no segmentation must produce nothing rather than
    /// an all-covering mask, which would apply the edit to the whole frame.
    #[test]
    fn a_region_layer_without_labels_is_skipped() {
        let Some(ctx) = ctx() else {
            eprintln!("no adapter; skipping");
            return;
        };
        let mut stack = MaskStack::new();
        stack.push(lit(MaskSource::Regions {
            signature: 1,
            level: 4,
            ids: vec![0],
        }));

        let mut pass = MaskPass::new(&ctx).expect("mask pass");
        assert!(pass.render(&stack, None, 8, 8).is_ok());
    }

    #[test]
    fn gradients_need_no_segmentation() {
        let Some(ctx) = ctx() else {
            eprintln!("no adapter; skipping");
            return;
        };
        let mut stack = MaskStack::new();
        stack.push(lit(MaskSource::Linear {
            centre: (0.5, 0.5),
            angle: 0.0,
            width: 0.2,
        }));
        stack.push(lit(MaskSource::Radial {
            centre: (0.5, 0.5),
            radii: (0.3, 0.2),
            angle: 0.0,
            feather: 0.5,
        }));

        let mut pass = MaskPass::new(&ctx).expect("mask pass");
        let array = pass.render(&stack, None, 16, 16).expect("render");
        assert_eq!(array.layers(), 2, "one slice per active layer");
    }

    #[test]
    fn an_empty_stack_still_yields_a_bindable_array() {
        let Some(ctx) = ctx() else {
            eprintln!("no adapter; skipping");
            return;
        };
        let mut pass = MaskPass::new(&ctx).expect("mask pass");
        let array = pass
            .render(&MaskStack::new(), None, 8, 8)
            .expect("render");
        assert_eq!(
            array.layers(),
            1,
            "the adjust shader binds this slot whether or not it reads it"
        );
    }

    #[test]
    fn the_array_is_reused_when_nothing_changed() {
        let Some(ctx) = ctx() else {
            eprintln!("no adapter; skipping");
            return;
        };
        let mut stack = MaskStack::new();
        stack.push(lit(MaskSource::Linear {
            centre: (0.5, 0.5),
            angle: 0.0,
            width: 0.2,
        }));

        let mut pass = MaskPass::new(&ctx).expect("mask pass");
        pass.render(&stack, None, 32, 32).expect("render");
        assert_eq!(pass.allocations(), 1);

        pass.render(&stack, None, 32, 32).expect("render");
        assert_eq!(
            pass.allocations(),
            1,
            "same size and layer count should not reallocate"
        );

        pass.render(&stack, None, 64, 64).expect("render");
        assert_eq!(pass.allocations(), 2, "a resize must reallocate");
    }
}