use crate::atlas::{AtlasGrid, AtlasPage, TILE_PAD};
use crate::gpu::Gpu;
use crate::snapshot::decode_png;

#[repr(C, align(16))]
#[derive(Copy, Clone, Debug, PartialEq, bytemuck::Pod, bytemuck::Zeroable)]
pub struct IconInstance {
    pub rect_xywh_px: [f32; 4],
    pub tint_premul_rgba: [f32; 4],
    pub tile_index: u32,
    pub(crate) pad: [u32; 3],
}

impl IconInstance {
    #[must_use]
    pub const fn new(rect_xywh_px: [f32; 4], tile_index: u32, tint_premul_rgba: [f32; 4]) -> Self {
        Self {
            rect_xywh_px,
            tint_premul_rgba,
            tile_index,
            pad: [0, 0, 0],
        }
    }
}

const INSTANCE_STRIDE: u64 = core::mem::size_of::<IconInstance>() as u64;
const INITIAL_INSTANCE_CAP: u64 = 64;

#[repr(C, align(16))]
#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
struct IconFrame {
    viewport_px: [f32; 2],
    atlas_px: [f32; 2],
    grid_cols: f32,
    tile_pad: f32,
    _pad: [f32; 2],
    pages: [[f32; 4]; AtlasPage::ALL.len()],
}

const FRAME_SIZE: u64 = core::mem::size_of::<IconFrame>() as u64;

const ATLAS_16: &[u8] = include_bytes!("../../assets/icons_16.png");
const ATLAS_24: &[u8] = include_bytes!("../../assets/icons_24.png");
const ATLAS_32: &[u8] = include_bytes!("../../assets/icons_32.png");
const ATLAS_48: &[u8] = include_bytes!("../../assets/icons_48.png");

pub struct IconPipeline {
    device: wgpu::Device,
    queue: wgpu::Queue,
    pipeline: wgpu::RenderPipeline,
    uniform_buffer: wgpu::Buffer,
    bind_group: wgpu::BindGroup,
    instance_buffer: wgpu::Buffer,
    instance_capacity: u64,
    frame: IconFrame,
}

impl IconPipeline {
    #[must_use]
    pub fn new(gpu: &Gpu, color_format: wgpu::TextureFormat) -> Self {
        let device = gpu.device().clone();
        let queue = gpu.queue().clone();
        let atlas = CombinedAtlas::build();
        let (atlas_view, sampler) = atlas.upload(&device, &queue);
        let bind_group_layout = create_bind_group_layout(&device);
        let pipeline = create_pipeline(&device, &bind_group_layout, color_format);
        let uniform_buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("bone-render:icon-uniform"),
            size: FRAME_SIZE,
            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("bone-render:icon-bg"),
            layout: &bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: uniform_buffer.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: wgpu::BindingResource::TextureView(&atlas_view),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: wgpu::BindingResource::Sampler(&sampler),
                },
            ],
        });
        let instance_buffer = create_instance_buffer(&device, INITIAL_INSTANCE_CAP);
        Self {
            device,
            queue,
            pipeline,
            uniform_buffer,
            bind_group,
            instance_buffer,
            instance_capacity: INITIAL_INSTANCE_CAP,
            frame: atlas.frame(),
        }
    }

    pub fn upload(&mut self, viewport_px: [f32; 2], instances: &[IconInstance]) {
        self.frame.viewport_px = viewport_px;
        self.queue
            .write_buffer(&self.uniform_buffer, 0, bytemuck::bytes_of(&self.frame));
        if instances.is_empty() {
            return;
        }
        let needed = instances.len() as u64;
        if needed > self.instance_capacity {
            let new_cap = needed.next_power_of_two().max(self.instance_capacity * 2);
            self.instance_buffer = create_instance_buffer(&self.device, new_cap);
            self.instance_capacity = new_cap;
        }
        self.queue
            .write_buffer(&self.instance_buffer, 0, bytemuck::cast_slice(instances));
    }

    pub fn draw_range(
        &self,
        encoder: &mut wgpu::CommandEncoder,
        color_view: &wgpu::TextureView,
        range: core::ops::Range<u32>,
    ) {
        if range.start >= range.end {
            return;
        }
        let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: Some("bone-render:icon-pass"),
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: color_view,
                resolve_target: None,
                depth_slice: None,
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Load,
                    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, &self.bind_group, &[]);
        let start_bytes = u64::from(range.start) * INSTANCE_STRIDE;
        let end_bytes = u64::from(range.end) * INSTANCE_STRIDE;
        pass.set_vertex_buffer(0, self.instance_buffer.slice(start_bytes..end_bytes));
        pass.draw(0..6, 0..(range.end - range.start));
    }
}

