/*
    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, TextureId,
    epaint::{Primitive, Vertex},
};
use std::{collections::HashMap, num::NonZeroU64};
use wgpu::{
    BindGroup, BindGroupLayoutDescriptor, BindGroupLayoutEntry, BindingType, BlendState, Buffer,
    BufferDescriptor, BufferUsages, Color, ColorTargetState, ColorWrites, CommandEncoder,
    CompareFunction, DepthBiasState, DepthStencilState, Device, FragmentState, FrontFace,
    IndexFormat, LoadOp, MultisampleState, Operations, PipelineCompilationOptions,
    PipelineLayoutDescriptor, PolygonMode, PrimitiveState, PrimitiveTopology, PushConstantRange,
    Queue, RenderPassColorAttachment, RenderPassDescriptor, RenderPipeline,
    RenderPipelineDescriptor, SamplerBindingType, ShaderStages, StencilState, StoreOp, Texture,
    TextureFormat, TextureSampleType, TextureView, TextureViewDimension, VertexAttribute,
    VertexBufferLayout, VertexFormat, VertexState, VertexStepMode, include_wgsl,
};

pub struct UiRenderer {
    ctx: Context,

    pipeline: RenderPipeline,

    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: &[
                        VertexAttribute {
                            format: VertexFormat::Float32x2,
                            offset: 0,
                            shader_location: 0,
                        },
                        VertexAttribute {
                            format: VertexFormat::Float32x2,
                            offset: size_of::<f32>() as u64 * 2,
                            shader_location: 1,
                        },
                        VertexAttribute {
                            format: VertexFormat::Float32x3,
                            offset: size_of::<f32>() as u64 * 4,
                            shader_location: 2,
                        },
                    ],
                }],
                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 {
                depth_compare: CompareFunction::Greater,
                depth_write_enabled: true,
                format: TextureFormat::Depth32Float,
                bias: DepthBiasState::default(),
                stencil: StencilState::default(),
            }),
            multisample: MultisampleState::default(),
            multiview: None,
            cache: None,
        });
        Self {
            ctx: Context::default(),
            pipeline,
            index,
            vertex,
            textures: HashMap::new(),
        }
    }

    pub fn create_texture() {
        // let texture = device.create_texture_with_data(
        //     &queue,
        //     &TextureDescriptor {
        //         label: None,
        //         size: Extent3d {
        //             depth_or_array_layers: 1,
        //             width,
        //             height,
        //         },
        //         mip_level_count: 1,
        //         sample_count: 1,
        //         dimension: TextureDimension::D2,
        //         format: TextureFormat::Rgba8UnormSrgb,
        //         usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
        //         view_formats: &[],
        //     },
        //     TextureDataOrder::LayerMajor,
        //     data,
        // );
        // let textureview = texture.create_view(&TextureViewDescriptor::default());
        // let sampler = device.create_sampler(&SamplerDescriptor {
        //     ..Default::default()
        // });
        // let bindgroup = device.create_bind_group(&BindGroupDescriptor {
        //     label: None,
        //     layout: &bgl,
        //     entries: &[
        //         BindGroupEntry {
        //             binding: 0,
        //             resource: BindingResource::TextureView(&textureview),
        //         },
        //         BindGroupEntry {
        //             binding: 1,
        //             resource: BindingResource::Sampler(&sampler),
        //         },
        //     ],
        // });
    }

    pub fn draw(&self, queue: &Queue, commands: &mut CommandEncoder, target: &TextureView) {
        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.button("Click me").clicked();
            });
        });

        for (texid, delta) in full_output.textures_delta.set {}

        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 {
                vertex_count += mesh.vertices.len();
                index_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.index,
                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..index_offset + (mesh.indices.len() * size_of::<u32>())]
                    .copy_from_slice(bytemuck::cast_slice(&mesh.indices));
                mapped_vertex
                    [vertex_offset..vertex_offset + (mesh.vertices.len() * size_of::<Vertex>())]
                    .copy_from_slice(bytemuck::cast_slice(&mesh.vertices));
                index_offset += mesh.indices.len();
                vertex_offset += mesh.vertices.len();
                slices.push((
                    index_offset as u32..index_offset as u32 + mesh.indices.len() as u32,
                    vertex_offset as i32,
                ));
            }
        }

        let mut rpass = commands.begin_render_pass(&RenderPassDescriptor {
            label: None,
            color_attachments: &[Some(RenderPassColorAttachment {
                view: target,
                resolve_target: None,
                ops: Operations {
                    store: StoreOp::Store,
                    load: LoadOp::Clear(Color {
                        r: 0.01,
                        g: 0.01,
                        b: 0.01,
                        a: 1.,
                    }),
                },
            })],
            ..Default::default()
        });

        rpass.set_pipeline(&self.pipeline);
        rpass.set_index_buffer(self.index.slice(..), IndexFormat::Uint32);
        rpass.set_vertex_buffer(0, self.vertex.slice(..));
        for (index, base_vertex) in slices {
            // rpass.set_bind_group(0, &self.bind_group, &[]);
            rpass.draw_indexed(index, base_vertex, 0..1);
        }
    }
}