/*
    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, ViewportBuilder, ViewportId, ViewportInfo,
    epaint::{Primitive, Vertex},
};
use glam::{Affine3A, Mat2, Mat3, Mat4};
use log::info;
use std::{collections::HashMap, num::NonZeroU64, sync::Arc};
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,
    ImageCopyTexture, ImageDataLayout, IndexFormat, LoadOp, MultisampleState, Operations, Origin3d,
    PipelineCompilationOptions, PipelineLayoutDescriptor, PolygonMode, PrimitiveState,
    PrimitiveTopology, PushConstantRange, Queue, RenderPassColorAttachment,
    RenderPassDepthStencilAttachment, RenderPassDescriptor, RenderPipeline,
    RenderPipelineDescriptor, SamplerBindingType, SamplerDescriptor, ShaderStages, StoreOp,
    Texture, TextureAspect, 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, [u32; 2])>,

    pub surfaces: HashMap<ViewportId, UiSurface>,
}

pub struct UiSurface {
    pub transform: Affine3A,
    pub content: Arc<dyn Fn(&Context) + Send + Sync + 'static>,
}

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(),
            surfaces: 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 mut raw_input = egui::RawInput::default();
        raw_input.viewport_id = self
            .surfaces
            .keys()
            .next()
            .copied()
            .unwrap_or(ViewportId::ROOT);
        raw_input.viewports = self
            .surfaces
            .keys()
            .chain(Some(ViewportId::ROOT).iter())
            .map(|key| {
                (*key, ViewportInfo {
                    native_pixels_per_point: Some(2.),
                    ..Default::default()
                })
            })
            .collect();

        let full_output = self.ctx.run(raw_input, |ctx| {
            for (id, surf) in &self.surfaces {
                ctx.show_viewport_immediate(*id, ViewportBuilder::default(), |ctx, _class| {
                    (surf.content)(ctx)
                });
            }
        });

        for (texid, delta) in full_output.textures_delta.set {
            let size = Extent3d {
                depth_or_array_layers: 1,
                width: delta.image.width() as u32,
                height: delta.image.height() as u32,
            };
            let pixels = match &delta.image {
                ImageData::Color(color_image) => color_image.pixels.clone(),
                ImageData::Font(font_image) => font_image.srgba_pixels(None).collect(),
            };

            if let Some((_texbg, tex, texsize)) = self.textures.get_mut(&texid) {
                let pos = delta.pos.unwrap_or([0, 0]);
                queue.write_texture(
                    ImageCopyTexture {
                        texture: &tex,
                        mip_level: 0,
                        origin: Origin3d {
                            x: pos[0] as u32,
                            y: pos[1] as u32,
                            z: 0,
                        },
                        aspect: TextureAspect::All,
                    },
                    bytemuck::cast_slice::<_, u8>(&pixels),
                    ImageDataLayout {
                        offset: 0,
                        bytes_per_row: Some(texsize[0] * 4),
                        rows_per_image: None,
                    },
                    size,
                );
            } 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 texture = device.create_texture_with_data(
                    &queue,
                    &TextureDescriptor {
                        label: None,
                        size,
                        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, delta.image.size().map(|e| e as u32)),
                );
            }
        }

        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();
            }
        }

        if index_count == 0 || vertex_count == 0 {
            return;
        }

        // TODO realloc buffers and retry 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 affine = self
            .surfaces
            .values()
            .next()
            .map(|s| s.transform)
            .unwrap_or(Affine3A::IDENTITY);

        let scale = 0.03;
        let projection = projection
            * Mat4::from_mat3a(affine.matrix3)
            * Mat4::from_translation(affine.translation.into())
            * Mat4::from_mat3(Mat3::from_mat2(Mat2::from_cols_array(&[
                scale, 0., 0., -scale,
            ])));

        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);
        }
    }
}