/*
    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/>.
*/
pub mod demand_map;
pub mod draw;
pub mod meshops;
pub mod pipelines;
pub mod textures;
pub mod vertex_buffers;

use super::{GraphicsConfig, shaders::SceneShaders};
use crate::{armature::RArmature, download::Downloader};
use anyhow::Result;
use bytemuck::{Pod, Zeroable};
use demand_map::DemandMap;
use egui::{Grid, Widget};
use glam::{UVec3, UVec4, Vec2, Vec3, Vec3A, uvec3, uvec4};
use log::{debug, info, trace};
use pipelines::SceneBgLayouts;
use std::{
    hash::Hash,
    marker::PhantomData,
    sync::{Arc, RwLock},
    time::Instant,
};
use textures::MipGenerationPipeline;
use weareshared::{
    Affine3A,
    packets::Resource,
    resources::{Image, MeshPart, Prefab},
};
use wgpu::{
    BindGroup, BindGroupDescriptor, BindGroupEntry, Buffer, BufferUsages, Device, Queue,
    RenderPipeline, Texture, TextureFormat,
    util::{BufferInitDescriptor, DeviceExt},
};

pub struct ScenePreparer {
    device: Arc<Device>,
    queue: Arc<Queue>,
    layouts: SceneBgLayouts,
    shaders: SceneShaders,
    render_format: TextureFormat,
    config: RwLock<GraphicsConfig>,
    downloader: Arc<Downloader>,

    textures: DemandMap<TextureSpec, (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>, Result<Arc<RMeshPart>, String>>,
    materials: DemandMap<Material, Arc<BindGroup>>,
    pipelines: DemandMap<PipelineSpec, Arc<RenderPipeline>>,
    mip_generation_pipelines: DemandMap<TextureFormat, Arc<MipGenerationPipeline>>,
    pub prefabs: DemandMap<Resource<Prefab>, Arc<RPrefab>>,
}

pub struct RPrefab(pub Vec<(Affine3A, Arc<RMeshPart>)>);
pub struct RMeshPart {
    pub pipeline: Arc<RenderPipeline>,
    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 material: Arc<BindGroup>,
    pub double_sided: bool,

