use crate::{IntoSendStorage, entity::EntitiesStorage, storage::FastIndexExt, sync::{StorageWriteGuard, OrderedIndexGuard}};
use std::{marker, ops::DerefMut};
use crate::Component;
use crate::ComponentSend;
use crate::Bitmask;
use super::{IterOptionWrapper, OrderedData, SafeIter};
use crate::storage::{Storage, StorageRef, HierarchicalStorage, IntoOrderedIterMut};
use crate::idtree;

pub struct WriteHierarchical<'a, T: Component>{
    _marker: marker::PhantomData<&'a T>,
}

pub struct HierarchicalStorageWrite<'a, S, BS, T>{
    storage: BS,
    marker_t: marker::PhantomData<&'a T>,
    marker_s: marker::PhantomData<S>,
}

impl<'a, BS, T> FastIndexExt for HierarchicalStorageWrite<'a, <T as Component>::Storage, BS, T>
where
    <T as Component>::Storage: HierarchicalStorage<'a,T> + Storage<'a,T> + 'a,
    BS: DerefMut<Target = <T as Component>::Storage>,
    T: Component
{
    type FastIndex = idtree::NodeId;
    type StaticTypeId = T;

    fn fast_index(&self, guid: usize) -> idtree::NodeId {
        unsafe{ self.storage.ordered_fast_index(guid) }
    }
}

impl<'a, 'r, BS, T> StorageRef<'r> for HierarchicalStorageWrite<'a, <T as Component>::Storage, BS, T>
where
    <T as Component>::Storage: HierarchicalStorage<'a,T> + Storage<'a,T> + 'a,
    BS: DerefMut<Target = <T as Component>::Storage>,
    T: Component
{
    type Data = WriteHierarchical<'a, T>;
    type Component = idtree::NodeRefMut<'r, T>;

    unsafe fn get_fast_unchecked(&'r mut self, idx: idtree::NodeId) -> idtree::NodeRefMut<'r, T>{
        self.storage.fast_get_node_mut(idx)
    }

    unsafe fn get_unchecked(&'r mut self, guid: usize) -> idtree::NodeRefMut<'r, T>{
        self.storage.get_node_mut(guid)
    }

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

impl<'a, 'r, T> IntoSendStorage<'r> for HierarchicalStorageWrite<'a, <T as Component>::Storage, StorageWriteGuard<'a, T>, T>
where
    <T as Component>::Storage: HierarchicalStorage<'a,T> + Storage<'a,T> + 'a,
    T: Component
{
    type SendStorage = HierarchicalStorageWrite<'a, <T as Component>::Storage, &'r mut <T as Component>::Storage, T>;
    fn into_send_storage(&'r mut self) -> Self::SendStorage {
        HierarchicalStorageWrite {
            storage: &mut self.storage,
            marker_s: marker::PhantomData,
            marker_t: marker::PhantomData,
        }
    }
}

unsafe impl<T: Component> SafeIter<'_> for WriteHierarchical<'_, T>{}

impl<'a, T: ComponentSend> OrderedData<'a> for WriteHierarchical<'a,T>
    where for<'b> <T as Component>::Storage: HierarchicalStorage<'b,T>,
          for<'b> StorageWriteGuard<'b, T>: IntoOrderedIterMut
{
    type Iter = IterOptionWrapper<<StorageWriteGuard<'a, T> as IntoOrderedIterMut>::OrderedIterMut>;
    type Components = T;
    type ComponentsRef = idtree::NodeRefMut<'a, T>;
    type Storage = HierarchicalStorageWrite<'a, <T as Component>::Storage, StorageWriteGuard<'a, T>, Self::Components>;

    fn query_mask<E: EntitiesStorage>(entities: &E) -> Bitmask{
        Bitmask::has(entities.component_mask::<T>())
    }

    fn into_iter<E: EntitiesStorage>(entities: &'a E) -> Self::Iter{
        IterOptionWrapper{
            it: entities.storage_mut::<T>().map(|s| s.into_ordered_iter_mut())
        }
    }

    fn storage<E: EntitiesStorage>(entities: &'a E) -> Option<Self::Storage>{
        Some(HierarchicalStorageWrite{
            storage: entities.storage_mut::<T>()?,
            marker_t: marker::PhantomData,
            marker_s: marker::PhantomData,
        })
    }

    fn ordered_ids<S: FastIndexExt, E: EntitiesStorage>(entities: &'a E, query_mask: Bitmask, unordered_storage: &S) -> OrderedIndexGuard<'a>{
        entities.ordered_entities_for::<T,S>(query_mask, unordered_storage)
    }
}