/*
    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;
use anyhow::Result;
use image::ImageReader;
use log::debug;
use std::{
    collections::{HashMap, HashSet},
    hash::Hash,
    io::Cursor,
    sync::{Arc, RwLock},
};
use weareshared::{
    Affine3A,
    packets::Resource,
    resources::{AttributeArray, Image, IndexArray, 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>,
}
impl<K: Hash + Eq + Clone, V: Clone> DemandMap<K, V> {
    pub fn new() -> Self {
        Self {
            inner: DemandMapState {
                needed: HashSet::new(),
                values: HashMap::new(),
            }
            .into(),
        }
    }
    pub fn needed(&self) -> Vec<K> {
        self.inner.read().unwrap().needed.iter().cloned().collect()
    }
    pub fn insert(&self, key: K, value: V) {
        let mut s = self.inner.write().unwrap();
        s.needed.remove(&key);
        s.values.insert(key, value);
    }
    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>, (Arc<Texture>, Arc<BindGroup>)>,
    placeholder_textures: DemandMap<(), (Arc<Texture>, Arc<BindGroup>)>,
    index_buffers: DemandMap<Resource<IndexArray>, (Arc<Buffer>, u32)>,
    vertex_buffers: DemandMap<Resource<AttributeArray>, (Arc<Buffer>, u32)>,
    placeholder_vertex_buffers: DemandMap<(u32, bool), 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 position: [Arc<Buffer>; 3],
    pub normal: [Arc<Buffer>; 3],
    pub texcoord: [Arc<Buffer>; 2],
    pub texture: Arc<BindGroup>,
}

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_vertex_buffers: DemandMap::new(),
            placeholder_textures: DemandMap::new(),
            device,
            queue,
        }
    }
    pub fn update(&self, dls: &Downloader) -> Result<()> {
        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));
                    debug!("prefab created ({pres})");
                }
            }
        }
        for pres in self.index_buffers.needed() {
            if let Some(buf) = dls.try_get(pres.clone())? {
                let buf = buf
                    .0
                    .into_iter()
                    .flatten()
                    .flat_map(u16::to_le_bytes)
                    .collect::<Vec<_>>();
                let buffer = self.device.create_buffer_init(&BufferInitDescriptor {
                    contents: &buf,
                    label: None,
                    usage: BufferUsages::INDEX | BufferUsages::COPY_DST,
                });
                self.index_buffers
                    .insert(pres.clone(), (Arc::new(buffer), (buf.len() / 2) as u32));
                debug!("index buffer created (len={}) {pres}", buf.len() / 2);
            }
        }
        for pres in self.vertex_buffers.needed() {
            if let Some(buf) = dls.try_get(pres.clone())? {
                let buf = buf
                    .0
                    .into_iter()
                    .map(f32::to_le_bytes)
                    .flatten()
                    .collect::<Vec<_>>();
                let buffer = self.device.create_buffer_init(&BufferInitDescriptor {
                    contents: &buf,
                    label: None,
                    usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
                });
                self.vertex_buffers
                    .insert(pres.clone(), (Arc::new(buffer), (buf.len() / 4) as u32));
                debug!(
                    "vertex attribute buffer created (len={}) {pres}",
                    buf.len() / 4
                );
            }
        }
        for pres in self.textures.needed() {
            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 image = image.to_rgba8();
                let image_raw = image.to_vec();
                let tex_bg = create_texture(
                    &self.device,
                    &self.queue,
                    &self.texture_bgl,
                    &image_raw,
                    image.width(),
                    image.height(),
                );
                self.textures.insert(pres.clone(), tex_bg);
            }
        }
        for pres in self.placeholder_textures.needed() {
            let tex_bg = create_texture(
                &self.device,
                &self.queue,
                &self.texture_bgl,
                &[255, 255, 255, 255],
                1,
                1,
            );
            self.placeholder_textures.insert(pres, tex_bg);
        }
        for pres in self.mesh_parts.needed() {
            if let Some(part) = dls.try_get(pres.clone())? {
                if let (Some(indexres), Some(positionres)) = (part.index, part.va_position) {
                    let Some((index, index_count)) = self.index_buffers.try_get(indexres.clone())
                    else {
                        continue;
                    };
                    let mut position = Vec::new();
                    let mut vertex_count = 0;
                    for vr in positionres {
                        if let Some((vertex, n)) = self.vertex_buffers.try_get(vr) {
                            vertex_count = n;
                            position.push(vertex);
                        }
                    }
                    let mut normal = Vec::new();
                    if let Some(normalres) = part.va_normal {
                        for vr in normalres {
                            if let Some((vertex, _)) = self.vertex_buffers.try_get(vr) {
                                normal.push(vertex);
                            }
                        }
                    } else {
                        // TODO generate normals
                        for _ in 0..3 {
                            if let Some(buf) = self
                                .placeholder_vertex_buffers
                                .try_get((vertex_count, false))
                            {
                                normal.push(buf);
                            }
                        }
                    }
                    let mut texcoord = Vec::new();
                    if let Some(texcoordres) = part.va_texcoord {
                        for vr in texcoordres {
                            if let Some((vertex, _)) = self.vertex_buffers.try_get(vr) {
                                texcoord.push(vertex);
                            }
                        }
                    } else {
                        // TODO generate UVs
                        for _ in 0..3 {
                            if let Some(buf) = self
                                .placeholder_vertex_buffers
                                .try_get((vertex_count, false))
                            {
                                texcoord.push(buf);
                            }
                        }
                    }
                    let mut texture = None;
                    if let Some(albedores) = part.tex_albedo {
                        if let Some((_tex, bg)) = self.textures.try_get(albedores) {
                            texture = Some(bg)
                        }
                    } else {
                        if let Some((_tex, bg)) = self.placeholder_textures.try_get(()) {
                            texture = Some(bg)
                        }
                    }

                    if texcoord.len() == 2
                        && normal.len() == 3
                        && position.len() == 3
                        && texture.is_some()
                    {
                        debug!("part created ({pres})");
                        self.mesh_parts.insert(
                            pres,
                            Arc::new(RMeshPart {
                                index_count,
                                index,
                                texcoord: texcoord.try_into().unwrap(),
                                normal: normal.try_into().unwrap(),
                                position: position.try_into().unwrap(),
                                texture: texture.unwrap(),
                            }),
                        );
                    }
                }
            }
        }
        Ok(())
    }
}

fn create_texture(
    device: &Device,
    queue: &Queue,
    bgl: &BindGroupLayout,
    data: &[u8],
    width: u32,
    height: u32,
) -> (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: 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))
}