/*
    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, ViewportId, ViewportInfo,
    epaint::{ImageDelta, Primitive, Vertex},
};
use glam::{Affine3A, Mat2, Mat3, Mat4};
use log::info;
use rand::random;
use std::{
    collections::HashMap,
    num::NonZeroU64,
    sync::{Arc, RwLock},
};
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 {
    device: Arc<Device>,
    queue: Arc<Queue>,
    ctx: Context,
    pipeline: RenderPipeline,
    bind_group_layout: BindGroupLayout,
    textures: RwLock<HashMap<TextureId, (BindGroup, Texture, [u32; 2])>>,
    surfaces: RwLock<HashMap<ViewportId, UiSurface>>,
}

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

impl UiRenderer {
    pub fn new(device: Arc<Device>, queue: Arc<Queue>, format: TextureFormat) -> Self {
        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,
            device,
            queue,
            bind_group_layout,
            textures: HashMap::new().into(),
            surfaces: HashMap::new().into(),
        }
    }

    pub fn add_surface(
        &mut self,
        transform: Affine3A,
        content: impl Fn(&Context) + Send + Sync + 'static,
    ) {
        let index_capacity = 1024;
        let vertex_capacity = 1024;
        let index = self.device.create_buffer(&BufferDescriptor {
            label: None,
            size: (size_of::<f32>() * index_capacity) as u64,
            usage: BufferUsages::INDEX | BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        let vertex = self.device.create_buffer(&BufferDescriptor {
            label: None,
            size: (size_of::<Vertex>() * vertex_capacity) as u64,
            usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        self.surfaces.write().unwrap().insert(
            ViewportId::from_hash_of(random::<u128>()),
            UiSurface {
                transform,
                content: Arc::new(content),
                index,
                vertex,
                index_capacity,
                vertex_capacity,
            },
        );
    }

    pub fn apply_texture_delta(&self, texid: TextureId, delta: ImageDelta) {
        let mut textures = self.textures.write().unwrap();
        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)) = textures.get_mut(&texid) {
            let pos = delta.pos.unwrap_or([0, 0]);
            info!("updating UI texture at {pos:?}");
            self.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 = self.device.create_texture_with_data(
                &self.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 = self.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 = self.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),
                    },
                ],
            });
            textures.insert(
                texid,
                (bindgroup, texture, delta.image.size().map(|e| e as u32)),
            );
        }
    }

    pub fn draw(
        &mut self,
        commands: &mut CommandEncoder,
        target: &TextureView,
        depth: &TextureView,
        projection: Mat4,
    ) {
        let mut surfaces = self.surfaces.write().unwrap();
        if surfaces.is_empty() {
            return;
        }

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

        rpass.set_pipeline(&self.pipeline);

        let mut raw_input = egui::RawInput::default();
        raw_input.viewport_id = surfaces.keys().next().copied().unwrap();
        raw_input.viewports = surfaces
            .keys()
            .map(|k| {
                (*k, ViewportInfo {
                    native_pixels_per_point: Some(2.),
                    ..Default::default()
                })
            })
            .collect();

        for (_viewport_id, surf) in surfaces.iter_mut() {
            let full_output = self.ctx.run(raw_input.clone(), |ctx| (surf.content)(ctx));

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

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

            while index_count > surf.index_capacity {
                info!(
                    "index buffer overflow. expanding {} -> {}",
                    surf.index_capacity,
                    surf.index_capacity * 2
                );
                surf.index_capacity *= 2;
                surf.index = self.device.create_buffer(&BufferDescriptor {
                    label: None,
                    size: (size_of::<f32>() * surf.index_capacity) as u64,
                    usage: BufferUsages::INDEX | BufferUsages::COPY_DST,
                    mapped_at_creation: false,
                });
            }
            while vertex_count > surf.vertex_capacity {
                info!(
                    "vertex buffer overflow. expanding {} -> {}",
                    surf.vertex_capacity,
                    surf.vertex_capacity * 2
                );
                surf.vertex_capacity *= 2;
                surf.vertex = self.device.create_buffer(&BufferDescriptor {
                    label: None,
                    size: (size_of::<f32>() * surf.vertex_capacity) as u64,
                    usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
                    mapped_at_creation: false,
                });
            }

            let mut mapped_index = self
                .queue
                .write_buffer_with(
                    &surf.index,
                    0,
                    NonZeroU64::new((size_of::<u32>() * index_count) as u64).unwrap(),
                )
                .expect("ui index buffer overflow");
            let mut mapped_vertex = self
                .queue
                .write_buffer_with(
                    &surf.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 scale = 0.03;
            let projection = projection
                * Mat4::from_mat3a(surf.transform.matrix3)
                * Mat4::from_translation(surf.transform.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_push_constants(ShaderStages::VERTEX, 0, projection.as_flattened());
            rpass.set_index_buffer(surf.index.slice(..), IndexFormat::Uint32);
            rpass.set_vertex_buffer(0, surf.vertex.slice(..));
            for (index, base_vertex, texid) in slices {
                let tex_guard = self.textures.read().unwrap();
                let bind_group = &tex_guard.get(&texid).unwrap().0;
                rpass.set_bind_group(0, bind_group, &[]);
                rpass.draw_indexed(index, base_vertex, 0..1);
            }
        }
    }
}