use crate::staging::*;
use parking_lot::{RwLock, RwLockReadGuard};
use slab::Slab;
use smallvec::SmallVec;
use std::any::TypeId;
use std::collections::HashMap;
use std::sync::Arc;

/// An error that can be returned when allocating a mesh.
#[derive(Debug)]
pub enum PoolError {
    TooBig,
    NoMoreRoom,
    InvalidIndex,
    AttrTaken,
    AttrUnregistered,
    LayoutUnregistered,
    MismatchedId,
    MismatchedLayout,
}

/// An identifier for a mesh attribute.
#[repr(transparent)]
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub struct AttrId(usize);

/// A description of an attribute's layout in memory.
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub struct AttrLayout {
    /// The size (in bytes) of this attribute.
    pub size: usize,
}

/// Information about an [Attribute] registered in [AttrStore].
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub struct AttrInfo {
    pub layout: AttrLayout,
    pub usages: wgpu::BufferUsages,
    pub default_pool_size: usize,
}

/// The data single mesh attribute.
#[derive(Clone)]
pub struct AttrBuffer {
    pub id: AttrId,
    pub count: usize,
    pub data: Vec<u8>,
}

/// A compile-time attribute data type.
pub trait Attribute: Sized {
    /// The memory layout of this data.
    fn get_layout() -> AttrLayout {
        AttrLayout {
            size: std::mem::size_of::<Self>(),
        }
    }

    /// The [wgpu::BufferUsages] of this data.
    fn get_usages() -> wgpu::BufferUsages;

    /// The default size for new pools of this attribute.
    ///
    /// Defaults to 1024 * 1024. (Around one million.)
    fn get_default_pool_size() -> usize {
        1024 * 1024
    }
}

/// A store of [AttrIds][AttrId].
pub struct AttrStore {
    attributes: RwLock<Slab<AttrInfo>>,
    types: RwLock<HashMap<TypeId, AttrId>>,
}

impl AttrStore {
    pub fn new() -> Arc<Self> {
        Arc::new(Self {
            attributes: Default::default(),
            types: Default::default(),
        })
    }

    /// Dynamically creates a new [AttrId].
    pub fn add(&self, info: AttrInfo) -> AttrId {
        let index = self.attributes.write().insert(info);
        AttrId(index)
    }

    /// Gets the [AttrId] for a type implementing [Attribute].
    ///
    /// Creates a new [AttrId] for unrecognized types, otherwise reuses an
    /// existing [AttrId].
    pub fn get_type<T: 'static + Attribute>(&self) -> AttrId {
        let type_id = TypeId::of::<T>();
        let existing_id = self.types.read().get(&type_id).copied();

        if let Some(id) = existing_id {
            id
        } else {
            let layout = T::get_layout();
            let usages = T::get_usages();
            let default_pool_size = T::get_default_pool_size();
            let info = AttrInfo {
                layout,
                usages,
                default_pool_size,
            };
            let id = self.add(info);
            self.types.write().insert(type_id, id);
            id
        }
    }

    /// Gets the [AttrInfo] for an [AttrId].
    pub fn get_info(&self, id: &AttrId) -> Option<AttrInfo> {
        self.attributes.read().get(id.0).copied()
    }
}

/// Attribute pool allocation location.
#[derive(Clone)]
pub struct AttrAllocKey {
    /// The target attribute.
    pub attr: AttrId,

    /// The index of the attribute pool.
    pub pool: usize,

    /// The allocation within the attribute pool.
    pub alloc: usize,
}

/// Info about an array of attributes that has been allocated in an [AttrPool].
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub struct AttrAlloc {
    pub offset: usize,
    pub count: usize,
    pub offset_bytes: usize,
    pub count_bytes: usize,
}

/// An unused space range in an [AttrPool].
pub struct FreeSpace {
    offset: usize,
    count: usize,
}

/// A single GPU buffer containing linear arrays of attributes.
pub struct AttrPool {
    buffer: wgpu::Buffer,
    id: AttrId,
    pool_id: usize,
    layout: AttrLayout,
    size: usize,
    allocs: Slab<AttrAlloc>,
    free_space: Vec<FreeSpace>,
}

impl PartialEq for AttrPool {
    fn eq(&self, other: &Self) -> bool {
        self.pool_id == other.pool_id && self.id == other.id
    }
}

