use crate::idtree;
use crate::storage::{Storage, HierarchicalOneToNStorage, IntoIter, IntoIterMut};
use densevec::DenseVec;
use crate::sync::{ReadGuardRef, WriteGuardRef};
use crate::Component;

use std::slice;
use std::mem;
use std::fmt::{self, Debug};
use crate::utils::*;

pub struct OneToNForest<T>{
    arena: idtree::Arena<T>,
    entities_roots: DenseVec<Vec<idtree::NodeId>>,
    reverse_index: DenseVec<idtree::NodeId>,
}

impl<'a, T: 'a + Component> Storage<'a, T> for OneToNForest<T>{
    type Get = RootsIter<'a, T>;
    type GetMut = RootsIterMut<'a, T>;
    type DerefTarget = RootsIter<'a, T>;
    type IdedIter = RawIdedIter<'a, T>;
    type Iter = RawIter<'a, T>;
    type IterMut = RawIterMut<'a, T>;

    fn new() -> OneToNForest<T>{
        OneToNForest{
            arena: idtree::Arena::new(),
            entities_roots: DenseVec::new(),
            reverse_index: DenseVec::new(),
        }
    }

    fn with_capacity(capacity: usize) -> OneToNForest<T>{
        OneToNForest{
            arena: idtree::Arena::with_capacity(capacity),
            entities_roots: DenseVec::with_capacity(capacity),
            reverse_index: DenseVec::with_capacity(capacity),
        }
    }

    fn insert(&mut self, guid: usize, t: T){
        unsafe{ self.insert_root(guid, t) };
    }

    fn remove(&mut self, guid: usize){
        for id in unsafe{ self.entities_roots.get_unchecked(guid).iter() }{
            self.arena.remove_tree(*id);
        }
        self.entities_roots.remove(guid);
    }

    unsafe fn get(&self, guid: usize) -> RootsIter<'a,T>{
        RootsIter{
            iter: delete_lifetime(self.entities_roots.get_unchecked(guid)).iter(),
            arena: delete_lifetime(&self.arena)
        }
    }

    unsafe fn get_mut(&mut self, guid: usize) -> RootsIterMut<'a,T>{
        RootsIterMut{
            iter: delete_lifetime(self.entities_roots.get_unchecked(guid)).iter(),
            arena: delete_lifetime_mut(&mut self.arena)
        }
    }

    unsafe fn fast_get(&self, guid: usize) -> RootsIter<'a,T>{
        RootsIter{
            iter: delete_lifetime(self.entities_roots.get_fast_unchecked(guid)).iter(),
            arena: delete_lifetime(&self.arena)
        }
    }

    unsafe fn fast_get_mut(&mut self, guid: usize) -> RootsIterMut<'a,T>{
        RootsIterMut{
            iter: delete_lifetime(self.entities_roots.get_fast_unchecked(guid)).iter(),
            arena: delete_lifetime_mut(&mut self.arena)
        }
    }

    fn fast_index(&self, guid: usize) -> usize{
        self.entities_roots.fast_index(guid)
    }

    fn guid_from_fast_index(&self, idx: usize) -> usize{
        unsafe{ self.entities_roots.unchecked_guid_from_fast_index(idx) }
    }

    fn contains(&self, guid: usize) -> bool{
        self.reverse_index.contains(guid)
    }

    fn ided_iter(&self) -> Self::IdedIter {
        unsafe{
            RawIdedIter{
                arena: delete_lifetime(&self.arena),
                iter: mem::transmute(self.entities_roots.iter())
            }
        }
    }

    fn iter(&self) -> Self::Iter {
        unsafe{
            RawIter{
                arena: delete_lifetime(&self.arena),
                iter: mem::transmute(self.entities_roots.values())
            }
        }
    }

    fn iter_mut(&mut self) -> Self::IterMut {
        unsafe{
            RawIterMut{
                arena: delete_lifetime_mut(&mut self.arena),
                iter: mem::transmute(self.entities_roots.values())
            }
        }
    }
}

pub struct RawIdedIter<'a, T: 'a>{
    arena: &'a idtree::Arena<T>,
    iter: densevec::Iter<'a, usize, Vec<idtree::NodeId>>,
}

impl<'a, T: 'a> Iterator for RawIdedIter<'a, T>{
    type Item = (usize, RootsIter<'a, T>);
    fn next(&mut self) -> Option<(usize, RootsIter<'a, T>)>{
        unsafe{
            self.iter.next().map(|(id, roots)| (id, RootsIter{
                iter: roots.iter(),
                arena: delete_lifetime(self.arena),
            }))
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.iter.size_hint()
    }
}

impl<'a, T: 'a> ExactSizeIterator for RawIdedIter<'a, T>{
    #[inline]
    fn len(&self) -> usize {
        self.iter.len()
    }
}

pub struct RawIter<'a, T: 'a>{
    arena: &'a idtree::Arena<T>,
    iter: slice::Iter<'a, Vec<idtree::NodeId>>,
}

impl<'a, T: 'a> Iterator for RawIter<'a, T>{
    type Item = RootsIter<'a, T>;
    fn next(&mut self) -> Option<RootsIter<'a, T>>{
        unsafe{
            self.iter.next().map(|roots| RootsIter{
                iter: roots.iter(),
                arena: delete_lifetime(self.arena),
            })
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.iter.size_hint()
    }
}

impl<'a, T: 'a> ExactSizeIterator for RawIter<'a, T>{
    #[inline]
    fn len(&self) -> usize {
        self.iter.len()
    }
}