impl core::fmt::Debug for IconPipeline {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("IconPipeline").finish_non_exhaustive()
    }
}

fn create_instance_buffer(device: &wgpu::Device, capacity: u64) -> wgpu::Buffer {
    device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("bone-render:icon-instances"),
        size: capacity * INSTANCE_STRIDE,
        usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    })
}

const INSTANCE_ATTRS: [wgpu::VertexAttribute; 3] = wgpu::vertex_attr_array![
    0 => Float32x4,
    1 => Float32x4,
    2 => Uint32,
];

fn create_bind_group_layout(device: &wgpu::Device) -> wgpu::BindGroupLayout {
    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        label: Some("bone-render:icon-bgl"),
        entries: &[
            wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false,
                    min_binding_size: wgpu::BufferSize::new(FRAME_SIZE),
                },
                count: None,
            },
            wgpu::BindGroupLayoutEntry {
                binding: 1,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Texture {
                    multisampled: false,
                    view_dimension: wgpu::TextureViewDimension::D2,
                    sample_type: wgpu::TextureSampleType::Float { filterable: true },
                },
                count: None,
            },
            wgpu::BindGroupLayoutEntry {
                binding: 2,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
                count: None,
            },
        ],
    })
}

fn create_pipeline(
    device: &wgpu::Device,
    bind_group_layout: &wgpu::BindGroupLayout,
    color_format: wgpu::TextureFormat,
) -> wgpu::RenderPipeline {
    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("bone-render:icon-shader"),
        source: wgpu::ShaderSource::Wgsl(include_str!("icon.wgsl").into()),
    });
    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: Some("bone-render:icon-layout"),
        bind_group_layouts: &[Some(bind_group_layout)],
        immediate_size: 0,
    });
    device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
        label: Some("bone-render:icon-pipeline"),
        layout: Some(&pipeline_layout),
        vertex: wgpu::VertexState {
            module: &shader,
            entry_point: Some("vs"),
            compilation_options: wgpu::PipelineCompilationOptions::default(),
            buffers: &[wgpu::VertexBufferLayout {
                array_stride: INSTANCE_STRIDE,
                step_mode: wgpu::VertexStepMode::Instance,
                attributes: &INSTANCE_ATTRS,
            }],
        },
        fragment: Some(wgpu::FragmentState {
            module: &shader,
            entry_point: Some("fs"),
            compilation_options: wgpu::PipelineCompilationOptions::default(),
            targets: &[Some(wgpu::ColorTargetState {
                format: color_format,
                blend: Some(wgpu::BlendState::PREMULTIPLIED_ALPHA_BLENDING),
                write_mask: wgpu::ColorWrites::ALL,
            })],
        }),
        primitive: wgpu::PrimitiveState {
            topology: wgpu::PrimitiveTopology::TriangleList,
            strip_index_format: None,
            front_face: wgpu::FrontFace::Ccw,
            cull_mode: None,
            polygon_mode: wgpu::PolygonMode::Fill,
            conservative: false,
            unclipped_depth: false,
        },
        depth_stencil: None,
        multisample: wgpu::MultisampleState::default(),
        multiview_mask: None,
        cache: None,
    })
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
struct PageGeom {
    origin_y: u32,
    page_px: u32,
    cell_px: u32,
    side_px: u32,
}

