use crate::{
    UiData, UiRenderState,
    render::{data::PrimitiveInstance, util::ArrBuf},
    util::{HashMap, Vec2},
};
use data::WindowUniform;
use wgpu::{
    util::{BufferInitDescriptor, DeviceExt},
    *,
};

mod atlas;
mod data;
mod page;
mod primitive;
mod texture;
mod util;

pub use atlas::*;
pub use data::{Mask, MaskIdx, MoveIdx, MoveOffset};
pub use primitive::*;

const PRELUDE: &str = include_str!("./shader/prelude.wgsl");

fn module_source(wgsl: &str) -> String {
    // The steps come from the same constants the CPU counts in, rather than
    // a second copy of them written into the shader: a grid the two disagree
    // about puts every coordinate somewhere else.
    format!(
        "const PX_STEP: f32 = 1.0 / {}.0;\nconst REL_STEP: f32 = 1.0 / {}.0;\n{PRELUDE}\n{wgsl}",
        1u32 << crate::PX_SHIFT,
        1u32 << crate::REL_SHIFT,
    )
}

pub struct UiRenderNode {
    shared_layout: BindGroupLayout,
    shared_group: BindGroup,
    format: TextureFormat,

    /// One per registered primitive, in id order.
    primitives: Vec<PrimitivePipeline>,

    layers: HashMap<usize, RenderLayer>,
    active: Vec<usize>,
    window_buffer: Buffer,
    masks: ArrBuf<Mask>,
    moves: ArrBuf<MoveOffset>,
}

struct RenderLayer {
    /// One per registered primitive, `None` where this layer draws none.
    primitives: Vec<Option<ListBuffers>>,
}

/// What draws one registered primitive.
struct PrimitivePipeline {
    data_layout: BindGroupLayout,
    pipeline: RenderPipeline,
    render: Box<dyn PrimitiveRender>,
}

/// One list's vertex buffer and the data its shader reads.
struct ListBuffers {
    instance: ArrBuf<PrimitiveInstance>,
    data: ArrBuf<u8>,
    group: Option<BindGroup>,
    /// What the primitive asked to keep per instance, if anything.
    bindings: Vec<u32>,
}