pub struct RawIterMut<'a, T: 'a>{
    arena: &'a mut idtree::Arena<T>,
    iter: slice::Iter<'a, Vec<idtree::NodeId>>,
}

impl<'a, T: 'a> Iterator for RawIterMut<'a, T>{
    type Item = RootsIterMut<'a, T>;
    fn next(&mut self) -> Option<RootsIterMut<'a, T>>{
        unsafe{
            self.iter.next().map(|roots| RootsIterMut{
                iter: roots.iter(),
                arena: delete_lifetime_mut(self.arena),
            })
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.iter.size_hint()
    }
}

impl<'a, T: 'a> ExactSizeIterator for RawIterMut<'a, T>{
    #[inline]
    fn len(&self) -> usize {
        self.iter.len()
    }
}

impl<'a, T: 'a + Component> HierarchicalOneToNStorage<'a,T> for OneToNForest<T>{
    unsafe fn insert_root(&mut self, guid: usize, t: T) -> idtree::NodeRefMut<T>{
        let root = self.arena.new_node(t);
        self.entities_roots.entry(guid)
            .or_insert_with(|| vec![])
            .push(root.id());

        //TODO: is it enough to insert in the reverse index here? or also with childs
        self.reverse_index.insert(guid, root.id());
        root
    }

    unsafe fn insert_child(&mut self, parent: idtree::NodeId, t: T) -> idtree::NodeRefMut<T>{
        parent.append_new(t, &mut self.arena)
    }
}


pub struct RootsIter<'a, T: 'a>{
    iter: slice::Iter<'a, idtree::NodeId>,
    arena: &'a idtree::Arena<T>,
}

impl<'a, T: 'a + Component> Debug for RootsIter<'a,T>{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result{
        fmt.write_str("Roots iter for Unknown component")
    }
}

impl<'a, T: 'a> Iterator for RootsIter<'a, T>{
    type Item = idtree::NodeRef<'a,T>;

    fn next(&mut self) -> Option<idtree::NodeRef<'a, T>>{
        self.iter.next().map(|id| self.arena.get(*id))
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.iter.size_hint()
    }
}

impl<'a, T: 'a> ExactSizeIterator for RootsIter<'a, T>{
    #[inline]
    fn len(&self) -> usize {
        self.iter.len()
    }
}

pub struct RootsIterMut<'a, T: 'a>{
    iter: slice::Iter<'a, idtree::NodeId>,
    arena: &'a mut idtree::Arena<T>,
}

impl<'a, T: 'a> Iterator for RootsIterMut<'a, T>{
    type Item = idtree::NodeRefMut<'a,T>;

    fn next(&mut self) -> Option<idtree::NodeRefMut<'a, T>>{
        self.iter.next().map(|id| unsafe{ mem::transmute::<idtree::NodeRefMut<T>, idtree::NodeRefMut<T>>(self.arena.get_mut(*id)) } )
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.iter.size_hint()
    }
}

impl<'a, T: 'a> ExactSizeIterator for RootsIterMut<'a, T>{
    #[inline]
    fn len(&self) -> usize {
        self.iter.len()
    }
}

pub struct Iter<'a, T: 'a>{
    storage: ReadGuardRef<'a, OneToNForest<T>>,
    iter: densevec::Values<'a, Vec<idtree::NodeId>>,
}

impl<'a, T: 'a> Iterator for Iter<'a, T>{
    type Item = RootsIter<'a, T>;
    fn next(&mut self) -> Option<RootsIter<'a, T>>{
        self.iter.next().map(|roots| RootsIter{
            iter: roots.iter(),
            arena: unsafe{ mem::transmute::<&idtree::Arena<T>, &idtree::Arena<T>>(&self.storage.arena) },
        })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.iter.size_hint()
    }
}

impl<'a, T: 'a> ExactSizeIterator for Iter<'a, T>{
    #[inline]
    fn len(&self) -> usize {
        self.iter.len()
    }
}

impl<'a, T> IntoIter for ReadGuardRef<'a, OneToNForest<T>>{
    type Iter = Iter<'a, T>;
    fn into_iter(self) -> Iter<'a, T>{
        Iter{
            iter: unsafe{ mem::transmute::<densevec::Values<Vec<idtree::NodeId>>, densevec::Values<Vec<idtree::NodeId>>>(self.entities_roots.values()) },
            storage: self,
        }
    }
}

pub struct IterMut<'a, T: 'a>{
    storage: WriteGuardRef<'a, OneToNForest<T>>,
    iter: densevec::Values<'a, Vec<idtree::NodeId>>,
}

impl<'a, T: 'a> Iterator for IterMut<'a, T>{
    type Item = RootsIterMut<'a, T>;
    fn next(&mut self) -> Option<RootsIterMut<'a, T>>{
        self.iter.next().map(|roots| RootsIterMut{
            iter: roots.iter(),
            arena: unsafe{ mem::transmute::<&mut idtree::Arena<T>, &mut idtree::Arena<T>>(&mut self.storage.arena) },
        })
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.iter.size_hint()
    }
}

impl<'a, T: 'a> ExactSizeIterator for IterMut<'a, T>{
    #[inline]
    fn len(&self) -> usize {
        self.iter.len()
    }
}

impl<'a, T> IntoIterMut for WriteGuardRef<'a, OneToNForest<T>>{
    type IterMut = IterMut<'a, T>;
    fn into_iter_mut(self) -> IterMut<'a, T>{
        IterMut{
            iter: unsafe{ mem::transmute::<densevec::Values<Vec<idtree::NodeId>>, densevec::Values<Vec<idtree::NodeId>>>(self.entities_roots.values()) },
            storage: self,
        }
    }
}