struct CombinedAtlas {
    width: u32,
    height: u32,
    rgba: Vec<u8>,
    geoms: [PageGeom; AtlasPage::ALL.len()],
}

fn committed_page_png(page: AtlasPage) -> &'static [u8] {
    match page {
        AtlasPage::Px16 => ATLAS_16,
        AtlasPage::Px24 => ATLAS_24,
        AtlasPage::Px32 => ATLAS_32,
        AtlasPage::Px48 => ATLAS_48,
    }
}

fn page_geoms() -> [PageGeom; AtlasPage::ALL.len()] {
    let grid = AtlasGrid::DEFAULT;
    AtlasPage::ALL
        .into_iter()
        .scan(0u32, |origin_y, page| {
            let (page_px, _) = grid.pixel_size(page);
            let geom = PageGeom {
                origin_y: *origin_y,
                page_px,
                cell_px: page.cell_px(),
                side_px: page.side_px(),
            };
            *origin_y += page_px;
            Some(geom)
        })
        .collect::<Vec<_>>()
        .try_into()
        .unwrap_or_else(|_| panic!("page_geoms must yield one entry per AtlasPage"))
}

fn decode_page(page: AtlasPage, expected_px: u32) -> Vec<u8> {
    let Ok((extent, rgba)) = decode_png(committed_page_png(page)) else {
        panic!("committed icon atlas page {page:?} failed to decode; repo state is broken");
    };
    assert_eq!(
        (extent.width().value(), extent.height().value()),
        (expected_px, expected_px),
        "committed icon atlas page {page:?} extent must equal its grid pixel size",
    );
    rgba
}

fn px_to_f32(px: u32) -> f32 {
    let Ok(small) = u16::try_from(px) else {
        panic!("atlas dimension {px} exceeds u16; icon atlas layout is broken");
    };
    f32::from(small)
}

fn page_row_texel(data: &[u8], geom: PageGeom, local_y: u32, x: u32) -> [u8; 4] {
    if x >= geom.page_px {
        return [0, 0, 0, 0];
    }
    let base = ((local_y * geom.page_px + x) * 4) as usize;
    [data[base], data[base + 1], data[base + 2], data[base + 3]]
}

impl CombinedAtlas {
    fn build() -> Self {
        let geoms = page_geoms();
        let pages: Vec<Vec<u8>> = AtlasPage::ALL
            .into_iter()
            .zip(geoms)
            .map(|(page, geom)| decode_page(page, geom.page_px))
            .collect();
        let width = geoms.iter().map(|g| g.page_px).max().unwrap_or(0);
        let height = geoms.iter().map(|g| g.page_px).sum();
        let rgba: Vec<u8> = (0..height)
            .flat_map(|y| {
                let index = geoms
                    .iter()
                    .position(|g| y >= g.origin_y && y < g.origin_y + g.page_px)
                    .unwrap_or_else(|| panic!("row {y} falls in no icon atlas page"));
                let geom = geoms[index];
                let data = &pages[index];
                let local_y = y - geom.origin_y;
                (0..width).flat_map(move |x| page_row_texel(data, geom, local_y, x))
            })
            .collect();
        Self {
            width,
            height,
            rgba,
            geoms,
        }
    }

    fn frame(&self) -> IconFrame {
        let pages = self.geoms.map(|g| {
            [
                px_to_f32(g.origin_y),
                px_to_f32(g.page_px),
                px_to_f32(g.cell_px),
                px_to_f32(g.side_px),
            ]
        });
        IconFrame {
            viewport_px: [0.0, 0.0],
            atlas_px: [px_to_f32(self.width), px_to_f32(self.height)],
            grid_cols: px_to_f32(grid_cols()),
            tile_pad: px_to_f32(TILE_PAD),
            _pad: [0.0, 0.0],
            pages,
        }
    }

