use std::mem;
use crate::idtree;
use densevec::DenseVec;
use crate::sync::{ReadGuardRef, WriteGuardRef, Lock};
use crate::storage::{
    Storage, IntoIter, IntoIterMut, HierarchicalStorage,
    IntoOrderedIter, IntoOrderedIterMut, IntoIdedIterMut, IntoOrderedIdedIterMut,
    OrderedId,
};
use crate::utils::*;
use std::slice;

pub struct Forest<T>{
    arena: idtree::Arena<T>,
    roots: Vec<idtree::NodeId>,
    index: DenseVec<idtree::NodeId>,
    reverse_index: DenseVec<usize>,
    ordered_ids: Lock<Vec<OrderedId<idtree::NodeId>>>,
}

// Is this actually Send + Sync?
// Only needed for Mutcell in Lock which should be safe?
unsafe impl<T: Send> Send for Forest<T>{}
unsafe impl<T: Sync> Sync for Forest<T>{}

impl<'a, T: 'a> Storage<'a, T> for Forest<T>{
    type Get = &'a T;
    type GetMut = &'a mut T;
    type DerefTarget = T;
    type IdedIter = idtree::AllIdedData<'a,T>;
    type Iter = idtree::AllData<'a, T>;
    type IterMut = idtree::AllDataMut<'a, T>;

    fn new() -> Forest<T>{
        Forest{
            arena: idtree::Arena::new(),
            roots: Vec::new(),
            index: DenseVec::new(),
            reverse_index: DenseVec::new(),
            ordered_ids: Lock::new(vec![]),
        }
    }

    fn with_capacity(capacity: usize) -> Forest<T>{
        Forest{
            arena: idtree::Arena::with_capacity(capacity),
            roots: Vec::with_capacity(capacity),
            index: DenseVec::with_capacity(capacity),
            reverse_index: DenseVec::with_capacity(capacity),
            ordered_ids: Lock::new(Vec::with_capacity(capacity)),
        }
    }

    fn insert(&mut self, guid: usize, t: T){
        let node_id = self.arena.new_node(t);
        self.index.insert(guid, node_id.id());
        self.reverse_index.insert(node_id.id().id(), guid);
        self.roots.push(node_id.id());
        self.ordered_ids.get_mut().clear();
    }

    fn remove(&mut self, guid: usize){
        {
            let id = unsafe{ self.index.get_unchecked(guid) };
            self.arena.remove(*id);
            if let Some(pos) = self.roots.iter().position(|i| i == id){
                self.roots.remove(pos);
            }
            self.reverse_index.remove(id.id());
        }
        self.index.remove(guid);
        self.ordered_ids.get_mut().clear();
    }

    unsafe fn get(&self, guid: usize) -> &'a T{
        let node_id = self.index.get_unchecked(guid);
        delete_lifetime(&self.arena[*node_id])
    }

    unsafe fn get_mut(&mut self, guid: usize) -> &'a mut T{
        let node_id = self.index.get_unchecked(guid);
        delete_lifetime_mut(&mut self.arena[*node_id])
    }

    unsafe fn fast_get(&self, guid: usize) -> &'a T{
        let node_id = self.index.get_fast_unchecked(guid);
        delete_lifetime(&self.arena[*node_id])
    }

    unsafe fn fast_get_mut(&mut self, guid: usize) -> &'a mut T{
        let node_id = self.index.get_fast_unchecked(guid);
        delete_lifetime_mut(&mut self.arena[*node_id])
    }

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

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

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

    fn ided_iter(&self) -> Self::IdedIter{
        unsafe{ mem::transmute(self.arena.all_ided_data()) }
    }

    fn iter(&self) -> Self::Iter{
        unsafe{ mem::transmute(self.arena.all_data()) }
    }

    fn iter_mut(&mut self) -> Self::IterMut{
        unsafe{ mem::transmute(self.arena.all_data_mut()) }
    }
}

pub struct Iter<'a, T: 'a>{
    _guard: ReadGuardRef<'a, Forest<T>>,
    it: idtree::AllNodes<'a, T>,
}

impl<'a, T: 'a> Iterator for Iter<'a, T>{
    type Item = &'a T;
    #[inline]
    fn next(&mut self) -> Option<&'a T>{
        self.it.next().map(|node| &node.data)
    }

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

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

impl<'a, T> IntoIter for ReadGuardRef<'a, Forest<T>>{
    type Iter = Iter<'a, T>;
    fn into_iter(self) -> Iter<'a, T>{
        Iter{
            it: unsafe{mem::transmute::<idtree::AllNodes<T>, idtree::AllNodes<T>>(self.arena.all_nodes())},
            _guard: self,
        }
    }
}

pub struct IterMut<'a, T: 'a>{
    _guard: WriteGuardRef<'a, Forest<T>>,
    it: idtree::AllNodesMut<'a, T>
}

