use qecs_core::{
    Component, Id, Valid, StoreBase, ComponentStore, ComponentStoreIter,
    _ComponentStore, _ComponentStoreIter,
};

use std::collections::btree_map::{self, BTreeMap};
use std::iter;

pub struct BTreeMapStore<ID, C> {
    storage: BTreeMap<ID, C>,
}

impl<ID, C> From<BTreeMap<ID, C>> for BTreeMapStore<ID, C>
    where ID: Id, C: Component
{
    fn from(storage: BTreeMap<ID, C>) -> Self {
        BTreeMapStore{ storage: storage }
    }
}

impl<ID, C> Default for BTreeMapStore<ID, C>
    where ID: Id, C: Component
{
    fn default() -> Self {
        Self::from(BTreeMap::default())
    }
}

impl<ID, C> StoreBase for BTreeMapStore<ID, C>
    where ID: Id, C: Component
{
    type Id = ID;

    fn len(&self) -> usize { self.storage.len() }

    fn is_assigned<'id>(&self, id: Valid<'id, Self::Id>) -> bool {
        self.storage.contains_key(&id)
    }

    fn assign_cloned<'id>(
        &mut self, 
        dest: Valid<'id, Self::Id>, 
        src: Valid<'id, Self::Id>,
    ) -> bool {
        match self.get(src).map(|c| c.clone()) {
            Some(com) => { self.assign(dest, com); true },
            None => false,
        }
    }

    fn release_drop<'id>(&mut self, id: Valid<'id, Self::Id>) -> bool {
        self.release(id).is_some()
    }

    fn clear(&mut self){ self.storage.clear() }
}

pub type Components<'a, ID, C> = btree_map::Values<'a, ID, C>;

pub type ComponentsMut<'a, ID, C>
= iter::Map<btree_map::IterMut<'a, ID, C>, fn((&'a ID, &'a mut C)) -> &'a mut C>;

pub type Iter<'a, ID, C>
= iter::Map<btree_map::Iter<'a, ID, C>, fn((&'a ID, &'a C)) -> (Valid<'a, ID>, &'a C)>;

pub type IterMut<'a, ID, C>
= iter::Map<btree_map::IterMut<'a, ID, C>, fn((&'a ID, &'a mut C)) -> (Valid<'a, ID>, &'a mut C)>;

impl<'a, ID, C> _ComponentStore<'a> for BTreeMapStore<ID, C>
    where ID: Id, C: Component
{
    type _Component = C;
    
    type Components = Components<'a, ID, C>;
    
    type ComponentsMut = ComponentsMut<'a, ID, C>;
}

impl<ID, C> ComponentStore for BTreeMapStore<ID, C>
    where ID: Id, C: Component
{
    type Component = C;

    fn assign<'id, T>(&mut self, id: Valid<'id, Self::Id>, com: T) -> Option<Self::Component>
        where T: Into<Self::Component>
    {
        self.storage.insert(*id, com.into())
    }

    fn release<'id>(&mut self, id: Valid<'id, Self::Id>) -> Option<Self::Component> {
        self.storage.remove(&id)
    }

    fn get<'id>(&self, id: Valid<'id, Self::Id>) -> Option<&Self::Component> {
        self.storage.get(&id)
    }

    fn get_mut<'id>(&mut self, id: Valid<'id, Self::Id>) -> Option<&mut Self::Component> {
        self.storage.get_mut(&id)
    }

    fn components<'a>(&'a self) -> <Self as _ComponentStore<'a>>::Components {
        self.storage.values()
    }

    fn components_mut<'a>(&'a mut self) -> <Self as _ComponentStore<'a>>::ComponentsMut {
        fn map_fn<'a, ID, C>((_, com): (&'a ID, &'a mut C)) -> &'a mut C {
            com
        }

        self.storage.iter_mut().map(map_fn)
    }
}

impl<'a, ID, C> _ComponentStoreIter<'a> for BTreeMapStore<ID, C>
    where ID: Id, C: Component
{
    type Iter = Iter<'a, ID, C>;

    type IterMut = IterMut<'a, ID, C>;
}

impl<ID, C> ComponentStoreIter for BTreeMapStore<ID, C>
    where ID: Id, C: Component
{
    fn iter<'a>(&'a self) -> <Self as _ComponentStoreIter<'a>>::Iter {
        fn map_fn<'a, ID: Id, C>((id, com): (&'a ID, &'a C)) -> (Valid<'a, ID>, &'a C) {
            (unsafe { Valid::new(*id) }, com)
        }

        self.storage.iter().map(map_fn)
    }

    fn iter_mut<'a>(&'a mut self) -> <Self as _ComponentStoreIter<'a>>::IterMut {
        fn map_fn<'a, ID: Id, C>((id, com): (&'a ID, &'a mut C)) -> (Valid<'a, ID>, &'a mut C) {
            (unsafe { Valid::new(*id) }, com)
        }

        self.storage.iter_mut().map(map_fn)
    }
}

impl<'id, ID, C, T> Extend<(Valid<'id, ID>, T)> for BTreeMapStore<ID, C> 
    where ID: Id, C: Component, T: Into<C>
{
    fn extend<I>(&mut self, iterable: I)
        where I: IntoIterator<Item = (Valid<'id, ID>, T)>
    {
        for (id, com) in iterable {
            self.assign(id, com.into());
        }
    }
}

impl<'a, ID, C> IntoIterator for &'a BTreeMapStore<ID, C>
    where ID: Id, C: Component
{
    type Item = <Self::IntoIter as Iterator>::Item;
    type IntoIter = Iter<'a, ID, C>;
    fn into_iter(self) -> Self::IntoIter { self.iter() }
}

impl<'a, ID, C> IntoIterator for &'a mut BTreeMapStore<ID, C>
    where ID: Id, C: Component
{
    type Item = <Self::IntoIter as Iterator>::Item;
    type IntoIter = IterMut<'a, ID, C>;
    fn into_iter(self) -> Self::IntoIter { self.iter_mut() }
}

#[test]
fn test(){
    use qecs_core::{Component, Id, Valid};

    // id-type

    #[derive(Debug, Hash, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
    struct MyId(u8);

    impl Component for MyId {}
    impl Id for MyId {}

    fn validate(id: MyId) -> Valid<'static, MyId> {
        unsafe { Valid::new(id) }
    }

    // component-type

    #[derive(Debug, Clone, PartialEq, Eq)]
    struct MyComponent(String);
    
    impl Component for MyComponent {}

    // store-type

    qecs_newtype!(type MyStore: BTreeMapStore<MyId, MyComponent>);

    // code

    let mut mystore: MyStore = MyStore::default();

    let entity_a = validate(MyId(0));
    let entity_b = validate(MyId(1));

    let com_a = MyComponent(String::from("foo"));
    let com_b = MyComponent(String::from("bar"));

    assert!(mystore.assign(entity_a, com_a.clone()).is_none());
    assert_eq!(&mystore.assign(entity_a, com_b).unwrap(), &com_a);
    
    assert!(mystore.assign_cloned(entity_b, entity_a));
    assert_eq!(mystore.get(entity_a).unwrap(), mystore.get(entity_b).unwrap());

    assert_eq!(mystore.iter().count(), 2);
    
    mystore.clear();

    assert!(!mystore.is_assigned(entity_a));
    assert!(!mystore.is_assigned(entity_b));
}