    fn upload(
        &self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
    ) -> (wgpu::TextureView, wgpu::Sampler) {
        assert_eq!(
            self.rgba.len(),
            (self.width * self.height * 4) as usize,
            "combined icon atlas rgba length mismatch",
        );
        let texture = device.create_texture(&wgpu::TextureDescriptor {
            label: Some("bone-render:icon-atlas"),
            size: wgpu::Extent3d {
                width: self.width,
                height: self.height,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: wgpu::TextureFormat::Rgba8Unorm,
            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
            view_formats: &[],
        });
        queue.write_texture(
            wgpu::TexelCopyTextureInfo {
                texture: &texture,
                mip_level: 0,
                origin: wgpu::Origin3d::ZERO,
                aspect: wgpu::TextureAspect::All,
            },
            &self.rgba,
            wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(self.width * 4),
                rows_per_image: Some(self.height),
            },
            wgpu::Extent3d {
                width: self.width,
                height: self.height,
                depth_or_array_layers: 1,
            },
        );
        let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
            label: Some("bone-render:icon-sampler"),
            address_mode_u: wgpu::AddressMode::ClampToEdge,
            address_mode_v: wgpu::AddressMode::ClampToEdge,
            address_mode_w: wgpu::AddressMode::ClampToEdge,
            mag_filter: wgpu::FilterMode::Linear,
            min_filter: wgpu::FilterMode::Linear,
            mipmap_filter: wgpu::MipmapFilterMode::Nearest,
            ..Default::default()
        });
        (view, sampler)
    }
}

fn grid_cols() -> u32 {
    let Ok(cols) = u32::try_from(AtlasGrid::DEFAULT.cols()) else {
        panic!("atlas grid columns fit a u32");
    };
    cols
}

#[cfg(test)]
mod tests {
    use super::{CombinedAtlas, IconInstance, page_geoms};
    use crate::atlas::AtlasPage;

    #[test]
    fn instance_layout_is_tight() {
        assert_eq!(core::mem::size_of::<IconInstance>(), 48);
        assert_eq!(core::mem::align_of::<IconInstance>(), 16);
    }

    #[test]
    fn page_geoms_stack_without_gaps() {
        let geoms = page_geoms();
        assert_eq!(geoms[0].origin_y, 0);
        geoms.windows(2).for_each(|pair| {
            assert_eq!(pair[1].origin_y, pair[0].origin_y + pair[0].page_px);
        });
        assert_eq!(geoms[0].side_px, AtlasPage::Px16.side_px());
        assert_eq!(geoms[3].side_px, AtlasPage::Px48.side_px());
    }

    #[test]
    fn combined_atlas_dims_match_pages() {
        let atlas = CombinedAtlas::build();
        let geoms = page_geoms();
        let widest = geoms.iter().map(|g| g.page_px).max().unwrap_or(0);
        let total = geoms.iter().map(|g| g.page_px).sum::<u32>();
        assert_eq!(atlas.width, widest);
        assert_eq!(atlas.height, total);
        assert_eq!(atlas.rgba.len(), (widest * total * 4) as usize);
    }

    #[test]
    fn every_page_region_has_visible_pixels() {
        let atlas = CombinedAtlas::build();
        let geoms = page_geoms();
        let empty: Vec<usize> = geoms
            .iter()
            .enumerate()
            .filter(|(_, g)| region_alpha_sum(&atlas.rgba, atlas.width, **g) == 0)
            .map(|(i, _)| i)
            .collect();
        assert!(
            empty.is_empty(),
            "icon atlas pages rasterized empty: {empty:?}"
        );
    }

    fn region_alpha_sum(rgba: &[u8], width: u32, geom: super::PageGeom) -> u64 {
        (geom.origin_y..geom.origin_y + geom.page_px)
            .flat_map(|y| (0..geom.page_px).map(move |x| (x, y)))
            .map(|(x, y)| u64::from(rgba[((y * width + x) * 4 + 3) as usize]))
            .sum()
    }
}
