use std::num::NonZero;

use crate::{
    layout::Ui,
    render::{data::PrimitiveInstance, texture::GpuTextures, util::ArrBuf},
    util::HashMap,
};
use data::WindowUniform;
use wgpu::{
    util::{BufferInitDescriptor, DeviceExt},
    *,
};
use winit::dpi::PhysicalSize;

mod data;
mod primitive;
mod texture;
mod util;

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

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

pub struct UiRenderer {
    uniform_group: BindGroup,
    primitive_layout: BindGroupLayout,
    rsc_layout: BindGroupLayout,
    rsc_group: BindGroup,

    pipeline: RenderPipeline,

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

struct RenderLayer {
    instance: ArrBuf<PrimitiveInstance>,
    primitives: PrimitiveBuffers,
    primitive_group: BindGroup,
}

impl UiRenderer {
    pub fn draw<'a>(&'a self, pass: &mut RenderPass<'a>) {
        pass.set_pipeline(&self.pipeline);
        pass.set_bind_group(0, &self.uniform_group, &[]);
        pass.set_bind_group(2, &self.rsc_group, &[]);
        for i in &self.active {
            let layer = &self.layers[i];
            if layer.instance.len() == 0 {
                continue;
            }
            pass.set_bind_group(1, &layer.primitive_group, &[]);
            pass.set_vertex_buffer(0, layer.instance.buffer.slice(..));
            pass.draw(0..4, 0..layer.instance.len() as u32);
        }
    }

    pub fn update(&mut self, device: &Device, queue: &Queue, ui: &mut Ui) {
        self.active.clear();
        for (i, primitives) in ui.data.layers.iter_mut() {
            self.active.push(i);
            for change in primitives.apply_free() {
                if let Some(inst) = ui.data.active.get_mut(&change.id) {
                    for h in &mut inst.primitives {
                        if h.layer == i && h.inst_idx == change.old {
                            h.inst_idx = change.new;
                            break;
                        }
                    }
                }
            }
            let rlayer = self.layers.entry(i).or_insert_with(|| {
                let primitives = PrimitiveBuffers::new(device);
                let primitive_group =
                    Self::primitive_group(device, &self.primitive_layout, primitives.buffers());
                RenderLayer {
                    instance: ArrBuf::new(
                        device,
                        BufferUsages::VERTEX | BufferUsages::COPY_DST,
                        "instance",
                    ),
                    primitives,
                    primitive_group,
                }
            });
            if primitives.updated {
                rlayer
                    .instance
                    .update(device, queue, primitives.instances());
                rlayer.primitives.update(device, queue, primitives.data());
                rlayer.primitive_group = Self::primitive_group(
                    device,
                    &self.primitive_layout,
                    rlayer.primitives.buffers(),
                );
                primitives.updated = false;
            }
        }
        let mut changed = false;
        changed |= self.textures.update(&mut ui.data.textures);
        if ui.data.masks.changed {
            ui.data.masks.changed = false;
            self.masks.update(device, queue, &ui.data.masks[..]);
            changed = true;
        }
        if changed {
            self.rsc_group = Self::rsc_group(device, &self.rsc_layout, &self.textures, &self.masks);
        }
    }

    pub fn resize(&mut self, size: &PhysicalSize<u32>, queue: &Queue) {
        let slice = &[WindowUniform {
            width: size.width as f32,
            height: size.height as f32,
        }];
        queue.write_buffer(&self.window_buffer, 0, bytemuck::cast_slice(slice));
    }

