/*
    wearechat - generic multiplayer game with voip
    Copyright (C) 2025 metamuffin

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU Affero General Public License as published by
    the Free Software Foundation, version 3 of the License only.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU Affero General Public License for more details.

    You should have received a copy of the GNU Affero General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/
use egui::{
    Context, ImageData, TextureId,
    epaint::{Primitive, Vertex},
};
use glam::{Mat2, Mat3, Mat4};
use log::info;
use std::{collections::HashMap, num::NonZeroU64};
use wgpu::{
    AddressMode, BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayout,
    BindGroupLayoutDescriptor, BindGroupLayoutEntry, BindingResource, BindingType, BlendState,
    Buffer, BufferDescriptor, BufferUsages, ColorTargetState, ColorWrites, CommandEncoder,
    CompareFunction, DepthStencilState, Device, Extent3d, FilterMode, FragmentState, FrontFace,
    IndexFormat, LoadOp, MultisampleState, Operations, PipelineCompilationOptions,
    PipelineLayoutDescriptor, PolygonMode, PrimitiveState, PrimitiveTopology, PushConstantRange,
    Queue, RenderPassColorAttachment, RenderPassDepthStencilAttachment, RenderPassDescriptor,
    RenderPipeline, RenderPipelineDescriptor, SamplerBindingType, SamplerDescriptor, ShaderStages,
    StoreOp, Texture, TextureDescriptor, TextureDimension, TextureFormat, TextureSampleType,
    TextureUsages, TextureView, TextureViewDescriptor, TextureViewDimension, VertexBufferLayout,
    VertexState, VertexStepMode, include_wgsl,
    util::{DeviceExt, TextureDataOrder},
    vertex_attr_array,
};

pub struct UiRenderer {
    ctx: Context,

    pipeline: RenderPipeline,
    bind_group_layout: BindGroupLayout,

    index: Buffer,
    vertex: Buffer,
    textures: HashMap<TextureId, (BindGroup, Texture)>,
}

impl UiRenderer {
    pub fn new(device: &Device, format: TextureFormat) -> Self {
        let index = device.create_buffer(&BufferDescriptor {
            label: None,
            size: size_of::<f32>() as u64 * 1024 * 1024,
            usage: BufferUsages::INDEX | BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let vertex = device.create_buffer(&BufferDescriptor {
            label: None,
            size: size_of::<f32>() as u64 * 1024 * 1024,
            usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let module = device.create_shader_module(include_wgsl!("ui.wgsl"));
        let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
            entries: &[
                BindGroupLayoutEntry {
                    binding: 0,
                    count: None,
                    visibility: ShaderStages::FRAGMENT,
                    ty: BindingType::Texture {
                        sample_type: TextureSampleType::Float { filterable: true },
                        view_dimension: TextureViewDimension::D2,
                        multisampled: false,
                    },
                },
                BindGroupLayoutEntry {
                    binding: 1,
                    count: None,
                    visibility: ShaderStages::FRAGMENT,
                    ty: BindingType::Sampler(SamplerBindingType::Filtering),
                },
            ],
            label: None,
        });
        let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
            label: None,
            bind_group_layouts: &[&bind_group_layout],
            push_constant_ranges: &[PushConstantRange {
                range: 0..(4 * 4 * size_of::<f32>() as u32),
                stages: ShaderStages::VERTEX,
            }],
        });
        let pipeline = device.create_render_pipeline(&RenderPipelineDescriptor {
            label: None,
            layout: Some(&pipeline_layout),
            fragment: Some(FragmentState {
                module: &module,
                entry_point: Some("fs_main"),
                targets: &[Some(ColorTargetState {
                    blend: Some(BlendState::PREMULTIPLIED_ALPHA_BLENDING),
                    format,
                    write_mask: ColorWrites::all(),
                })],
                compilation_options: PipelineCompilationOptions::default(),
            }),
            vertex: VertexState {
                module: &module,
                entry_point: Some("vs_main"),
                buffers: &[VertexBufferLayout {
                    array_stride: size_of::<Vertex>() as u64,
                    step_mode: VertexStepMode::Vertex,
                    attributes: &vertex_attr_array![0 => Float32x2, 1 => Float32x2, 2 => Uint32],
                }],
                compilation_options: PipelineCompilationOptions::default(),
            },
            primitive: PrimitiveState {
                topology: PrimitiveTopology::TriangleList,
                front_face: FrontFace::Ccw,
                cull_mode: None, //Some(Face::Back),
                polygon_mode: PolygonMode::Fill,
                ..Default::default()
            },
            depth_stencil: Some(DepthStencilState {
                format: TextureFormat::Depth32Float,
                depth_write_enabled: false,
                depth_compare: CompareFunction::Always,
                stencil: Default::default(),
                bias: Default::default(),
            }),
            multisample: MultisampleState::default(),
            multiview: None,
            cache: None,
        });
        Self {
            ctx: Context::default(),
            pipeline,
            index,
            vertex,
            bind_group_layout,
            textures: HashMap::new(),
        }
    }

    pub fn create_texture() {}

    pub fn draw(
        &mut self,
        device: &Device,
        queue: &Queue,
        commands: &mut CommandEncoder,
        target: &TextureView,
        depth: &TextureView,
        projection: Mat4,
    ) {
        let raw_input = egui::RawInput::default();
        let full_output = self.ctx.run(raw_input, |ctx| {
            egui::CentralPanel::default().show(&ctx, |ui| {
                ui.label("Hello world!");
                ui.label("I dont think this will ever work...");
                ui.button("Click me").clicked();
            });
        });

        for (texid, delta) in full_output.textures_delta.set {
            if let Some((texbg, tex)) = self.textures.get_mut(&texid) {
                drop((texbg, tex));
                todo!()
            } else {
                assert_eq!(
                    delta.pos, None,
                    "partial update impossible; texture does not yet exist"
                );
                info!(
                    "uploading new UI texture: width={}, height={}",
                    delta.image.width(),
                    delta.image.height()
                );
                let pixels = match &delta.image {
                    ImageData::Color(color_image) => color_image.pixels.clone(),
                    ImageData::Font(font_image) => font_image.srgba_pixels(None).collect(),
                };

                let texture = device.create_texture_with_data(
                    &queue,
                    &TextureDescriptor {
                        label: None,
                        size: Extent3d {
                            depth_or_array_layers: 1,
                            width: delta.image.width() as u32,
                            height: delta.image.height() as u32,
                        },
                        mip_level_count: 1,
                        sample_count: 1,
                        dimension: TextureDimension::D2,
                        format: TextureFormat::Rgba8UnormSrgb,
                        usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
                        view_formats: &[],
                    },
                    TextureDataOrder::LayerMajor,
                    bytemuck::cast_slice::<_, u8>(&pixels),
                );
                let textureview = texture.create_view(&TextureViewDescriptor::default());
                let sampler = device.create_sampler(&SamplerDescriptor {
                    address_mode_u: AddressMode::ClampToEdge,
                    address_mode_v: AddressMode::ClampToEdge,
                    mag_filter: FilterMode::Linear,
                    min_filter: FilterMode::Linear,
                    ..Default::default()
                });
                let bindgroup = device.create_bind_group(&BindGroupDescriptor {
                    label: None,
                    layout: &self.bind_group_layout,
                    entries: &[
                        BindGroupEntry {
                            binding: 0,
                            resource: BindingResource::TextureView(&textureview),
                        },
                        BindGroupEntry {
                            binding: 1,
                            resource: BindingResource::Sampler(&sampler),
                        },
                    ],
                });
                self.textures.insert(texid, (bindgroup, texture));
            }
        }

        let clipped_primitives = self
            .ctx
            .tessellate(full_output.shapes, full_output.pixels_per_point);

        let mut index_count = 0;
        let mut vertex_count = 0;
        for p in &clipped_primitives {
            if let Primitive::Mesh(mesh) = &p.primitive {
                index_count += mesh.indices.len();
                vertex_count += mesh.vertices.len();
            }
        }

        // TODO realloc buffers if overflowing

        let mut mapped_index = queue
            .write_buffer_with(
                &self.index,
                0,
                NonZeroU64::new((size_of::<u32>() * index_count) as u64).unwrap(),
            )
            .expect("ui index buffer overflow");
        let mut mapped_vertex = queue
            .write_buffer_with(
                &self.vertex,
                0,
                NonZeroU64::new((size_of::<Vertex>() * vertex_count) as u64).unwrap(),
            )
            .expect("ui vertex buffer overflow");

        let mut index_offset = 0;
        let mut vertex_offset = 0;
        let mut slices = Vec::new();
        for p in clipped_primitives {
            if let Primitive::Mesh(mesh) = p.primitive {
                mapped_index[index_offset * size_of::<u32>()
                    ..(index_offset + mesh.indices.len()) * size_of::<u32>()]
                    .copy_from_slice(bytemuck::cast_slice(&mesh.indices));
                mapped_vertex[vertex_offset * size_of::<Vertex>()
                    ..(vertex_offset + mesh.vertices.len()) * size_of::<Vertex>()]
                    .copy_from_slice(bytemuck::cast_slice(&mesh.vertices));
                slices.push((
                    index_offset as u32..index_offset as u32 + mesh.indices.len() as u32,
                    vertex_offset as i32,
                    mesh.texture_id,
                ));
                index_offset += mesh.indices.len();
                vertex_offset += mesh.vertices.len();
            }
        }

        assert_eq!(index_count, index_offset);
        assert_eq!(vertex_count, vertex_offset);

        let mut rpass = commands.begin_render_pass(&RenderPassDescriptor {
            label: None,
            color_attachments: &[Some(RenderPassColorAttachment {
                view: target,
                resolve_target: None,
                ops: Operations {
                    load: LoadOp::Load,
                    store: StoreOp::Store,
                },
            })],
            depth_stencil_attachment: Some(RenderPassDepthStencilAttachment {
                view: &depth,
                depth_ops: Some(Operations {
                    load: LoadOp::Load,
                    store: StoreOp::Store,
                }),
                stencil_ops: None,
            }),
            ..Default::default()
        });

        let projection = projection
            * Mat4::from_mat3(Mat3::from_mat2(Mat2::from_cols_array(&[1., 0., 0., -1.])));

        let projection = projection.to_cols_array().map(|v| v.to_le_bytes());

        rpass.set_pipeline(&self.pipeline);
        rpass.set_push_constants(ShaderStages::VERTEX, 0, projection.as_flattened());
        rpass.set_index_buffer(self.index.slice(..), IndexFormat::Uint32);
        rpass.set_vertex_buffer(0, self.vertex.slice(..));
        for (index, base_vertex, texid) in slices {
            rpass.set_bind_group(0, &self.textures.get(&texid).unwrap().0, &[]);
            rpass.draw_indexed(index, base_vertex, 0..1);
        }
    }
}