use crate::download::Downloader;
use anyhow::{Context, Result};
use log::debug;
use std::{
    collections::{HashMap, HashSet},
    sync::Arc,
};
use weareshared::{
    Affine3A,
    packets::{ReadWrite, Resource},
    resources::{Attribute, Part, Prefab},
};
use wgpu::{
    Buffer, BufferUsages, Device,
    util::{BufferInitDescriptor, DeviceExt},
};

pub struct ScenePreparer {
    index_buffers: HashMap<Resource, (Arc<Buffer>, u32)>,
    index_buffers_needed: HashSet<Resource>,
    vertex_buffers: HashMap<Resource, Arc<Buffer>>,
    vertex_buffers_needed: HashSet<Resource>,
    parts: HashMap<Resource, Arc<RPart>>,
    parts_needed: HashSet<Resource>,
    pub prefabs: HashMap<Resource, RPrefab>,
    pub prefabs_needed: HashSet<Resource>,

    device: Arc<Device>,
}

pub struct RPrefab(pub Vec<(Affine3A, Arc<RPart>)>);
pub struct RPart {
    pub index_count: u32,
    pub index: Arc<Buffer>,
    pub position: [Arc<Buffer>; 3],
    pub normal: [Arc<Buffer>; 3],
    pub texcoord: [Arc<Buffer>; 2],
}

impl ScenePreparer {
    pub fn new(device: Arc<Device>) -> Self {
        Self {
            index_buffers: HashMap::new(),
            vertex_buffers: HashMap::new(),
            vertex_buffers_needed: HashSet::new(),
            parts: HashMap::new(),
            parts_needed: HashSet::new(),
            prefabs: HashMap::new(),
            prefabs_needed: HashSet::new(),
            index_buffers_needed: HashSet::new(),
            device,
        }
    }
    pub fn update(&mut self, dls: &mut Downloader) -> Result<()> {
        let mut done = Vec::new();
        for pres in &self.prefabs_needed {
            if let Some(buf) = dls.try_get(*pres)? {
                let prefab = Prefab::read(&mut buf.as_slice()).context("parsing prefab")?;
                let mut rprefab = RPrefab(Vec::new());
                for (aff, partres) in &prefab.0 {
                    if let Some(part) = self.parts.get(partres) {
                        rprefab.0.push((*aff, part.clone()));
                    } else {
                        self.parts_needed.insert(*partres);
                    }
                }
                if rprefab.0.len() == prefab.0.len() {
                    self.prefabs.insert(*pres, rprefab);
                    debug!("prefab created ({pres})");
                    done.push(*pres);
                }
            }
        }
        for pres in &self.index_buffers_needed {
            if let Some(buf) = dls.try_get(*pres)? {
                let buf = buf
                    .into_iter()
                    .array_chunks::<2>()
                    .map(u16::from_be_bytes)
                    .map(u16::to_le_bytes)
                    .flatten()
                    .collect::<Vec<_>>();
                let buffer = self.device.create_buffer_init(&BufferInitDescriptor {
                    contents: &buf,
                    label: None,
                    usage: BufferUsages::INDEX | BufferUsages::COPY_DST,
                });
                self.index_buffers
                    .insert(*pres, (Arc::new(buffer), buf.len() as u32 / 2));
                debug!("index buffer created (len={}) {pres}", buf.len() / 6);
                done.push(*pres);
            }
        }
        for pres in &self.vertex_buffers_needed {
            if let Some(buf) = dls.try_get(*pres)? {
                let buf = buf
                    .into_iter()
                    .array_chunks::<4>()
                    .map(f32::from_be_bytes)
                    .map(f32::to_le_bytes)
                    .flatten()
                    .collect::<Vec<_>>();
                let buffer = self.device.create_buffer_init(&BufferInitDescriptor {
                    contents: &buf,
                    label: None,
                    usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
                });
                self.vertex_buffers.insert(*pres, Arc::new(buffer));
                debug!(
                    "vertex attribute buffer created (len={}) {pres}",
                    buf.len() / 4
                );
                done.push(*pres);
            }
        }
        for pres in &self.parts_needed {
            if let Some(buf) = dls.try_get(*pres)? {
                let part = Part::read(&mut buf.as_slice()).context("parsing part")?;
                if let (Some(indexres), Some(positionres), Some(normalres), Some(texcoordres)) = (
                    part.index,
                    part.va_position,
                    part.va_normal,
                    part.va_texcoord,
                ) {
                    let Some((index, index_count)) = self.index_buffers.get(&indexres).cloned()
                    else {
                        self.index_buffers_needed.insert(indexres);
                        continue;
                    };
                    let mut position = Vec::new();
                    for vr in positionres {
                        match vr {
                            Attribute::Constant(_) => todo!(),
                            Attribute::Vertex(resource) => {
                                if let Some(vertex) = self.vertex_buffers.get(&resource).cloned() {
                                    position.push(vertex);
                                } else {
                                    self.vertex_buffers_needed.insert(resource);
                                };
                            }
                            Attribute::Texture(_resource, _ch) => todo!(),
                        }
                    }
                    let mut normal = Vec::new();
                    for vr in normalres {
                        match vr {
                            Attribute::Constant(_) => todo!(),
                            Attribute::Vertex(resource) => {
                                if let Some(vertex) = self.vertex_buffers.get(&resource).cloned() {
                                    normal.push(vertex);
                                } else {
                                    self.vertex_buffers_needed.insert(resource);
                                };
                            }
                            Attribute::Texture(_resource, _ch) => todo!(),
                        }
                    }
                    let mut texcoord = Vec::new();
                    for vr in texcoordres {
                        match vr {
                            Attribute::Constant(_) => todo!(),
                            Attribute::Vertex(resource) => {
                                if let Some(vertex) = self.vertex_buffers.get(&resource).cloned() {
                                    texcoord.push(vertex);
                                } else {
                                    self.vertex_buffers_needed.insert(resource);
                                };
                            }
                            Attribute::Texture(_resource, _ch) => todo!(),
                        }
                    }
                    if texcoord.len() == 2 && normal.len() == 3 && position.len() == 3 {
                        debug!("part created ({pres})");
                        self.parts.insert(
                            *pres,
                            Arc::new(RPart {
                                index_count,
                                index,
                                texcoord: texcoord.try_into().unwrap(),
                                normal: normal.try_into().unwrap(),
                                position: position.try_into().unwrap(),
                            }),
                        );
                        done.push(*pres);
                    }
                }
            }
        }

        for d in done {
            self.parts_needed.remove(&d);
            self.prefabs_needed.remove(&d);
            self.index_buffers_needed.remove(&d);
            self.vertex_buffers_needed.remove(&d);
        }
        Ok(())
    }
}