    pub va_joint_index: Option<Arc<Buffer>>,
    pub va_joint_weight: Option<Arc<Buffer>>,
    pub joint_uniform: Option<Arc<Buffer>>,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub struct PipelineSpec {
    format: TextureFormat,
    skin: bool,
    backface_culling: bool,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct TextureSpec {
    data: Resource<Image<'static>>,
    linear: bool,
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct TangentBufferSpec {
    index: Resource<Vec<[u32; 3]>>,
    position: Resource<Vec<Vec3>>,
    texcoord: Option<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,
}

#[derive(Debug, Clone, Copy, Pod, Zeroable, Hash, PartialEq, Eq)]
#[repr(C)]
struct Material {
    roughness: u32,
    metallic: u32,
    _pad1: [u32; 2],
    albedo_alpha: UVec4,
    emission: UVec3,
    _pad2: u32,
}

impl ScenePreparer {
    pub fn new(
        device: Arc<Device>,
        queue: Arc<Queue>,
        render_format: TextureFormat,
        downloader: Arc<Downloader>,
        config: GraphicsConfig,
    ) -> Self {
        Self {
            render_format,
            config: config.into(),
            layouts: SceneBgLayouts::load(&device),
            shaders: SceneShaders::load(&device),
            device,
            queue,
            downloader,
            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(),
            materials: DemandMap::new(),
            pipelines: DemandMap::new(),
            mip_generation_pipelines: DemandMap::new(),
        }
    }
    pub fn reconfigure(&self, config: &GraphicsConfig) {
        let mut cc = self.config.write().unwrap();
        if cc.max_anisotropy != config.max_anisotropy || cc.max_mip_count != config.max_mip_count {
            info!("clear all scene textures");
            self.textures.clear();
            self.mesh_parts.clear();
            self.prefabs.clear();
        }
        if cc.sample_count != config.sample_count {
            info!("clear all scene pipelines");
            self.pipelines.clear();
            self.mesh_parts.clear();
            self.prefabs.clear();
        }
        *cc = config.clone();
    }
    pub fn update(&self) -> Result<usize> {
        let mut num_done = 0;

        self.update_textures(&mut num_done)?;
        self.update_vertex_buffers(&mut num_done)?;

        for pres in self.prefabs.needed() {
            if let Some(prefab) = self.downloader.try_get(pres.clone())? {
                let mut rprefab = RPrefab(Vec::new());
                let mut done_or_failed = 0;
                for (aff, partres) in &prefab.mesh {
                    if let Some(part) = self.mesh_parts.try_get(partres.clone()) {
                        if let Ok(part) = part.clone() {
                            rprefab.0.push((*aff, part));
                        }
                        done_or_failed += 1;
                    }
                }
                if done_or_failed == prefab.mesh.len() {
                    self.prefabs.insert(pres.clone(), Arc::new(rprefab), 0);
                    debug!("prefab created ({pres})");
                    num_done += 1;
                }
            }
        }

        for spec in self.materials.needed() {
            let buffer = self.device.create_buffer_init(&BufferInitDescriptor {
                label: Some("material props"),
                usage: BufferUsages::COPY_DST | BufferUsages::UNIFORM,
                contents: bytemuck::cast_slice(&[spec]),
            });
            let bind_group = self.device.create_bind_group(&BindGroupDescriptor {
                label: Some("material"),
                layout: &self.layouts.material,
                entries: &[BindGroupEntry {
                    binding: 0,
                    resource: buffer.as_entire_binding(),
                }],
            });
            self.materials.insert(spec, Arc::new(bind_group), 0);
        }
        for spec in self.pipelines.needed() {
            self.pipelines.insert(
                spec.clone(),
                Arc::new(spec.create(
                    &self.device,
                    &self.layouts,
                    &self.shaders,
                    &self.config.read().unwrap().clone(),
                )),
                0,
            );
        }
        for pres in self.mesh_parts.needed() {
            let start = Instant::now();
            if let Some(part) = self.downloader.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,
                        })
                    };

                    let joint_weight = if let Some(res) = part.va_joint_weight.clone() {
                        self.vertex_buffers
                            .try_get(Resource(res.0, PhantomData))
                            .map(Some)
                    } else {
                        Some(None)
                    };
                    let joint_index = if let Some(res) = part.va_joint_index.clone() {
                        self.vertex_buffers
                            .try_get(Resource(res.0, PhantomData))
                            .map(Some)
                    } else {
                        Some(None)
                    };

                    let mut tex_albedo = None;
                    if let Some(albedores) = part.tex_albedo {
                        if let Some((_tex, bg)) = self.textures.try_get(TextureSpec {
                            data: albedores,
                            linear: 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(TextureSpec {
                            data: normalres,
                            linear: true,
                        }) {
                            tex_normal = Some(bg)
                        }
                    } else if let Some((_tex, bg)) = self
                        .placeholder_textures
                        .try_get(TextureIdentityKind::Normal)
                    {
                        tex_normal = Some(bg)
                    }

                    let material = self.materials.try_get({
                        let albedo = part.g_albedo.unwrap_or(Vec3A::ONE);
                        let emission = part.g_emission.unwrap_or(Vec3A::ONE);
                        Material {
                            roughness: part.g_roughness.unwrap_or(1.).to_bits(),
                            metallic: part.g_metallic.unwrap_or(0.).to_bits(),
                            _pad1: [0, 0],
                            albedo_alpha: uvec4(
                                albedo.x.to_bits(),
                                albedo.y.to_bits(),
                                albedo.z.to_bits(),
                                part.g_alpha.unwrap_or(1.).to_bits(),
                            ),
                            emission: uvec3(
                                emission.x.to_bits(),
                                emission.y.to_bits(),
                                emission.z.to_bits(),
                            ),
                            _pad2: 0,
                        }
                    });

                    let armature = Some(None); // TODO

                    let pipeline = self.pipelines.try_get(PipelineSpec {
                        format: self.render_format,
                        skin: false,
                        backface_culling: part.g_double_sided.is_none(),
                    });

                    if let (
                        Some(pipeline),
                        Some((index, index_count)),
                        Some(va_normal),
                        Some(va_tangent),
                        Some(va_texcoord),
                        Some(va_position),
                        Some(va_joint_index),
                        Some(va_joint_weight),
                        Some(armature),
                        Some(tex_normal),
                        Some(tex_albedo),
                        Some(material),
                    ) = (
                        pipeline,
                        index,
                        normal,
                        tangent,
                        texcoord,
                        position,
                        joint_index,
                        joint_weight,
                        armature,
                        tex_normal,
                        tex_albedo,
                        material,
                    ) {
                        let double_sided = part.g_double_sided.is_some();

                        let joint_uniform = if let Some(a) = armature {
                            let ra = RArmature::new(&self.device, a);
                            Some(ra.joint_mat_uniform_buffer.clone())
                        } else {
                            None
                        };

                        debug!("part created (took {:?}) {pres}", start.elapsed());
                        self.mesh_parts.insert(
                            pres,
                            Ok(Arc::new(RMeshPart {
                                pipeline,
                                index_count,
                                index,
                                va_normal,
                                va_tangent,
                                va_position,
                                va_texcoord,
                                va_joint_index,
                                va_joint_weight,
                                tex_albedo,
                                tex_normal,
                                material,
                                double_sided,
                                joint_uniform,
                            })),
                            0,
                        );
                        num_done += 1;
                    }
                } else {
                    self.mesh_parts
                        .insert(pres, Err("index or positions missing".to_string()), 0);
                }
            }
        }
        self.print_missing();
        Ok(num_done)
    }
}

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);
        visit("materials", &self.materials);
        visit("pipelines", &self.pipelines);
        visit("mip_generation_pipelines", &self.mip_generation_pipelines);
    }
}

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);
                ui.label("materials");
                self.materials.ui(ui);
                ui.label("pipelines");
                self.pipelines.ui(ui);
                ui.label("mip_generation_pipelines");
                self.mip_generation_pipelines.ui(ui);
            })
            .response
    }
}