/*
    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::{helper::ReadWrite, packets::Resource};
use anyhow::Result;
use glam::{Affine3A, Vec2, Vec3A, Vec4};
use log::warn;
use std::{
    borrow::Cow,
    collections::BTreeMap,
    io::{Read, Write},
};

#[derive(Debug, Default, Clone)]
pub struct Prefab {
    pub name: Option<String>,
    pub mesh: Vec<(Affine3A, Resource<MeshPart>)>,
    pub collision: Vec<(Affine3A, Resource<CollisionPart>)>,
    pub light: Vec<(Vec3A, Resource<LightPart>)>,
    pub armature: Vec<Resource<ArmaturePart>>,
    pub environment: Option<Resource<EnvironmentPart>>,
}

#[derive(Debug, Default, Clone)]
pub struct LightPart {
    pub name: Option<String>,
    pub emission: Option<Vec3A>,
    pub radius: Option<f32>,
}

#[derive(Debug, Default, Clone)]
pub struct EnvironmentPart {
    pub skybox: Option<Resource<Image<'static>>>,
    pub sun: Option<(Vec3A, Vec3A)>,
}

#[derive(Debug, Default, Clone)]
pub struct MeshPart {
    pub name: Option<String>,
    pub index: Option<Resource<Vec<[u32; 3]>>>,
    pub armature: Option<u32>,
    pub g_metallic: Option<f32>,
    pub g_roughness: Option<f32>,
    pub g_albedo: Option<Vec3A>,
    pub g_transmission: Option<f32>,
    pub g_alpha: Option<f32>,
    pub g_emission: Option<Vec3A>,
    pub g_thickness: Option<f32>,
    pub g_refractive_index: Option<f32>,
    pub g_attenuation: Option<Vec3A>,
    pub g_dispersion: Option<f32>,
    pub g_unlit: Option<()>,
    pub g_double_sided: Option<()>,
    pub va_position: Option<Resource<Vec<Vec3A>>>,
    pub va_normal: Option<Resource<Vec<Vec3A>>>,
    pub va_tangent: Option<Resource<Vec<Vec4>>>,
    pub va_texcoord: Option<Resource<Vec<Vec2>>>,
    pub va_roughness: Option<Resource<Vec<f32>>>,
    pub va_metallic: Option<Resource<Vec<f32>>>,
    pub va_albedo: Option<Resource<Vec<Vec3A>>>,
    pub va_transmission: Option<Resource<Vec<f32>>>,
    pub va_alpha: Option<Resource<Vec<f32>>>,
    pub va_emission: Option<Resource<Vec<Vec3A>>>,
    pub va_joint_index: Option<Resource<Vec<[u32; 4]>>>,
    pub va_joint_weight: Option<Resource<Vec<[f32; 4]>>>,
    pub tex_normal: Option<Resource<Image<'static>>>,
    pub tex_roughness: Option<Resource<Image<'static>>>,
    pub tex_metallic: Option<Resource<Image<'static>>>,
    pub tex_albedo: Option<Resource<Image<'static>>>,
    pub tex_transmission: Option<Resource<Image<'static>>>,
    pub tex_alpha: Option<Resource<Image<'static>>>,
    pub tex_emission: Option<Resource<Image<'static>>>,
    pub tex_thickness: Option<Resource<Image<'static>>>,
    pub tex_occlusion: Option<Resource<Image<'static>>>,
    pub hint_mirror: Option<()>,
    pub hint_static: Option<()>,
}

#[derive(Debug, Default, Clone)]
pub struct ArmaturePart {
    pub parent: Option<Vec<u16>>,
    pub transform: Option<Vec<Affine3A>>,
    pub name: Option<Vec<String>>,
}

#[derive(Debug, Default, Clone)]
pub struct CollisionPart {
    pub name: Option<String>,
    pub restitution_coeff: Option<f32>,
    pub friction_kinetic: Option<f32>,
    pub friction_static: Option<f32>,
    pub sh_box: Option<Vec3A>,
    pub sh_sphere: Option<f32>,
    pub sh_cylinder: Option<(f32, f32, f32)>,
    pub sh_capsule: Option<(f32, f32, f32)>,
    pub sh_convex_hull: Option<Resource<Vec<Vec3A>>>,
    pub sh_mesh: Option<(Resource<Vec<[u32; 3]>>, Resource<Vec<Vec3A>>)>,
}

#[derive(Debug, Default, Clone)]
pub struct PrefabIndex(pub BTreeMap<String, Resource<Prefab>>);

#[derive(Debug, Clone)]
pub struct Image<'a>(pub Cow<'a, [u8]>);

impl ReadWrite for PrefabIndex {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        for (k, v) in &self.0 {
            write_kv(w, k.as_bytes(), &v.0)?;
        }
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut s = Self(BTreeMap::new());
        let mut g = Vec::new();
        r.read_to_end(&mut g)?;
        let mut g = g.as_slice();
        while !g.is_empty() {
            let (k, v) = read_kv(&mut g)?;
            s.0.insert(String::from_utf8(k)?, Resource::read(&mut v.as_slice())?);
        }
        Ok(s)
    }
}

impl ReadWrite for ArmaturePart {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        write_kv_opt(w, b"parent", &self.parent)?;
        write_kv_opt(w, b"transform", &self.transform)?;
        write_kv_opt(w, b"name", &self.name)?;
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut s = Self::default();
        read_kv_iter(r, |k, v| match k {
            b"parent" => Ok(s.parent = Some(read_slice(v)?)),
            b"name" => Ok(s.name = Some(read_slice(v)?)),
            b"transform" => Ok(s.transform = Some(read_slice(v)?)),
            x => Ok(warn!(
                "unknown armature key: {:?}",
                String::from_utf8_lossy(x)
            )),
        })?;
        Ok(s)
    }
}
impl ReadWrite for CollisionPart {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        write_kv_opt(w, b"name", &self.name)?;
        write_kv_opt(w, b"restitution_coeff", &self.restitution_coeff)?;
        write_kv_opt(w, b"friction_kinetic", &self.friction_kinetic)?;
        write_kv_opt(w, b"friction_static", &self.friction_static)?;
        write_kv_opt(w, b"sh_box", &self.sh_box)?;
        write_kv_opt(w, b"sh_sphere", &self.sh_sphere)?;
        write_kv_opt(w, b"sh_cylinder", &self.sh_cylinder)?;
        write_kv_opt(w, b"sh_capsule", &self.sh_capsule)?;
        write_kv_opt(w, b"sh_convex_hull", &self.sh_convex_hull)?;
        write_kv_opt(w, b"sh_mesh", &self.sh_mesh)?;
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut s = Self::default();
        read_kv_iter(r, |k, v| match k {
            b"name" => Ok(s.name = Some(read_slice(v)?)),
            b"restitution_coeff" => Ok(s.restitution_coeff = Some(read_slice(v)?)),
            b"friction_kinetic" => Ok(s.friction_kinetic = Some(read_slice(v)?)),
            b"friction_static" => Ok(s.friction_static = Some(read_slice(v)?)),
            b"sh_box" => Ok(s.sh_box = Some(read_slice(v)?)),
            b"sh_sphere" => Ok(s.sh_sphere = Some(read_slice(v)?)),
            b"sh_cylinder" => Ok(s.sh_cylinder = Some(read_slice(v)?)),
            b"sh_capsule" => Ok(s.sh_capsule = Some(read_slice(v)?)),
            b"sh_convex_hull" => Ok(s.sh_convex_hull = Some(read_slice(v)?)),
            b"sh_mesh" => Ok(s.sh_mesh = Some(read_slice(v)?)),
            x => Ok(warn!(
                "unknown prefab key: {:?}",
                String::from_utf8_lossy(x)
            )),
        })?;
        Ok(s)
    }
}
impl ReadWrite for Prefab {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        write_kv_opt(w, b"name", &self.name)?;
        for x in &self.mesh {
            write_kv_opt(w, b"mesh", &Some(x.clone()))?;
        }
        for x in &self.collision {
            write_kv_opt(w, b"collision", &Some(x.clone()))?;
        }
        for x in &self.light {
            write_kv_opt(w, b"light", &Some(x.clone()))?;
        }
        for x in &self.armature {
            write_kv_opt(w, b"armature", &Some(x.clone()))?;
        }
        write_kv_opt(w, b"environment", &self.environment)?;
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut s = Self::default();
        read_kv_iter(r, |k, v| match k {
            b"name" => Ok(s.name = Some(read_slice(v)?)),
            b"mesh" => Ok(s.mesh.push(read_slice(v)?)),
            b"collision" => Ok(s.collision.push(read_slice(v)?)),
            b"light" => Ok(s.light.push(read_slice(v)?)),
            b"armature" => Ok(s.armature.push(read_slice(v)?)),
            b"environment" => Ok(s.environment = Some(read_slice(v)?)),
            x => Ok(warn!(
                "unknown prefab key: {:?}",
                String::from_utf8_lossy(x)
            )),
        })?;
        Ok(s)
    }
}
impl ReadWrite for LightPart {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        write_kv_opt(w, b"name", &self.name)?;
        write_kv_opt(w, b"emission", &self.emission)?;
        write_kv_opt(w, b"radius", &self.radius)?;
        Ok(())
    }

    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut s = Self::default();
        read_kv_iter(r, |k, v| match k {
            b"name" => Ok(s.name = Some(read_slice(v)?)),
            b"emission" => Ok(s.emission = Some(read_slice(v)?)),
            b"radius" => Ok(s.radius = Some(read_slice(v)?)),
            x => Ok(warn!(
                "unknown light part key: {:?}",
                String::from_utf8_lossy(x)
            )),
        })?;
        Ok(s)
    }
}
impl ReadWrite for EnvironmentPart {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        write_kv_opt(w, b"skybox", &self.skybox)?;
        write_kv_opt(w, b"sun", &self.sun)?;
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut s = Self::default();
        read_kv_iter(r, |k, v| match k {
            b"skybox" => Ok(s.skybox = Some(read_slice(v)?)),
            b"sun" => Ok(s.sun = Some(read_slice(v)?)),
            x => Ok(warn!(
                "unknown environment part key: {:?}",
                String::from_utf8_lossy(x)
            )),
        })?;
        Ok(s)
    }
}
impl ReadWrite for MeshPart {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        write_kv_opt(w, b"name", &self.name)?;
        write_kv_opt(w, b"index", &self.index)?;
        write_kv_opt(w, b"armature", &self.armature)?;
        write_kv_opt(w, b"g_metallic", &self.g_metallic)?;
        write_kv_opt(w, b"g_roughness", &self.g_roughness)?;
        write_kv_opt(w, b"g_albedo", &self.g_albedo)?;
        write_kv_opt(w, b"g_transmission", &self.g_transmission)?;
        write_kv_opt(w, b"g_alpha", &self.g_alpha)?;
        write_kv_opt(w, b"g_emission", &self.g_emission)?;
        write_kv_opt(w, b"g_unlit", &self.g_unlit)?;
        write_kv_opt(w, b"g_double_sided", &self.g_double_sided)?;
        write_kv_opt(w, b"va_position", &self.va_position)?;
        write_kv_opt(w, b"va_normal", &self.va_normal)?;
        write_kv_opt(w, b"va_tangent", &self.va_tangent)?;
        write_kv_opt(w, b"va_texcoord", &self.va_texcoord)?;
        write_kv_opt(w, b"va_roughness", &self.va_roughness)?;
        write_kv_opt(w, b"va_metallic", &self.va_metallic)?;
        write_kv_opt(w, b"va_albedo", &self.va_albedo)?;
        write_kv_opt(w, b"va_transmission", &self.va_transmission)?;
        write_kv_opt(w, b"va_alpha", &self.va_transmission)?;
        write_kv_opt(w, b"va_emission", &self.va_emission)?;
        write_kv_opt(w, b"va_joint_weight", &self.va_joint_weight)?;
        write_kv_opt(w, b"va_joint_index", &self.va_joint_index)?;
        write_kv_opt(w, b"tex_normal", &self.tex_normal)?;
        write_kv_opt(w, b"tex_roughness", &self.tex_roughness)?;
        write_kv_opt(w, b"tex_metallic", &self.tex_metallic)?;
        write_kv_opt(w, b"tex_albedo", &self.tex_albedo)?;
        write_kv_opt(w, b"tex_transmission", &self.tex_transmission)?;
        write_kv_opt(w, b"tex_alpha", &self.tex_alpha)?;
        write_kv_opt(w, b"tex_emission", &self.tex_emission)?;
        write_kv_opt(w, b"tex_occlusion", &self.tex_occlusion)?;
        write_kv_opt(w, b"hint_mirror", &self.hint_mirror)?;
        write_kv_opt(w, b"hint_static", &self.hint_static)?;
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut s = Self::default();
        read_kv_iter(r, |k, v| match k {
            b"name" => Ok(s.name = Some(read_slice(v)?)),
            b"index" => Ok(s.index = Some(read_slice(v)?)),
            b"armature" => Ok(s.armature = Some(read_slice(v)?)),
            b"g_metallic" => Ok(s.g_metallic = Some(read_slice(v)?)),
            b"g_roughness" => Ok(s.g_roughness = Some(read_slice(v)?)),
            b"g_albedo" => Ok(s.g_albedo = Some(read_slice(v)?)),
            b"g_transmission" => Ok(s.g_transmission = Some(read_slice(v)?)),
            b"g_alpha" => Ok(s.g_alpha = Some(read_slice(v)?)),
            b"g_emission" => Ok(s.g_emission = Some(read_slice(v)?)),
            b"g_unlit" => Ok(s.g_unlit = Some(read_slice(v)?)),
            b"g_double_sided" => Ok(s.g_double_sided = Some(read_slice(v)?)),
            b"va_position" => Ok(s.va_position = Some(read_slice(v)?)),
            b"va_normal" => Ok(s.va_normal = Some(read_slice(v)?)),
            b"va_tangent" => Ok(s.va_tangent = Some(read_slice(v)?)),
            b"va_texcoord" => Ok(s.va_texcoord = Some(read_slice(v)?)),
            b"va_roughness" => Ok(s.va_roughness = Some(read_slice(v)?)),
            b"va_metallic" => Ok(s.va_metallic = Some(read_slice(v)?)),
            b"va_albedo" => Ok(s.va_albedo = Some(read_slice(v)?)),
            b"va_transmission" => Ok(s.va_transmission = Some(read_slice(v)?)),
            b"va_alpha" => Ok(s.va_alpha = Some(read_slice(v)?)),
            b"va_emission" => Ok(s.va_emission = Some(read_slice(v)?)),
            b"va_joint_weight" => Ok(s.va_joint_weight = Some(read_slice(v)?)),
            b"va_joint_index" => Ok(s.va_joint_index = Some(read_slice(v)?)),
            b"tex_normal" => Ok(s.tex_normal = Some(read_slice(v)?)),
            b"tex_roughness" => Ok(s.tex_roughness = Some(read_slice(v)?)),
            b"tex_metallic" => Ok(s.tex_metallic = Some(read_slice(v)?)),
            b"tex_albedo" => Ok(s.tex_albedo = Some(read_slice(v)?)),
            b"tex_transmission" => Ok(s.tex_transmission = Some(read_slice(v)?)),
            b"tex_alpha" => Ok(s.tex_alpha = Some(read_slice(v)?)),
            b"tex_emission" => Ok(s.tex_emission = Some(read_slice(v)?)),
            b"tex_occlusion" => Ok(s.tex_occlusion = Some(read_slice(v)?)),
            b"hint_mirror" => Ok(s.hint_mirror = Some(read_slice(v)?)),
            b"hint_static" => Ok(s.hint_static = Some(read_slice(v)?)),
            x => Ok(warn!(
                "unknown mesh part key: {:?}",
                String::from_utf8_lossy(x)
            )),
        })?;
        Ok(s)
    }
}

fn read_kv(r: &mut &[u8]) -> Result<(Vec<u8>, Vec<u8>)> {
    let mut key_size = [0; { size_of::<u16>() }];
    let mut value_size = [0; { size_of::<u16>() }];
    r.read_exact(&mut key_size)?;
    r.read_exact(&mut value_size)?;
    let key_size = u16::from_le_bytes(key_size);
    let value_size = u16::from_le_bytes(value_size);
    let mut key = vec![0; key_size as usize];
    let mut value = vec![0; value_size as usize];
    r.read_exact(&mut key)?;
    r.read_exact(&mut value)?;
    Ok((key, value))
}

fn read_kv_iter(r: &mut dyn Read, mut cb: impl FnMut(&[u8], &[u8]) -> Result<()>) -> Result<()> {
    let mut g = Vec::new();
    r.read_to_end(&mut g)?;
    let mut g = g.as_slice();
    while !g.is_empty() {
        let (k, v) = read_kv(&mut g)?;
        cb(&k, &v)?
    }
    Ok(())
}
fn read_slice<T: ReadWrite>(mut r: &[u8]) -> Result<T> {
    T::read(&mut r)
}

fn write_kv_opt(w: &mut dyn Write, key: &[u8], value: &Option<impl ReadWrite>) -> Result<()> {
    if let Some(v) = value {
        write_kv(w, key, &v.write_alloc())?;
    }
    Ok(())
}
fn write_kv(w: &mut dyn Write, key: &[u8], value: &[u8]) -> Result<()> {
    w.write_all(&(key.len() as u16).to_le_bytes())?;
    w.write_all(&(value.len() as u16).to_le_bytes())?;
    w.write_all(key)?;
    w.write_all(value)?;
    Ok(())
}

impl ReadWrite for Image<'_> {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        self.0.write(w)
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        Ok(Self(<Vec<u8> as ReadWrite>::read(r)?.into()))
    }
    fn write_alloc(&self) -> Cow<'_, [u8]> {
        self.0.write_alloc()
    }
}