impl UiRenderNode {
    pub fn draw<'a>(&'a self, pass: &mut RenderPass<'a>) {
        pass.set_bind_group(0, &self.shared_group, &[]);
        for i in &self.active {
            let layer = &self.layers[i];
            for (id, list) in layer.primitives.iter().enumerate() {
                let Some(list) = list else { continue };
                let Some(group) = &list.group else { continue };
                let primitive = &self.primitives[id];
                pass.set_pipeline(&primitive.pipeline);
                pass.set_bind_group(1, group, &[]);
                pass.set_vertex_buffer(0, list.instance.buffer.slice(..));
                primitive.render.draw(
                    pass,
                    ListDraw {
                        instances: list.instance.len() as u32,
                        bindings: &list.bindings,
                    },
                );
            }
        }
    }

    pub fn update(
        &mut self,
        device: &Device,
        queue: &Queue,
        ui: &mut UiData,
        ui_render: &mut UiRenderState,
    ) {
        // Before the layers: each list is given its pipeline's data layout.
        self.build_pipelines(device, queue, &ui.primitives);
        self.active.clear();
        for (i, draws) in ui_render.layers.iter_mut() {
            self.active.push(i);
            for change in draws.apply_free() {
                if let Some(inst) = ui_render.active.get_mut(&change.id) {
                    for primitive in &mut inst.primitives {
                        let h = &mut primitive.handle;
                        if h.layer == i && h.kind == change.kind && h.inst_idx == change.old {
                            h.inst_idx = change.new;
                            break;
                        }
                    }
                }
            }
            let rlayer = self.layers.entry(i).or_insert_with(RenderLayer::new);
            if draws.updated {
                let lists = draws.primitives();
                // The zip would otherwise skip a list with no pipeline.
                assert!(lists.len() <= self.primitives.len());
                rlayer.primitives.resize_with(lists.len(), || None);
                for ((buffers, list), primitive) in rlayer
                    .primitives
                    .iter_mut()
                    .zip(lists)
                    .zip(&self.primitives)
                {
                    let Some(list) = list else {
                        continue;
                    };
                    buffers
                        .get_or_insert_with(|| ListBuffers::new(device))
                        .update(device, queue, primitive, list);
                }
                draws.updated = false;
            }
        }
        for primitive in &mut self.primitives {
            primitive.render.update(ui);
        }
        let mut regroup = false;
        if ui.masks.changed {
            ui.masks.changed = false;
            regroup |= self.masks.update(device, queue, &ui.masks[..]);
        }
        if ui_render.moves.changed {
            ui_render.moves.changed = false;
            regroup |= self.moves.update(device, queue, ui_render.moves.entries());
        }
        if regroup {
            self.shared_group = Self::shared_group(
                device,
                &self.shared_layout,
                &self.window_buffer,
                &self.masks,
                &self.moves,
            );
        }
    }

    pub fn resize(&mut self, size: impl Into<Vec2>, queue: &Queue) {
        let size = size.into();
        let slice = &[WindowUniform { dim: size }];
        queue.write_buffer(&self.window_buffer, 0, bytemuck::cast_slice(slice));
    }

    pub fn new(device: &Device, config: &SurfaceConfiguration) -> Self {
        let window_uniform = WindowUniform {
            dim: Vec2::new(config.width as f32, config.height as f32),
        };
        let window_buffer = device.create_buffer_init(&BufferInitDescriptor {
            label: Some("window"),
            contents: bytemuck::cast_slice(&[window_uniform]),
            usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
        });

        let shared_layout = Self::shared_layout(device);
        let masks = ArrBuf::new(
            device,
            BufferUsages::STORAGE | BufferUsages::COPY_DST,
            "ui masks",
        );
        let moves = ArrBuf::new(
            device,
            BufferUsages::STORAGE | BufferUsages::COPY_DST,
            "ui move offsets",
        );
        let shared_group =
            Self::shared_group(device, &shared_layout, &window_buffer, &masks, &moves);

        Self {
            shared_layout,
            shared_group,
            format: config.format,
            primitives: Vec::new(),
            window_buffer,
            layers: HashMap::default(),
            active: Vec::new(),
            masks,
            moves,
        }
    }

    /// Compiles a pipeline for every primitive registered since the last call.
    /// Sources only ever arrive at the end, so an id keeps its pipeline.
    fn build_pipelines(&mut self, device: &Device, queue: &Queue, registry: &PrimitiveRegistry) {
        for source in &registry.sources()[self.primitives.len()..] {
            let render = (source.render)(device, queue);
            let data_layout = Self::data_layout(device, source.stride);
            let mut groups = vec![Some(&self.shared_layout), Some(&data_layout)];
            groups.extend(render.layout().map(Some));
            let layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
                label: Some(source.label),
                bind_group_layouts: &groups,
                immediate_size: 0,
            });
            let pipeline = Self::pipeline(device, &layout, self.format, source.wgsl, source.label);
            self.primitives.push(PrimitivePipeline {
                data_layout,
                pipeline,
                render,
            });
        }
    }

    fn pipeline(
        device: &Device,
        layout: &PipelineLayout,
        format: TextureFormat,
        wgsl: &str,
        label: &str,
    ) -> RenderPipeline {
        let module = device.create_shader_module(ShaderModuleDescriptor {
            label: Some(label),
            source: ShaderSource::Wgsl(module_source(wgsl).into()),
        });
        device.create_render_pipeline(&RenderPipelineDescriptor {
            label: Some(label),
            layout: Some(layout),
            vertex: VertexState {
                module: &module,
                entry_point: Some("vs_main"),
                buffers: &[Some(PrimitiveInstance::desc())],
                compilation_options: Default::default(),
            },
            fragment: Some(FragmentState {
                module: &module,
                entry_point: Some("fs_main"),
                targets: &[Some(ColorTargetState {
                    format,
                    blend: Some(BlendState::ALPHA_BLENDING),
                    write_mask: ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: PrimitiveState {
                topology: PrimitiveTopology::TriangleStrip,
                strip_index_format: None,
                front_face: FrontFace::Cw,
                cull_mode: Some(Face::Back),
                polygon_mode: PolygonMode::Fill,
                unclipped_depth: false,
                conservative: false,
            },
            depth_stencil: None,
            multisample: MultisampleState {
                count: 1,
                mask: !0,
                alpha_to_coverage_enabled: false,
            },
            multiview_mask: None,
            cache: None,
        })
    }

    /// What every draw in the ui is given: the window, the masks and the
    /// move chain every position is resolved through.
    fn shared_layout(device: &Device) -> BindGroupLayout {
        device.create_bind_group_layout(&BindGroupLayoutDescriptor {
            entries: &[
                BindGroupLayoutEntry {
                    binding: 0,
                    visibility: ShaderStages::VERTEX | ShaderStages::FRAGMENT,
                    ty: BindingType::Buffer {
                        ty: BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: BufferSize::new(size_of::<WindowUniform>() as u64),
                    },
                    count: None,
                },
                BindGroupLayoutEntry {
                    binding: 1,
                    visibility: ShaderStages::FRAGMENT,
                    ty: BindingType::Buffer {
                        ty: BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: BufferSize::new(size_of::<Mask>() as u64),
                    },
                    count: None,
                },
                BindGroupLayoutEntry {
                    binding: 2,
                    visibility: ShaderStages::VERTEX | ShaderStages::FRAGMENT,
                    ty: BindingType::Buffer {
                        ty: BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: BufferSize::new(size_of::<MoveOffset>() as u64),
                    },
                    count: None,
                },
            ],
            label: Some("ui shared"),
        })
    }

    fn shared_group(
        device: &Device,
        layout: &BindGroupLayout,
        window: &Buffer,
        masks: &ArrBuf<Mask>,
        moves: &ArrBuf<MoveOffset>,
    ) -> BindGroup {
        device.create_bind_group(&BindGroupDescriptor {
            layout,
            entries: &[
                BindGroupEntry {
                    binding: 0,
                    resource: window.as_entire_binding(),
                },
                BindGroupEntry {
                    binding: 1,
                    resource: masks.buffer.as_entire_binding(),
                },
                BindGroupEntry {
                    binding: 2,
                    resource: moves.buffer.as_entire_binding(),
                },
            ],
            label: Some("ui shared"),
        })
    }

    /// Layout for a list of one primitive's data. Every size in the ui is
    /// stated, so "is the buffer big enough for one entry?" is answered when
    /// the bind group is made; a `None` size is wgpu's to check on every draw.
    fn data_layout(device: &Device, stride: u64) -> BindGroupLayout {
        device.create_bind_group_layout(&BindGroupLayoutDescriptor {
            entries: &[BindGroupLayoutEntry {
                binding: 0,
                visibility: ShaderStages::FRAGMENT,
                ty: BindingType::Buffer {
                    ty: BufferBindingType::Storage { read_only: true },
                    has_dynamic_offset: false,
                    min_binding_size: BufferSize::new(stride),
                },
                count: None,
            }],
            label: Some("ui primitive data"),
        })
    }
}

impl RenderLayer {
    fn new() -> Self {
        Self {
            primitives: Vec::new(),
        }
    }
}

impl ListBuffers {
    fn new(device: &Device) -> Self {
        Self {
            instance: ArrBuf::new(
                device,
                BufferUsages::VERTEX | BufferUsages::COPY_DST,
                "instance",
            ),
            data: ArrBuf::new(
                device,
                BufferUsages::STORAGE | BufferUsages::COPY_DST,
                "primitive data",
            ),
            group: None,
            bindings: Vec::new(),
        }
    }

    fn update(
        &mut self,
        device: &Device,
        queue: &Queue,
        primitive: &PrimitivePipeline,
        list: &InstanceList,
    ) {
        self.bindings.clear();
        primitive.render.instance_bindings(list, &mut self.bindings);
        self.instance.update(device, queue, list.instances());
        let resized = self.data.update(device, queue, list.data());
        if list.instances().is_empty() {
            self.group = None;
        } else if resized || self.group.is_none() {
            self.group = Some(device.create_bind_group(&BindGroupDescriptor {
                layout: &primitive.data_layout,
                entries: &[BindGroupEntry {
                    binding: 0,
                    resource: self.data.buffer.as_entire_binding(),
                }],
                label: Some("ui primitive data"),
            }));
        }
    }
}

#[cfg(test)]
mod tests {
    use super::module_source;
    use wgpu::naga::{
        front::wgsl,
        valid::{Capabilities, ValidationFlags, Validator},
    };

    /// Every shader file, composed as the renderer composes it, parses and
    /// validates with no device -- so an edit that breaks one fails here and
    /// not in the first window opened.
    #[test]
    fn every_shader_validates() {
        let dir = concat!(env!("CARGO_MANIFEST_DIR"), "/src/render/shader");
        let mut checked = 0;
        for entry in std::fs::read_dir(dir).unwrap() {
            let path = entry.unwrap().path();
            if path.extension().is_none_or(|e| e != "wgsl") || path.ends_with("prelude.wgsl") {
                continue;
            }
            let source = module_source(&std::fs::read_to_string(&path).unwrap());
            let module = wgsl::parse_str(&source)
                .unwrap_or_else(|e| panic!("{}: {}", path.display(), e.emit_to_string(&source)));
            Validator::new(ValidationFlags::all(), Capabilities::all())
                .validate(&module)
                .unwrap_or_else(|e| panic!("{}: {e:?}", path.display()));
            checked += 1;
        }
        assert!(checked > 0, "no shaders found in {dir}");
    }
}