impl AttrPool {
    pub fn new(device: &wgpu::Device, id: AttrId, pool_id: usize, info: AttrInfo) -> Self {
        let layout = info.layout;
        let size = info.default_pool_size;

        // TODO debug strings for attributes + pools + buffers
        let buffer = device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            usage: wgpu::BufferUsages::COPY_DST | info.usages,
            mapped_at_creation: false,
            size: (size * layout.size) as wgpu::BufferAddress,
        });

        Self {
            buffer,
            id,
            pool_id,
            layout,
            size,
            free_space: vec![FreeSpace {
                offset: 0,
                count: size,
            }],
            allocs: Default::default(),
        }
    }

    /// Loads an [AttrBuffer].
    ///
    /// Returns the new [AttrAllocKey].
    pub fn load(&mut self, buf: &AttrBuffer) -> Result<AttrAllocKey, PoolError> {
        let best_index = self.can_load(buf)?;
        self.alloc_at(best_index, buf.count)
    }

    /// Allocates an array of attributes.
    ///
    /// Returns the new [AttrAllocKey].
    pub fn alloc(&mut self, count: usize) -> Result<AttrAllocKey, PoolError> {
        let best_index = self.can_alloc(count)?;
        self.alloc_at(best_index, count)
    }

    /// Tests if an [AttrBuffer] can be loaded.
    ///
    /// Returns the result of [Self::best_fit].
    pub fn can_load(&self, buf: &AttrBuffer) -> Result<usize, PoolError> {
        if buf.id != self.id {
            Err(PoolError::MismatchedId)
        } else if buf.count * self.layout.size != buf.data.len() {
            Err(PoolError::MismatchedLayout)
        } else {
            self.can_alloc(buf.count)
        }
    }

    /// Tests if an array of attributes can be allocated.
    ///
    /// Returns the result of [Self::best_fit].
    pub fn can_alloc(&self, count: usize) -> Result<usize, PoolError> {
        if count > self.size {
            Err(PoolError::TooBig)
        } else if let Some(best_index) = self.best_fit(count) {
            Ok(best_index)
        } else {
            Err(PoolError::NoMoreRoom)
        }
    }

    /// Finds the index of the best-fit free space for an array of attributes.
    ///
    /// TODO: use a binary tree to find best-fit free space in logarithmic time
    fn best_fit(&self, count: usize) -> Option<usize> {
        let mut best_index = None;
        let mut best_count = usize::MAX;
        for (index, space) in self.free_space.iter().enumerate() {
            if space.count >= count && space.count < best_count {
                best_index = Some(index);
                best_count = space.count;
            }
        }

        best_index
    }

    /// Allocates an array of attribute at a specific free space index.
    ///
    /// Returns the new [AttrAllocKey].
    fn alloc_at(&mut self, index: usize, count: usize) -> Result<AttrAllocKey, PoolError> {
        let free_space = match self.free_space.get_mut(index) {
            Some(index) => index,
            None => return Err(PoolError::InvalidIndex),
        };

        let offset = free_space.offset;

        let alloc = AttrAlloc {
            offset,
            count,
            offset_bytes: offset * self.layout.size,
            count_bytes: count * self.layout.size,
        };

        let index = self.allocs.insert(alloc);

        let key = AttrAllocKey {
            attr: self.id,
            pool: self.pool_id,
            alloc: index,
        };

        use std::cmp::Ordering;
        match free_space.count.cmp(&count) {
            Ordering::Less => {
                return Err(PoolError::TooBig);
            }
            Ordering::Equal => {
                self.free_space.remove(index);
            }
            Ordering::Greater => {
                free_space.count -= count;
                free_space.offset += count;
            }
        }

        Ok(key)
    }

    /// Frees an allocation (by key) from the pool.
    pub fn free(&mut self, key: usize) -> Result<(), PoolError> {
        let alloc = self.allocs.try_remove(key).ok_or(PoolError::InvalidIndex)?;

        let free_space = FreeSpace {
            offset: alloc.offset,
            count: alloc.count,
        };

        // TODO merge free spaces
        self.free_space.push(free_space);

        Ok(())
    }

    /// Retrieves an [AttrAlloc] by key.
    pub fn get(&self, key: usize) -> Option<AttrAlloc> {
        self.allocs.get(key).copied()
    }

    /// Gets this pool's internal GPU buffer.
    pub fn get_buffer(&self) -> &wgpu::Buffer {
        &self.buffer
    }

    /// Gets a [CopyDest] for an allocation, by key.
    pub fn get_copy_dest(&self, key: usize) -> Result<CopyDest, PoolError> {
        let offset = self.get(key).ok_or(PoolError::InvalidIndex)?.offset_bytes;
        Ok(CopyDest {
            buffer: self.get_buffer(),
            offset,
        })
    }
}

