use super::scene::MeshInstance;
use std::collections::HashMap;
use std::ops::Range;

#[derive(Clone)]
struct DrawState {
    group_id: usize,
    material_id: usize,
    sub_id: usize,
    instance_range: Range<u32>,
}

pub enum Command {
    BindMeshGroup {
        group_id: usize,
    },
    BindMaterial {
        material_id: usize,
    },
    Draw {
        sub_id: usize,
        instance_range: Range<u32>,
    },
}

pub struct CommandIterator<'a> {
    command_set: &'a CommandSet,
    command_idx: usize, // TODO use iterator instead lol
    last_group_id: usize,
    last_material_id: usize,
}

impl<'a> CommandIterator<'a> {
    pub fn new(command_set: &'a CommandSet) -> Self {
        Self {
            command_set,
            command_idx: 0,
            last_group_id: usize::MAX,
            last_material_id: usize::MAX,
        }
    }
}

impl<'a> Iterator for CommandIterator<'a> {
    type Item = Command;

    fn next(&mut self) -> Option<Self::Item> {
        let DrawState {
            group_id,
            sub_id,
            material_id,
            instance_range,
        } = self.command_set.draws.get(self.command_idx)?.clone();

        if group_id != self.last_group_id {
            self.last_group_id = group_id;
            Some(Command::BindMeshGroup { group_id })
        } else if material_id != self.last_material_id {
            self.last_material_id = material_id;
            Some(Command::BindMaterial { material_id })
        } else {
            self.command_idx += 1;
            Some(Command::Draw {
                sub_id,
                instance_range,
            })
        }
    }
}

pub struct CommandSet {
    transforms: Vec<[f32; 16]>,
    draws: Vec<DrawState>,
}

impl CommandSet {
    pub fn build(instances: &[MeshInstance]) -> Self {
        let mut sorted_meshes = HashMap::<(usize, usize), Vec<MeshInstance>>::new();
        for instance in instances.iter() {
            let group_id = instance.mesh.group_id;
            let material_id = instance.material.id;
            let key = (group_id, material_id);
            if let Some(by_state) = sorted_meshes.get_mut(&key) {
                by_state.push(*instance);
            } else {
                sorted_meshes.insert(key, vec![*instance]);
            }
        }

        let mut transforms = Vec::new();
        let mut draws = Vec::new();

        for ((group_id, material_id), mut meshes) in sorted_meshes.drain() {
            let mut sorted_subs = HashMap::<usize, Vec<MeshInstance>>::new();
            for mesh in meshes.drain(..) {
                let sub_id = mesh.mesh.sub_id;
                if let Some(by_sub) = sorted_subs.get_mut(&sub_id) {
                    by_sub.push(mesh);
                } else {
                    sorted_subs.insert(sub_id, vec![mesh]);
                }
            }

            for (sub_id, meshes) in sorted_subs.drain() {
                let start_idx = transforms.len() as u32;
                let as_arrays = meshes.iter().map(|i| i.transform.to_cols_array());
                transforms.extend(as_arrays);
                let end_idx = transforms.len() as u32;
                let instance_range = start_idx..end_idx;
                draws.push(DrawState {
                    group_id,
                    material_id,
                    sub_id,
                    instance_range,
                });
            }
        }

        Self { transforms, draws }
    }

    pub fn get_storage(&self) -> &[u8] {
        bytemuck::cast_slice(&self.transforms)
    }

    pub fn iter(&self) -> CommandIterator {
        CommandIterator::new(self)
    }
}