/*
    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 crate::{
    download::Downloader,
    meshops::{generate_normals, generate_tangents, generate_texcoords},
};
use anyhow::Result;
use egui::{Grid, Widget};
use glam::{Vec2, Vec3};
use humansize::DECIMAL;
use image::ImageReader;
use log::{debug, trace};
use std::{
    collections::{HashMap, HashSet},
    hash::Hash,
    io::Cursor,
    marker::PhantomData,
    sync::{Arc, RwLock},
    time::Instant,
};
use weareshared::{
    Affine3A,
    packets::Resource,
    resources::{Image, MeshPart, Prefab},
};
use wgpu::{
    AddressMode, BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayout, BindingResource,
    Buffer, BufferUsages, Device, Extent3d, FilterMode, Queue, SamplerDescriptor, Texture,
    TextureDescriptor, TextureDimension, TextureFormat, TextureUsages, TextureViewDescriptor,
    util::{BufferInitDescriptor, DeviceExt, TextureDataOrder},
};

pub struct DemandMap<K, V> {
    inner: RwLock<DemandMapState<K, V>>,
}
struct DemandMapState<K, V> {
    values: HashMap<K, V>,
    needed: HashSet<K>,
    size_metric: usize,
}
impl<K: Hash + Eq + Clone, V: Clone> DemandMap<K, V> {
    pub fn new() -> Self {
        Self {
            inner: DemandMapState {
                needed: HashSet::new(),
                values: HashMap::new(),
                size_metric: 0,
            }
            .into(),
        }
    }
    pub fn needed(&self) -> Vec<K> {
        self.inner.read().unwrap().needed.iter().cloned().collect()
    }
    pub fn insert(&self, key: K, value: V, size: usize) {
        let mut s = self.inner.write().unwrap();
        s.needed.remove(&key);
        s.values.insert(key, value);
        s.size_metric += size;
    }
    pub fn try_get(&self, key: K) -> Option<V> {
        let mut s = self.inner.write().unwrap();
        if let Some(k) = s.values.get(&key) {
            Some(k.to_owned())
        } else {
            s.needed.insert(key);
            None
        }
    }
}

pub struct ScenePreparer {
    device: Arc<Device>,
    queue: Arc<Queue>,
    texture_bgl: BindGroupLayout,

    textures: DemandMap<(Resource<Image<'static>>, bool), (Arc<Texture>, Arc<BindGroup>)>,
    placeholder_textures: DemandMap<TextureIdentityKind, (Arc<Texture>, Arc<BindGroup>)>,
    index_buffers: DemandMap<Resource<Vec<[u32; 3]>>, (Arc<Buffer>, u32)>,
    vertex_buffers: DemandMap<Resource<Vec<f32>>, Arc<Buffer>>,
    generated_tangent_buffers: DemandMap<TangentBufferSpec, Arc<Buffer>>,
    generated_normal_buffers: DemandMap<NormalBufferSpec, Arc<Buffer>>,
    generated_texcoord_buffers: DemandMap<TexcoordBufferSpec, Arc<Buffer>>,
    mesh_parts: DemandMap<Resource<MeshPart>, Arc<RMeshPart>>,
    pub prefabs: DemandMap<Resource<Prefab>, Arc<RPrefab>>,
}

pub struct RPrefab(pub Vec<(Affine3A, Arc<RMeshPart>)>);
pub struct RMeshPart {
    pub index_count: u32,
    pub index: Arc<Buffer>,
    pub va_position: Arc<Buffer>,
    pub va_normal: Arc<Buffer>,
    pub va_tangent: Arc<Buffer>,
    pub va_texcoord: Arc<Buffer>,
    pub tex_albedo: Arc<BindGroup>,
    pub tex_normal: Arc<BindGroup>,
    pub double_sided: bool,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct TangentBufferSpec {
    index: Resource<Vec<[u32; 3]>>,
    position: Resource<Vec<Vec3>>,
    texcoord: Resource<Vec<Vec2>>,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct NormalBufferSpec {
    index: Resource<Vec<[u32; 3]>>,
    position: Resource<Vec<Vec3>>,
}
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct TexcoordBufferSpec {
    index: Resource<Vec<[u32; 3]>>,
    position: Resource<Vec<Vec3>>,
}

#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
enum TextureIdentityKind {
    Normal,
    Multiply,
}

impl ScenePreparer {
    pub fn new(device: Arc<Device>, queue: Arc<Queue>, texture_bgl: BindGroupLayout) -> Self {
        Self {
            texture_bgl,
            index_buffers: DemandMap::new(),
            vertex_buffers: DemandMap::new(),
            mesh_parts: DemandMap::new(),
            prefabs: DemandMap::new(),
            textures: DemandMap::new(),
            placeholder_textures: DemandMap::new(),
            generated_tangent_buffers: DemandMap::new(),
            generated_normal_buffers: DemandMap::new(),
            generated_texcoord_buffers: DemandMap::new(),
            device,
            queue,
        }
    }
    pub fn update(&self, dls: &Downloader) -> Result<usize> {
        let mut num_done = 0;
        for pres in self.prefabs.needed() {
            if let Some(prefab) = dls.try_get(pres.clone())? {
                let mut rprefab = RPrefab(Vec::new());
                for (aff, partres) in &prefab.mesh {
                    if let Some(part) = self.mesh_parts.try_get(partres.clone()) {
                        rprefab.0.push((*aff, part.clone()));
                    }
                }
                if rprefab.0.len() == prefab.mesh.len() {
                    self.prefabs.insert(pres.clone(), Arc::new(rprefab), 0);
                    debug!("prefab created ({pres})");
                    num_done += 1;
                }
            }
        }
        for pres in self.index_buffers.needed() {
            let start = Instant::now();
            if let Some(buf) = dls.try_get(pres.clone())? {
                let buffer = self.device.create_buffer_init(&BufferInitDescriptor {
                    label: Some("index"),
                    contents: bytemuck::cast_slice(buf.as_slice()),
                    usage: BufferUsages::INDEX | BufferUsages::COPY_DST,
                });
                self.index_buffers.insert(
                    pres.clone(),
                    (Arc::new(buffer), (buf.len() * 3) as u32),
                    buf.len() * size_of::<u32>() * 3,
                );
                debug!(
                    "index buffer created (len={}, took {:?}) {pres}",
                    buf.len() / size_of::<u32>(),
                    start.elapsed(),
                );
                num_done += 1;
            }
        }
        for pres in self.vertex_buffers.needed() {
            let start = Instant::now();
            if let Some(buf) = dls.try_get(pres.clone())? {
                let buffer = self.device.create_buffer_init(&BufferInitDescriptor {
                    contents: bytemuck::cast_slice(buf.as_slice()),
                    label: Some("vertex attribute"),
                    usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
                });
                self.vertex_buffers.insert(
                    pres.clone(),
                    Arc::new(buffer),
                    buf.len() * size_of::<f32>(),
                );
                debug!(
                    "vertex attribute buffer created (len={}, took {:?}) {pres}",
                    buf.len() / size_of::<f32>(),
                    start.elapsed()
                );
                num_done += 1;
            }
        }
        for (pres, linear) in self.textures.needed() {
            let start = Instant::now();
            if let Some(buf) = dls.try_get(pres.clone())? {
                let image = ImageReader::new(Cursor::new(buf.0)).with_guessed_format()?;
                let image = image.decode()?;
                let dims = (image.width(), image.height());
                let image = image.into_rgba8();
                let image = image.into_vec();
                let tex_bg = create_texture(
                    &self.device,
                    &self.queue,
                    &self.texture_bgl,
                    &image,
                    dims.0,
                    dims.1,
                    linear,
                );
                self.textures
                    .insert((pres.clone(), linear), tex_bg, image.len());
                debug!(
                    "texture created (res={}x{}, took {:?})",
                    dims.0,
                    dims.1,
                    start.elapsed()
                );
                num_done += 1;
            }
        }
        for kind in self.placeholder_textures.needed() {
            let (linear, color) = match kind {
                TextureIdentityKind::Normal => (true, [128, 128, 255, 255]),
                TextureIdentityKind::Multiply => (false, [255, 255, 255, 255]),
            };
            let tex_bg = create_texture(
                &self.device,
                &self.queue,
                &self.texture_bgl,
                &color,
                1,
                1,
                linear,
            );
            self.placeholder_textures.insert(kind, tex_bg, 4);
            num_done += 1;
        }
        for spec in self.generated_tangent_buffers.needed() {
            if let (Some(index), Some(position), Some(texcoord)) = (
                dls.try_get(spec.index.clone())?,
                dls.try_get(spec.position.clone())?,
                dls.try_get(spec.texcoord.clone())?,
            ) {
                let tangents = generate_tangents(&index, &position, &texcoord);
                let buffer = self.device.create_buffer_init(&BufferInitDescriptor {
                    label: Some("generated tangent"),
                    usage: BufferUsages::COPY_DST | BufferUsages::VERTEX,
                    contents: bytemuck::cast_slice(tangents.as_slice()),
                });
                self.generated_tangent_buffers.insert(
                    spec,
                    Arc::new(buffer),
                    size_of::<f32>() * tangents.len() * 3,
                );
            }
        }
        for spec in self.generated_normal_buffers.needed() {
            if let (Some(index), Some(position)) = (
                dls.try_get(spec.index.clone())?,
                dls.try_get(spec.position.clone())?,
            ) {
                let normals = generate_normals(&index, &position);
                let buffer = self.device.create_buffer_init(&BufferInitDescriptor {
                    label: Some("generated normal"),
                    usage: BufferUsages::COPY_DST | BufferUsages::VERTEX,
                    contents: bytemuck::cast_slice(normals.as_slice()),
                });
                self.generated_normal_buffers.insert(
                    spec,
                    Arc::new(buffer),
                    size_of::<f32>() * normals.len() * 3,
                );
            }
        }
        for spec in self.generated_texcoord_buffers.needed() {
            if let (Some(index), Some(position)) = (
                dls.try_get(spec.index.clone())?,
                dls.try_get(spec.position.clone())?,
            ) {
                let texcoords = generate_texcoords(&index, &position);
                let buffer = self.device.create_buffer_init(&BufferInitDescriptor {
                    label: Some("generated texcoord"),
                    usage: BufferUsages::COPY_DST | BufferUsages::VERTEX,
                    contents: bytemuck::cast_slice(texcoords.as_slice()),
                });
                self.generated_texcoord_buffers.insert(
                    spec,
                    Arc::new(buffer),
                    size_of::<f32>() * texcoords.len() * 3,
                );
            }
        }
        for pres in self.mesh_parts.needed() {
            let start = Instant::now();
            if let Some(part) = dls.try_get(pres.clone())? {
                if let (Some(indexres), Some(positionres)) = (part.index, part.va_position) {
                    let index = self.index_buffers.try_get(indexres.clone());
                    let position = self
                        .vertex_buffers
                        .try_get(Resource(positionres.0, PhantomData));

                    let normal = if let Some(res) = part.va_normal.clone() {
                        self.vertex_buffers.try_get(Resource(res.0, PhantomData))
                    } else {
                        self.generated_normal_buffers.try_get(NormalBufferSpec {
                            index: indexres.clone(),
                            position: Resource(positionres.0, PhantomData),
                        })
                    };

                    let texcoord = if let Some(res) = part.va_texcoord.clone() {
                        self.vertex_buffers.try_get(Resource(res.0, PhantomData))
                    } else {
                        self.generated_texcoord_buffers.try_get(TexcoordBufferSpec {
                            index: indexres.clone(),
                            position: Resource(positionres.0, PhantomData),
                        })
                    };
                    let tangent = if let Some(res) = part.va_tangent.clone() {
                        self.vertex_buffers.try_get(Resource(res.0, PhantomData))
                    } else {
                        self.generated_tangent_buffers.try_get(TangentBufferSpec {
                            index: indexres,
                            position: Resource(positionres.0, PhantomData),
                            texcoord: part.va_texcoord.expect("TODO"),
                        })
                    };

                    let mut tex_albedo = None;
                    if let Some(albedores) = part.tex_albedo {
                        if let Some((_tex, bg)) = self.textures.try_get((albedores, false)) {
                            tex_albedo = Some(bg)
                        }
                    } else {
                        if let Some((_tex, bg)) = self
                            .placeholder_textures
                            .try_get(TextureIdentityKind::Multiply)
                        {
                            tex_albedo = Some(bg)
                        }
                    }
                    let mut tex_normal = None;
                    if let Some(normalres) = part.tex_normal {
                        if let Some((_tex, bg)) = self.textures.try_get((normalres, true)) {
                            tex_normal = Some(bg)
                        }
                    } else {
                        if let Some((_tex, bg)) = self
                            .placeholder_textures
                            .try_get(TextureIdentityKind::Normal)
                        {
                            tex_normal = Some(bg)
                        }
                    }

                    if let (
                        Some((index, index_count)),
                        Some(va_normal),
                        Some(va_tangent),
                        Some(va_texcoord),
                        Some(va_position),
                        Some(tex_normal),
                        Some(tex_albedo),
                    ) = (
                        index, normal, tangent, texcoord, position, tex_normal, tex_albedo,
                    ) {
                        debug!("part created (took {:?}) {pres}", start.elapsed());
                        self.mesh_parts.insert(
                            pres,
                            Arc::new(RMeshPart {
                                index_count,
                                index,
                                va_normal,
                                va_tangent,
                                va_position,
                                va_texcoord,
                                tex_albedo,
                                tex_normal,
                                double_sided: part.g_double_sided.is_some(),
                            }),
                            0,
                        );
                        num_done += 1;
                    }
                }
            }
        }
        self.print_missing();
        Ok(num_done)
    }
}

fn create_texture(
    device: &Device,
    queue: &Queue,
    bgl: &BindGroupLayout,
    data: &[u8],
    width: u32,
    height: u32,
    linear: bool,
) -> (Arc<Texture>, Arc<BindGroup>) {
    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: if linear {
                TextureFormat::Rgba8Unorm
            } else {
                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 {
        address_mode_u: AddressMode::Repeat,
        address_mode_v: AddressMode::Repeat,
        mag_filter: FilterMode::Linear,
        min_filter: FilterMode::Linear,
        ..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),
            },
        ],
    });
    (Arc::new(texture), Arc::new(bindgroup))
}

impl<K, V> Widget for &DemandMap<K, V> {
    fn ui(self, ui: &mut egui::Ui) -> egui::Response {
        let state = self.inner.read().unwrap();
        ui.label(state.needed.len().to_string());
        ui.label(state.values.len().to_string());
        ui.label(humansize::format_size(state.size_metric, DECIMAL));
        ui.end_row();
        ui.response()
    }
}
impl ScenePreparer {
    pub fn print_missing(&self) {
        fn visit<K, V>(name: &str, m: &DemandMap<K, V>)
        where
            K: Clone,
            K: Hash,
            K: std::cmp::Eq,
            V: Clone,
        {
            let nl = m.needed().len();
            if nl > 0 {
                trace!("{name}: need {nl}")
            }
        }
        visit("prefabs", &self.prefabs);
        visit("mesh_parts", &self.mesh_parts);
        visit("vertex_buffers", &self.vertex_buffers);
        visit("index_buffers", &self.index_buffers);
        visit("placeholder_textures", &self.placeholder_textures);
        visit("generated_tangent_buffers", &self.generated_tangent_buffers);
        visit("generated_normal_buffers", &self.generated_normal_buffers);
        visit(
            "generated_texcoord_buffers",
            &self.generated_texcoord_buffers,
        );
        visit("textures", &self.textures);
    }
}

impl Widget for &ScenePreparer {
    fn ui(self, ui: &mut egui::Ui) -> egui::Response {
        Grid::new("sp")
            .num_columns(4)
            .show(ui, |ui| {
                ui.label("prefabs");
                self.prefabs.ui(ui);
                ui.label("mesh_parts");
                self.mesh_parts.ui(ui);
                ui.label("vertex_buffers");
                self.vertex_buffers.ui(ui);
                ui.label("index_buffers");
                self.index_buffers.ui(ui);
                ui.label("placeholder_textures");
                self.placeholder_textures.ui(ui);
                ui.label("generated_tangent_buffers");
                self.generated_tangent_buffers.ui(ui);
                ui.label("generated_normal_buffers");
                self.generated_normal_buffers.ui(ui);
                ui.label("generated_texcoord_buffers");
                self.generated_texcoord_buffers.ui(ui);
                ui.label("textures");
                self.textures.ui(ui);
            })
            .response
    }
}