/// The number of attributes a mesh can have before they're moved to the heap.
pub const MAX_MESH_INLINE_ATTRIBUTES: usize = 16;

/// A mesh and all of its attributes.
#[derive(Clone, Default)]
pub struct MeshBuffer {
    pub attributes: SmallVec<[AttrBuffer; MAX_MESH_INLINE_ATTRIBUTES]>,
}

/// An allocated mesh.
pub struct MeshAlloc {
    pub attributes: SmallVec<[AttrAllocKey; MAX_MESH_INLINE_ATTRIBUTES]>,
}

/// A handle to an allocated mesh.
#[repr(transparent)]
#[derive(Copy, Clone, Debug)]
pub struct MeshHandle(usize);

/// A reusable set of [AttrIds][AttrId], for use with querying compatible meshes.
pub type MeshLayoutDesc = smallmap::Set<AttrId>;

/// The ID of a [MeshLayoutDesc] registered in a [MeshPool].
#[repr(transparent)]
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub struct MeshLayoutId(usize);

/// Mappings of the attributes in a [MeshLayout] to specific pool indices.
pub type MeshLayoutBindingIndices = smallmap::Map<AttrId, usize>;

/// Mappings of the attributes in a [MeshLayout] to specific pools.
pub struct MeshLayoutBindings<'a> {
    lock: RwLockReadGuard<'a, HashMap<AttrId, Vec<AttrPool>>>,
    indices: MeshLayoutBindingIndices,
}

impl<'a> MeshLayoutBindings<'a> {
    pub fn get(&self, attr: AttrId) -> Option<&AttrPool> {
        let pool_id = self.indices.get(&attr)?;
        let pools = self.lock.get(&attr)?;
        pools.get(*pool_id)
    }
}

/// Mappings of the attributes in a [MeshAlloc] to specific pool offsets.
pub type MeshAllocInfos = SmallVec<[(AttrId, AttrAlloc); MAX_MESH_INLINE_ATTRIBUTES]>;

/// A set of mesh instances fitting a common [MeshLayoutDesc].
#[derive(Default)]
pub struct MeshLayoutInstances<T> {
    pub bindings: MeshLayoutBindingIndices,
    pub instances: Vec<(T, MeshAllocInfos)>,
}

/// A mesh data pool.
pub struct MeshPool {
    device: Arc<wgpu::Device>,
    staging: StagingPool<AttrAllocKey>,
    attr_store: Arc<AttrStore>,
    allocs: RwLock<Slab<MeshAlloc>>,
    mesh_layouts: RwLock<Slab<MeshLayoutDesc>>,
    pools: RwLock<HashMap<AttrId, Vec<AttrPool>>>,
}

impl MeshPool {
    pub fn new(device: Arc<wgpu::Device>, attr_store: Arc<AttrStore>) -> Arc<Self> {
        Arc::new(Self {
            device: device.clone(),
            staging: StagingPool::new(device, 1024 * 1024),
            attr_store,
            mesh_layouts: Default::default(),
            allocs: Default::default(),
            pools: Default::default(),
        })
    }

    /// Registers a [MeshLayoutDesc] with this pool.
    ///
    /// TODO: keep track of mesh allocations that fit each layout
    pub fn add_layout(&self, layout: MeshLayoutDesc) -> Result<MeshLayoutId, PoolError> {
        let idx = self.mesh_layouts.write().insert(layout);
        Ok(MeshLayoutId(idx))
    }