impl<'a, T: 'a> Iterator for IterMut<'a, T>{
    type Item = &'a mut T;
    #[inline]
    fn next(&mut self) -> Option<&'a mut T>{
        self.it.next().map(|node| &mut node.data)
    }

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

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

impl<'a, T> IntoIterMut for WriteGuardRef<'a, Forest<T>>{
    type IterMut = IterMut<'a, T>;
    fn into_iter_mut(mut self) -> IterMut<'a, T>{
        IterMut{
            it: unsafe{mem::transmute::<idtree::AllNodesMut<T>, idtree::AllNodesMut<T>>(self.arena.all_nodes_mut())},
            _guard: self,
        }
    }
}

pub struct IdedIterMut<'a, T: 'a>{
    _guard: WriteGuardRef<'a, Forest<T>>,
    iter: densevec::IterMut<'a, usize, idtree::Node<T>>,
}

impl<'a, T: 'a> Iterator for IdedIterMut<'a, T>{
    type Item = (usize, &'a mut T);
    #[inline]
    fn next(&mut self) -> Option<(usize, &'a mut T)>{
        self.iter.next().map(|(id, node)| (id, &mut node.data))
    }

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

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

impl<'a, T> IntoIdedIterMut for WriteGuardRef<'a, Forest<T>>{
    type IdedIterMut = IdedIterMut<'a, T>;
    fn into_ided_iter_mut(mut self) -> IdedIterMut<'a, T>{
        IdedIterMut{
            iter: unsafe{mem::transmute(self.arena.all_nodes_ided_mut())},
            _guard: self,
        }
    }
}

impl<'a,T: 'a> HierarchicalStorage<'a,T> for Forest<T>{
    unsafe fn insert_child(&mut self, parent_guid: usize, guid: usize, value: T){
        let parent_id = *self.index.get_unchecked(parent_guid);
        let node_id = self.arena.get_mut(parent_id).append_new(value);
        self.index.insert(guid, node_id.id());
        self.reverse_index.insert(node_id.id().id(), guid);
        self.ordered_ids.get_mut().clear();
    }

    unsafe fn get_node(&self, guid: usize) -> idtree::NodeRef<T>{
        let node_id = self.index.get_unchecked(guid);
        self.arena.get(*node_id)
    }

    unsafe fn get_node_mut(&mut self, guid: usize) -> idtree::NodeRefMut<T>{
        let node_id = self.index.get_unchecked(guid);
        self.arena.get_mut(*node_id)
    }

    unsafe fn fast_get_node(&self, idx: idtree::NodeId) -> idtree::NodeRef<T>{
        self.arena.get(idx)
    }

    unsafe fn fast_get_node_mut(&mut self, idx: idtree::NodeId) -> idtree::NodeRefMut<T>{
        self.arena.get_mut(idx)
    }

    //TODO: this is terribly unoptimized
    fn ordered_ids(&self) -> ReadGuardRef<Vec<OrderedId<idtree::NodeId>>>{
        if self.ordered_ids.read().is_empty() {
            let iter_tree = ForestHierarchicalIdsIter{
                forest: self,
                current: if self.roots.is_empty(){ None } else { Some(0) },
                iter: if self.roots.is_empty(){ None } else { Some(self.roots[0].descendants(&self.arena)) }
            };
            self.ordered_ids
                .write()
                .extend(iter_tree
                    .map(|id| OrderedId{
                        guid: unsafe{ *self.reverse_index.get_unchecked(id.id()) },
                        fast: id,
                    })
                );
        }
        ReadGuardRef::new(self.ordered_ids.read())
    }

    unsafe fn guid_for(&self, nodeid: idtree::NodeId) -> usize{
        *self.reverse_index.get_unchecked(nodeid.id())
    }

    unsafe fn ordered_fast_index(&self, guid: usize) -> idtree::NodeId{
        *self.index.get_unchecked(guid)
    }
}

impl<'a, T> IntoOrderedIter for ReadGuardRef<'a, Forest<T>>{
    type OrderedIter = ForestHierarchicalIter<'a,T>;
    fn into_ordered_iter(self) -> Self::OrderedIter{
        // We need this cause there's no way to split a Send Lock, should be safe
        // since we are holding a read reference to the storage and we fill the ids
        // here so no one will modify it while we iterate and no one will need
        // to fill the ordered ids cache either
        let ordered_ids = {
            let ordered_ids = self.ordered_ids();
            let ptr = ordered_ids.as_ptr();
            let len = ordered_ids.len();
            unsafe{ slice::from_raw_parts(ptr, len) }
        };
        ForestHierarchicalIter{
            ordered_ids,
            forest: self,
            next: 0,
        }
    }
}

