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

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

#[derive(Debug, Default, Clone)]
pub struct Part {
    pub index: Option<Resource>,
    pub va_position: Option<[Attribute; 3]>,
    pub va_normal: Option<[Attribute; 3]>,
    pub va_texcoord: Option<[Attribute; 2]>,
    pub va_pbr_roughness: Option<Attribute>,
    pub va_pbr_metallic: Option<Attribute>,
    pub va_pbr_albedo: Option<[Attribute; 3]>,
    pub va_pbr_transmission: Option<Attribute>,
}

#[derive(Debug, Clone)]
pub enum Attribute {
    Constant(f32),
    Vertex(Resource),
    Texture(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)?;
        while !g.is_empty() {
            s.0.push((Affine3A::read(r)?, Resource::read(r)?))
        }
        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<()> {
        if let Some(a) = &self.index {
            write_kv(w, b"index", &a.write_alloc())?;
        }
        if let Some(a) = &self.va_position {
            write_kv(w, b"va_position", &a.write_alloc())?;
        }
        if let Some(a) = &self.va_normal {
            write_kv(w, b"va_normal", &a.write_alloc())?;
        }
        if let Some(a) = &self.va_texcoord {
            write_kv(w, b"va_texcoord", &a.write_alloc())?;
        }
        if let Some(a) = &self.va_pbr_roughness {
            write_kv(w, b"va_pbr_roughness", &a.write_alloc())?;
        }
        if let Some(a) = &self.va_pbr_metallic {
            write_kv(w, b"va_pbr_metallic", &a.write_alloc())?;
        }
        if let Some(a) = &self.va_pbr_albedo {
            write_kv(w, b"va_pbr_albedo", &a.write_alloc())?;
        }
        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"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_pbr_roughness" => s.va_pbr_roughness = Some(<_ as ReadWrite>::read(&mut v)?),
                b"va_pbr_metallic" => s.va_pbr_metallic = Some(<_ as ReadWrite>::read(&mut v)?),
                b"va_pbr_albedo" => s.va_pbr_albedo = Some(<_ as ReadWrite>::read(&mut v)?),
                b"va_pbr_transmission" => {
                    s.va_pbr_transmission = Some(<_ as ReadWrite>::read(&mut v)?)
                }
                _ => warn!("unknown part key"),
            }
        }
        Ok(s)
    }
}
fn read_kv(r: &mut dyn Read) -> 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 value)?;
    r.read_exact(&mut key)?;
    Ok((key, value))
}

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 Attribute {
    fn write(&self, w: &mut dyn Write) -> Result<()> {
        match self {
            Attribute::Constant(v) => {
                w.write_all(&[0x01])?;
                w.write_all(&v.to_be_bytes())?;
            }
            Attribute::Vertex(resource) => {
                w.write_all(&[0x02])?;
                resource.write(w)?;
            }
            Attribute::Texture(resource) => {
                w.write_all(&[0x03])?;
                resource.write(w)?;
            }
        }
        Ok(())
    }
    fn read(r: &mut dyn Read) -> Result<Self> {
        let mut tag = [0u8; 1];
        r.read_exact(&mut tag)?;
        Ok(match tag[0] {
            0x01 => Self::Constant(f32::read(r)?),
            0x02 => Self::Vertex(Resource::read(r)?),
            0x03 => Self::Texture(Resource::read(r)?),
            _ => bail!("unknown attribute tag"),
        })
    }
}

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))?))
    }
}