    pub fn new(
        device: &Device,
        queue: &Queue,
        config: &SurfaceConfiguration,
        limits: UiLimits,
    ) -> Self {
        let shader = device.create_shader_module(ShaderModuleDescriptor {
            label: Some("UI Shape Shader"),
            source: ShaderSource::Wgsl(SHAPE_SHADER.into()),
        });

        let window_uniform = WindowUniform::default();
        let window_buffer = device.create_buffer_init(&BufferInitDescriptor {
            label: Some("window"),
            contents: bytemuck::cast_slice(&[window_uniform]),
            usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
        });

        let uniform_layout = 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: None,
                },
                count: None,
            }],
            label: Some("window"),
        });

        let uniform_group = Self::bind_group_0(device, &uniform_layout, &window_buffer);

        let primitive_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
            entries: &core::array::from_fn::<_, { PrimitiveBuffers::LEN }, _>(|i| {
                BindGroupLayoutEntry {
                    binding: i as u32,
                    visibility: ShaderStages::FRAGMENT,
                    ty: BindingType::Buffer {
                        ty: BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                }
            }),
            label: Some("primitive"),
        });

        let tex_manager = GpuTextures::new(device, queue);
        let masks = ArrBuf::new(
            device,
            BufferUsages::STORAGE | BufferUsages::COPY_DST,
            "ui masks",
        );

        let rsc_layout = Self::rsc_layout(device, &limits);
        let rsc_group = Self::rsc_group(device, &rsc_layout, &tex_manager, &masks);

        let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
            label: Some("UI Shape Pipeline Layout"),
            bind_group_layouts: &[&uniform_layout, &primitive_layout, &rsc_layout],
            push_constant_ranges: &[],
        });
        let pipeline = device.create_render_pipeline(&RenderPipelineDescriptor {
            label: Some("UI Shape Pipeline"),
            layout: Some(&pipeline_layout),
            vertex: VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                buffers: &[PrimitiveInstance::desc()],
                compilation_options: Default::default(),
            },
            fragment: Some(FragmentState {
                module: &shader,
                entry_point: Some("fs_main"),
                targets: &[Some(ColorTargetState {
                    format: config.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: None,
            cache: None,
        });

        Self {
            uniform_group,
            primitive_layout,
            rsc_layout,
            rsc_group,
            pipeline,
            window_buffer,
            layers: HashMap::default(),
            active: Vec::new(),
            textures: tex_manager,
            masks,
        }
    }

    fn bind_group_0(
        device: &Device,
        layout: &BindGroupLayout,
        window_buffer: &Buffer,
    ) -> BindGroup {
        device.create_bind_group(&BindGroupDescriptor {
            layout,
            entries: &[BindGroupEntry {
                binding: 0,
                resource: window_buffer.as_entire_binding(),
            }],
            label: Some("ui window"),
        })
    }

    fn primitive_group(
        device: &Device,
        layout: &BindGroupLayout,
        buffers: [(u32, &Buffer); PrimitiveBuffers::LEN],
    ) -> BindGroup {
        device.create_bind_group(&BindGroupDescriptor {
            layout,
            entries: &buffers.map(|(binding, buf)| BindGroupEntry {
                binding,
                resource: buf.as_entire_binding(),
            }),
            label: Some("ui primitives"),
        })
    }

    fn rsc_layout(device: &Device, limits: &UiLimits) -> BindGroupLayout {
        device.create_bind_group_layout(&BindGroupLayoutDescriptor {
            entries: &[
                BindGroupLayoutEntry {
                    binding: 0,
                    visibility: ShaderStages::FRAGMENT,
                    ty: BindingType::Texture {
                        sample_type: TextureSampleType::Float { filterable: false },
                        view_dimension: TextureViewDimension::D2,
                        multisampled: false,
                    },
                    count: Some(NonZero::new(limits.max_textures).unwrap()),
                },
                BindGroupLayoutEntry {
                    binding: 1,
                    visibility: ShaderStages::FRAGMENT,
                    ty: BindingType::Sampler(SamplerBindingType::NonFiltering),
                    count: Some(NonZero::new(limits.max_samplers).unwrap()),
                },
                BindGroupLayoutEntry {
                    binding: 2,
                    visibility: ShaderStages::FRAGMENT,
                    ty: BindingType::Buffer {
                        ty: BufferBindingType::Storage { read_only: true },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                },
            ],
            label: Some("ui rsc"),
        })
    }

    fn rsc_group(
        device: &Device,
        layout: &BindGroupLayout,
        tex_manager: &GpuTextures,
        masks: &ArrBuf<Mask>,
    ) -> BindGroup {
        device.create_bind_group(&BindGroupDescriptor {
            layout,
            entries: &[
                BindGroupEntry {
                    binding: 0,
                    resource: BindingResource::TextureViewArray(&tex_manager.views()),
                },
                BindGroupEntry {
                    binding: 1,
                    resource: BindingResource::SamplerArray(&tex_manager.samplers()),
                },
                BindGroupEntry {
                    binding: 2,
                    resource: masks.buffer.as_entire_binding(),
                },
            ],
            label: Some("ui rsc"),
        })
    }

    pub fn view_count(&self) -> usize {
        self.textures.view_count()
    }
}

pub struct UiLimits {
    max_textures: u32,
    max_samplers: u32,
}

impl Default for UiLimits {
    fn default() -> Self {
        Self {
            max_textures: 100000,
            max_samplers: 1000,
        }
    }
}

impl UiLimits {
    pub fn max_binding_array_elements_per_shader_stage(&self) -> u32 {
        self.max_textures + self.max_samplers
    }
    pub fn max_binding_array_sampler_elements_per_shader_stage(&self) -> u32 {
        self.max_samplers
    }
}