    /// Loads a [MeshBuffer].
    pub fn load(&self, buf: MeshBuffer) -> Result<MeshHandle, PoolError> {
        let mut attrs = HashMap::new();

        for attr in buf.attributes.into_iter() {
            if let Some(_) = attrs.insert(attr.id, attr) {
                // TODO: support for loading duplicate attribute IDs
                return Err(PoolError::AttrTaken);
            }
        }

        let mut attr_allocs = SmallVec::new();
        let mut copies = Vec::new();
        for (id, buf) in attrs.drain() {
            let mut pools_write = self.pools.write();

            let pools = match pools_write.get_mut(&id) {
                Some(pools) => pools,
                None => {
                    let info = match self.attr_store.get_info(&id) {
                        Some(info) => info,
                        None => return Err(PoolError::AttrUnregistered),
                    };

                    let pool = AttrPool::new(&self.device, id, 0, info);
                    pools_write.insert(id, vec![pool]);
                    pools_write.get_mut(&id).unwrap()
                }
            };

            for pool in pools.iter_mut() {
                match pool.load(&buf) {
                    Ok(alloc) => {
                        let copy = CopyBuffer {
                            target: alloc.clone(),
                            offset: 0,
                            data: buf.data,
                        };

                        attr_allocs.push(alloc);
                        copies.push(copy);
                        break;
                    }
                    Err(PoolError::NoMoreRoom) => {}
                    Err(e) => return Err(e),
                }
            }

            // TODO create a new pool if no available pool was found
        }

        self.staging.queue_copies(copies);

        let alloc = MeshAlloc {
            attributes: attr_allocs,
        };
        let key = self.allocs.write().insert(alloc);
        let handle = MeshHandle(key);
        Ok(handle)
    }

    pub fn flush(&self, commands: &mut wgpu::CommandEncoder) {
        let pools_read = self.pools.read();

        let get_dst = |target: &AttrAllocKey| {
            pools_read
                .get(&target.attr)
                .unwrap()
                .get(target.pool)
                .unwrap()
                .get_copy_dest(target.alloc)
                .unwrap()
        };

        // TODO: keep track of loaded/unloaded meshes
        let on_complete = |_target| {};

        self.staging.flush(commands, get_dst, on_complete);
    }

    pub fn iter_meshes<T, I, F>(
        &self,
        layout: MeshLayoutId,
        meshes: I,
        get_handle: F,
    ) -> Result<Vec<MeshLayoutInstances<T>>, PoolError>
    where
        I: Iterator<Item = T>,
        F: Fn(&T) -> &MeshHandle,
    {
        let layout = self
            .mesh_layouts
            .read()
            .get(layout.0)
            .ok_or(PoolError::LayoutUnregistered)?
            .clone();

        let layout_attrs: Vec<AttrId> = layout.iter().map(|(id, _)| *id).collect();
        let mut layouts = Vec::<MeshLayoutInstances<T>>::new();

        let allocs_read = self.allocs.read();
        let mut attr_allocs: Vec<AttrAllocKey> = Vec::with_capacity(layout_attrs.len());
        for mesh in meshes {
            let handle = get_handle(&mesh);
            let alloc = allocs_read.get(handle.0).ok_or(PoolError::InvalidIndex)?;

            attr_allocs.clear();
            for layout_attr in layout_attrs.iter() {
                match alloc
                    .attributes
                    .iter()
                    .find(|attr_alloc| attr_alloc.attr == *layout_attr)
                {
                    Some(alloc) => attr_allocs.push(alloc.clone()),
                    None => break,
                }
            }

            if attr_allocs.len() != layout_attrs.len() {
                continue;
            }

            let mut layout_bindings = MeshLayoutBindingIndices::default();
            let mut alloc_infos = MeshAllocInfos::default();
            for alloc in attr_allocs.iter() {
                let pools_read = self.pools.read();
                let pools = pools_read
                    .get(&alloc.attr)
                    .ok_or(PoolError::AttrUnregistered)?;

                let pool = pools.get(alloc.pool).ok_or(PoolError::InvalidIndex)?;
                let alloc_info = pool.get(alloc.alloc).ok_or(PoolError::InvalidIndex)?;

                layout_bindings.insert(alloc.attr, alloc.pool);
                alloc_infos.push((alloc.attr, alloc_info));
            }

            let instance = (mesh, alloc_infos);

            match layouts
                .iter_mut()
                .find(|layout| layout.bindings == layout_bindings)
            {
                Some(layout) => layout.instances.push(instance),
                None => layouts.push(MeshLayoutInstances {
                    bindings: layout_bindings,
                    instances: vec![instance],
                }),
            }
        }

        Ok(layouts)
    }

    pub fn get_bindings<'a>(&'a self, indices: MeshLayoutBindingIndices) -> MeshLayoutBindings<'a> {
        let lock = self.pools.read();
        MeshLayoutBindings { lock, indices }
    }
}