struct ForestHierarchicalIdsIter<'a, T: 'a>{
    forest: &'a Forest<T>,
    current: Option<usize>,
    iter: Option<idtree::Descendants<'a, T>>,
}

impl<'a, T: 'a> Iterator for ForestHierarchicalIdsIter<'a, T>{
    type Item = idtree::NodeId;
    fn next(&mut self) -> Option<idtree::NodeId>{
        let iter = if let Some(ref mut iter) = self.iter{
            let next = iter.next();
            if let Some(next) = next{
                return Some(self.forest.arena.get(next).id());
            }else{
                if let Some(current) = self.current{
                    if current + 1 < self.forest.roots.len(){
                        self.current = Some(current + 1);
                        Some(self.forest.roots[current + 1].descendants(&self.forest.arena))
                    }else{
                        self.current = None;
                        None
                    }
                }else{
                    None
                }
            }
        }else{
            None
        };

        if let Some(mut iter) = iter{
            let next = iter.next().map(|id| self.forest.arena.get(id).id() );
            self.iter = Some(iter);
            next
        }else{
            None
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>){
        (self.forest.arena.len(), Some(self.forest.arena.len()))
    }
}

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

pub struct ForestHierarchicalIter<'a, T: 'a>{
    forest: ReadGuardRef<'a, Forest<T>>,
    ordered_ids: &'a [OrderedId<idtree::NodeId>],
    next: usize,
}

impl<'a, T: 'a> Iterator for ForestHierarchicalIter<'a, T>{
    type Item = idtree::NodeRef<'a,T>;
    fn next(&mut self) -> Option<idtree::NodeRef<'a,T>>{
        if self.next == self.ordered_ids.len() {
            None
        }else{
            let next = unsafe{ *self.ordered_ids.get_unchecked(self.next) };
            let node = unsafe{ self.forest.fast_get_node(next.fast) };
            self.next += 1;
            let node = unsafe{ mem::transmute::<idtree::NodeRef<T>, idtree::NodeRef<T>>(node) };
            Some(node)
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>){
        (self.forest.arena.len(), Some(self.forest.arena.len()))
    }
}

impl<'a, T> IntoOrderedIterMut for WriteGuardRef<'a, Forest<T>>{
    type OrderedIterMut = ForestHierarchicalIterMut<'a,T>;
    fn into_ordered_iter_mut(self) -> Self::OrderedIterMut{
        self.ordered_ids();
        ForestHierarchicalIterMut{
            next: 0,
            forest: self,
        }
    }
}

pub struct ForestHierarchicalIterMut<'a, T: 'a>{
    forest: WriteGuardRef<'a, Forest<T>>,
    next: usize,
}

impl<'a, T: 'a> Iterator for ForestHierarchicalIterMut<'a, T>{
    type Item = idtree::NodeRefMut<'a,T>;
    fn next(&mut self) -> Option<idtree::NodeRefMut<'a,T>>{
        if self.next == self.forest.ordered_ids.get_mut().len() {
            None
        }else{
            let next = unsafe{ *self.forest.ordered_ids.get_mut().get_unchecked(self.next) };
            let node = unsafe{ self.forest.fast_get_node_mut(next.fast) };
            self.next += 1;
            let node = unsafe{ mem::transmute::<idtree::NodeRefMut<T>, idtree::NodeRefMut<T>>(node) };
            Some(node)
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>){
        (self.forest.arena.len(), Some(self.forest.arena.len()))
    }
}

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

pub struct ForestHierarchicalIdedIterMut<'a, T: 'a>{
    forest: WriteGuardRef<'a, Forest<T>>,
    next: usize,
}

impl<'a, T: 'a> Iterator for ForestHierarchicalIdedIterMut<'a, T>{
    type Item = (usize, idtree::NodeRefMut<'a,T>);
    fn next(&mut self) -> Option<(usize, idtree::NodeRefMut<'a,T>)>{
        if self.next == self.forest.ordered_ids.get_mut().len() {
            None
        }else{
            let next = unsafe{ *self.forest.ordered_ids.get_mut().get_unchecked(self.next) };
            let node = unsafe{ self.forest.fast_get_node_mut(next.fast) };
            self.next += 1;
            let node = unsafe{ mem::transmute::<idtree::NodeRefMut<T>, idtree::NodeRefMut<T>>(node) };
            Some((next.guid, node))
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>){
        (self.forest.arena.len(), Some(self.forest.arena.len()))
    }
}

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

impl<'a, T> IntoOrderedIdedIterMut for WriteGuardRef<'a, Forest<T>>{
    type OrderedIdedIterMut = ForestHierarchicalIdedIterMut<'a, T>;
    fn into_ordered_ided_iter_mut(self) -> Self::OrderedIdedIterMut{
        self.ordered_ids();
        ForestHierarchicalIdedIterMut{
            forest: self,
            next: 0,
        }
    }
}