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

#[derive(Debug, Default, Clone)]
pub struct Prefab(pub Vec<(Affine3A, Resource)>);

/// Combinations of va_* and tex_* are multiplied except normal which is added.
/// Defaults should be the identity for that operation, so default is 1 (or white) except normals are Vec3::ZERO.
#[derive(Debug, Default, Clone)]
pub struct Part {
    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 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 tex_normal: Option<Resource>,
    pub tex_roughness: Option<Resource>,
    pub tex_metallic: Option<Resource>,
    pub tex_albedo: Option<Resource>,
    pub tex_transmission: Option<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 Prefab {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        for (aff, res) in self.0.clone() {
            aff.write(w)?;
            res.write(w)?;
        }
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut s = Prefab::default();
        let mut g = Vec::new();
        r.read_to_end(&mut g)?;
        let mut g = g.as_slice();
        while !g.is_empty() {
            s.0.push((Affine3A::read(&mut g)?, Resource::read(&mut g)?))
        }
        Ok(s)
    }
}
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 ReadWrite for Part {
    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"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"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)?;

        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut s = Self::default();
        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)?;
            let mut v = v.as_slice();
            match k.as_slice() {
                b"index" => s.index = Some(<_ as ReadWrite>::read(&mut v)?),
                b"g_metallic" => s.g_metallic = Some(<_ as ReadWrite>::read(&mut v)?),
                b"g_roughness" => s.g_roughness = Some(<_ as ReadWrite>::read(&mut v)?),
                b"g_albedo" => s.g_albedo = Some(<_ as ReadWrite>::read(&mut v)?),
                b"g_transmission" => s.g_transmission = Some(<_ as ReadWrite>::read(&mut v)?),
                b"va_position" => s.va_position = Some(<_ as ReadWrite>::read(&mut v)?),
                b"va_normal" => s.va_normal = Some(<_ as ReadWrite>::read(&mut v)?),
                b"va_texcoord" => s.va_texcoord = Some(<_ as ReadWrite>::read(&mut v)?),
                b"va_roughness" => s.va_roughness = Some(<_ as ReadWrite>::read(&mut v)?),
                b"va_metallic" => s.va_metallic = Some(<_ as ReadWrite>::read(&mut v)?),
                b"va_albedo" => s.va_albedo = Some(<_ as ReadWrite>::read(&mut v)?),
                b"va_transmission" => {
                    s.va_transmission = Some(<_ as ReadWrite>::read(&mut v)?)
                }
                b"tex_normal" => s.tex_normal = Some(<_ as ReadWrite>::read(&mut v)?),
                b"tex_roughness" => s.tex_roughness = Some(<_ as ReadWrite>::read(&mut v)?),
                b"tex_metallic" => s.tex_metallic = Some(<_ as ReadWrite>::read(&mut v)?),
                b"tex_albedo" => s.tex_albedo = Some(<_ as ReadWrite>::read(&mut v)?),
                b"tex_transmission" => {
                    s.tex_transmission = Some(<_ as ReadWrite>::read(&mut v)?)
                }
                x => warn!("unknown 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 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))?))
    }
}