use crate::packets::{ReadWrite, Resource};
use anyhow::Result;
use glam::{Affine3A, Vec3A};
use log::warn;
use std::{
    collections::HashMap,
    io::{Read, Write},
};

#[derive(Debug, Default, Clone)]
pub struct Prefab {
    pub mesh: Vec<(Affine3A, Resource)>,
    pub light: Vec<(Vec3A, Resource)>,
    pub environment: Option<EnvironmentPart>,
}

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

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

#[derive(Debug, Default, Clone)]
pub struct MeshPart {
    pub index: Option<Resource>,
    pub g_metallic: Option<f32>,
    pub g_roughness: Option<f32>,
    pub g_albedo: Option<Vec3A>,
    pub g_transmission: Option<f32>,
    pub g_emission: Option<Vec3A>,
    pub va_position: Option<[Resource; 3]>,
    pub va_normal: Option<[Resource; 3]>,
    pub va_texcoord: Option<[Resource; 2]>,
    pub va_roughness: Option<Resource>,
    pub va_metallic: Option<Resource>,
    pub va_albedo: Option<[Resource; 3]>,
    pub va_transmission: Option<Resource>,
    pub va_emission: Option<[Resource; 3]>,
    pub tex_normal: Option<Resource>,
    pub tex_roughness: Option<Resource>,
    pub tex_metallic: Option<Resource>,
    pub tex_albedo: Option<Resource>,
    pub tex_transmission: Option<Resource>,
    pub tex_emission: Option<Resource>,
}

pub struct PrefabIndex(pub HashMap<String, Resource>);

#[derive(Debug, Default, Clone)]
pub struct AttributeArray(pub Vec<f32>);
#[derive(Debug, Default, Clone)]
pub struct IndexArray(pub Vec<[u16; 3]>);

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(HashMap::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 Prefab {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        for x in &self.mesh {
            write_kv_opt(w, b"mesh", &Some(*x))?;
        }
        for x in &self.light {
            write_kv_opt(w, b"light", &Some(*x))?;
        }
        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"mesh" => Ok(s.mesh.push(read_slice(v)?)),
            b"light" => Ok(s.light.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"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"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"index", &self.index)?;
        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_emission", &self.g_emission)?;
        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_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_emission", &self.va_emission)?;
        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_emission", &self.tex_emission)?;
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut s = Self::default();
        read_kv_iter(r, |k, v| match k {
            b"index" => Ok(s.index = 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_emission" => Ok(s.g_emission = 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_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_emission" => Ok(s.va_emission = 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_emission" => Ok(s.tex_emission = 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; 2];
    let mut value_size = [0; 2];
    r.read_exact(&mut key_size)?;
    r.read_exact(&mut value_size)?;
    let key_size = u16::from_be_bytes(key_size);
    let value_size = u16::from_be_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_be_bytes())?;
    w.write_all(&(value.len() as u16).to_be_bytes())?;
    w.write_all(key)?;
    w.write_all(value)?;
    Ok(())
}

impl ReadWrite for u8 {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        w.write_all(&[*self])?;
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut buf = [0u8; 1];
        r.read_exact(&mut buf)?;
        Ok(buf[0])
    }
}

impl ReadWrite for Vec3A {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        w.write_all(&self.x.to_be_bytes())?;
        w.write_all(&self.y.to_be_bytes())?;
        w.write_all(&self.z.to_be_bytes())?;
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        Ok(Self::new(f32::read(r)?, f32::read(r)?, f32::read(r)?))
    }
}
impl ReadWrite for Affine3A {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        for v in self.to_cols_array() {
            v.write(w)?
        }
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        Ok(Self::from_cols_array(&[(); 12].try_map(|()| f32::read(r))?))
    }
}
impl ReadWrite for IndexArray {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        for x in self.0.clone() {
            w.write_all(x.map(|x| x.to_be_bytes()).as_flattened())?;
        }
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut s = Self(Vec::new());
        let mut g = Vec::new();
        r.read_to_end(&mut g)?;
        for x in g.iter().array_chunks::<2>().array_chunks::<3>() {
            s.0.push(x.map(|x| u16::from_be_bytes(x.map(|x| *x))))
        }
        Ok(s)
    }
}
impl ReadWrite for AttributeArray {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        for x in self.0.clone() {
            w.write_all(&x.to_be_bytes())?;
        }
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut s = Self(Vec::new());
        let mut g = Vec::new();
        r.read_to_end(&mut g)?;
        for x in g.iter().array_chunks::<4>() {
            s.0.push(f32::from_be_bytes(x.map(|x| *x)))
        }
        Ok(s)
    }
}
impl<A: ReadWrite, B: ReadWrite> ReadWrite for (A, B) {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        self.0.write(w)?;
        self.1.write(w)?;
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        Ok((A::read(r)?, B::read(r)?))
    }
}