use super::Entity;
use super::component::{Component, ComponentStorage};

pub trait Query<'a, W> {
    type Item;
    fn run(world: &'a W) -> impl Iterator<Item = Self::Item>;
}

pub trait QueryMut<'a, W> {
    type Item;
    fn run_mut(world: &'a mut W) -> impl Iterator<Item = Self::Item>;
}

impl<'a, W, A> Query<'a, W> for A
where
    W: ComponentStorage<A>,
    A: Component + 'a,
{
    type Item = &'a A;
    fn run(world: &'a W) -> impl Iterator<Item = Self::Item> {
        ComponentStorage::<A>::storage(world).iter().map(|(_, a)| a)
    }
}

impl<'a, W, A> QueryMut<'a, W> for A
where
    W: ComponentStorage<A>,
    A: Component + 'a,
{
    type Item = &'a mut A;

    fn run_mut(world: &'a mut W) -> impl Iterator<Item = Self::Item> {
        let a_ptr = <W as ComponentStorage<A>>::storage_ptr(world);

        unsafe {
            let a_ref = &mut *a_ptr;
            a_ref.iter_mut().map(|(_, a_comp)| a_comp)
        }
    }
}

impl<'a, W, A> Query<'a, W> for (A,)
where
    W: ComponentStorage<A>,
    A: Component + 'a,
{
    type Item = &'a A;
    fn run(world: &'a W) -> impl Iterator<Item = Self::Item> {
        ComponentStorage::<A>::storage(world).iter().map(|(_, a)| a)
    }
}

impl<'a, W, A> QueryMut<'a, W> for (A,)
where
    W: ComponentStorage<A>,
    A: Component + 'a,
{
    type Item = &'a mut A;

    fn run_mut(world: &'a mut W) -> impl Iterator<Item = Self::Item> {
        let a_ptr = <W as ComponentStorage<A>>::storage_ptr(world);

        unsafe {
            let a_ref = &mut *a_ptr;
            a_ref.iter_mut().map(|(_, a_comp)| a_comp)
        }
    }
}

impl<'a, W, A> Query<'a, W> for (Entity, A)
where
    W: ComponentStorage<A>,
    A: Component + 'a,
{
    type Item = (Entity, &'a A);
    fn run(world: &'a W) -> impl Iterator<Item = Self::Item> {
        ComponentStorage::<A>::storage(world).iter()
    }
}

macro_rules! impl_query {
    (
        $first:ident
        $(, $rest:ident)+
    ) => {
        #[allow(non_snake_case)]
        impl<'a, W, $first, $($rest),+> Query<'a, W>
            for ($first, $($rest),+)
        where
            $first: Component + 'a,
            $(
                $rest: Component + 'a,
            )+
            W: ComponentStorage<$first>
                $(+ ComponentStorage<$rest>)+,
        {
            type Item = (
                &'a $first,
                $(
                    &'a $rest,
                )+
            );

            fn run(world: &'a W) -> impl Iterator<Item = Self::Item> {
                let first_storage =
                    <W as ComponentStorage<$first>>::storage(world);

                $(
                    let $rest =
                        <W as ComponentStorage<$rest>>::storage(world);
                )+

                first_storage.iter().filter_map(move |(entity, first_comp)| {
                    $(
                        let $rest = $rest.get(entity)?;
                    )+

                    Some((
                        first_comp,
                        $(
                            $rest,
                        )+
                    ))
                })
            }
        }

        #[allow(non_snake_case)]
        impl<'a, W, $first, $($rest),+> Query<'a, W>
            for (Entity, $first, $($rest),+)
        where
            $first: Component + 'a,
            $(
                $rest: Component + 'a,
            )+
            W: ComponentStorage<$first>
                $(+ ComponentStorage<$rest>)+,
        {
            type Item = (
                Entity,
                &'a $first,
                $(
                    &'a $rest,
                )+
            );

            fn run(world: &'a W) -> impl Iterator<Item = Self::Item> {
                let first_storage =
                    <W as ComponentStorage<$first>>::storage(world);

                $(
                    let $rest =
                        <W as ComponentStorage<$rest>>::storage(world);
                )+

                first_storage.iter().filter_map(move |(entity, first_comp)| {
                    $(
                        let $rest = $rest.get(entity)?;
                    )+

                    Some((
                        entity,
                        first_comp,
                        $(
                            $rest,
                        )+
                    ))
                })
            }
        }
    };
}

macro_rules! impl_query_mut {
    (
        $first:ident
        $(, $rest:ident)*
    ) => {
        #[allow(non_snake_case)]
        impl<'a, W, $first, $($rest),*> QueryMut<'a, W>
            for ($first, $($rest),*)
        where
            $first: Component + 'a,
            $(
                $rest: Component + 'a,
            )*
            W: ComponentStorage<$first>
                $(+ ComponentStorage<$rest>)*,
        {
            type Item = (
                &'a mut $first,
                $(
                    &'a mut $rest,
                )*
            );

            fn run_mut(world: &'a mut W) -> impl Iterator<Item = Self::Item> {
                let first_ptr =
                    <W as ComponentStorage<$first>>::storage_ptr(world);

                $(
                    let $rest =
                        <W as ComponentStorage<$rest>>::storage_ptr(world);
                )*

                unsafe {
                    let first_ref = &mut *first_ptr;

                    first_ref.iter_mut().filter_map(move |(entity, first_comp)| {
                        $(
                            let $rest =
                                (*$rest).get_mut(entity)?;
                        )*

                        Some((
                            first_comp,
                            $(
                                $rest,
                            )*
                        ))
                    })
                }
            }
        }

        #[allow(non_snake_case)]
        impl<'a, W, $first, $($rest),*> QueryMut<'a, W>
            for (Entity, $first, $($rest),*)
        where
            $first: Component + 'a,
            $(
                $rest: Component + 'a,
            )*
            W: ComponentStorage<$first>
                $(+ ComponentStorage<$rest>)*,
        {
            type Item = (
                Entity,
                &'a mut $first,
                $(
                    &'a mut $rest,
                )*
            );

            fn run_mut(world: &'a mut W) -> impl Iterator<Item = Self::Item> {
                let first_ptr =
                    <W as ComponentStorage<$first>>::storage_ptr(world);

                $(
                    let $rest =
                        <W as ComponentStorage<$rest>>::storage_ptr(world);
                )*

                unsafe {
                    let first_ref = &mut *first_ptr;

                    first_ref.iter_mut().filter_map(move |(entity, first_comp)| {
                        $(
                            let $rest =
                                (*$rest).get_mut(entity)?;
                        )*

                        Some((
                            entity,
                            first_comp,
                            $(
                                $rest,
                            )*
                        ))
                    })
                }
            }
        }
    };
}

macro_rules! impl_query_mut_opt {
    (
        $first:ident
        $(, $req:ident)*
        =>
        $($opt:ident ?),* $(,)?
    ) => {
        #[allow(non_snake_case)]
        impl<'a, W, $first, $($req,)* $($opt,)*> QueryMut<'a, W>
            for (
                $first,
                $($req,)*
                $(Option<$opt>,)*
            )
        where
            W: ComponentStorage<$first>
                $(+ ComponentStorage<$req>)*
                $(+ ComponentStorage<$opt>)*,
            $first: Component + 'a,
            $(
                $req: Component + 'a,
            )*
            $(
                $opt: Component + 'a,
            )*
        {
            type Item = (
                &'a mut $first,
                $(
                    &'a mut $req,
                )*
                $(
                    Option<&'a mut $opt>,
                )*
            );

            fn run_mut(world: &'a mut W) -> impl Iterator<Item = Self::Item> {
                let first_ptr =
                    <W as ComponentStorage<$first>>::storage_ptr(world);

                $(
                    let $req =
                        <W as ComponentStorage<$req>>::storage_ptr(world);
                )*

                $(
                    let $opt =
                        <W as ComponentStorage<$opt>>::storage_ptr(world);
                )*

                unsafe {
                    let first_ref = &mut *first_ptr;

                    first_ref.iter_mut().filter_map(move |(entity, first_comp)| {
                        $(
                            let $req =
                                (*$req).get_mut(entity)?;
                        )*

                        Some((
                            first_comp,
                            $(
                                $req,
                            )*
                            $(
                                (*$opt).get_mut(entity),
                            )*
                        ))
                    })
                }
            }
        }

        #[allow(non_snake_case)]
        impl<'a, W, $first, $($req,)* $($opt,)*> QueryMut<'a, W>
            for (
                Entity,
                $first,
                $($req,)*
                $(Option<$opt>,)*
            )
        where
            W: ComponentStorage<$first>
                $(+ ComponentStorage<$req>)*
                $(+ ComponentStorage<$opt>)*,
            $first: Component + 'a,
            $(
                $req: Component + 'a,
            )*
            $(
                $opt: Component + 'a,
            )*
        {
            type Item = (
                Entity,
                &'a mut $first,
                $(
                    &'a mut $req,
                )*
                $(
                    Option<&'a mut $opt>,
                )*
            );

            fn run_mut(world: &'a mut W) -> impl Iterator<Item = Self::Item> {
                let first_ptr =
                    <W as ComponentStorage<$first>>::storage_ptr(world);

                $(
                    let $req =
                        <W as ComponentStorage<$req>>::storage_ptr(world);
                )*

                $(
                    let $opt =
                        <W as ComponentStorage<$opt>>::storage_ptr(world);
                )*

                unsafe {
                    let first_ref = &mut *first_ptr;

                    first_ref.iter_mut().filter_map(move |(entity, first_comp)| {
                        $(
                            let $req =
                                (*$req).get_mut(entity)?;
                        )*

                        Some((
                            entity,
                            first_comp,
                            $(
                                $req,
                            )*
                            $(
                                (*$opt).get_mut(entity),
                            )*
                        ))
                    })
                }
            }
        }
    };
}

impl_query!(A, B);
impl_query!(A, B, C);
impl_query!(A, B, C, D);
impl_query!(A, B, C, D, E);
impl_query!(A, B, C, D, E, F);

impl_query_mut!(A, B);
impl_query_mut!(A, B, C);
impl_query_mut!(A, B, C, D);
impl_query_mut!(A, B, C, D, E);
impl_query_mut!(A, B, C, D, E, F);

impl_query_mut_opt!(A => B?);
impl_query_mut_opt!(A, B => C?);
impl_query_mut_opt!(A, B, C => D?);
impl_query_mut_opt!(A, B, C, D => E?);
impl_query_mut_opt!(A, B, C, D, E => F?);
impl_query_mut_opt!(A => B?, C?);
impl_query_mut_opt!(A, B => C?, D?);
impl_query_mut_opt!(A, B, C => D?, E?);
impl_query_mut_opt!(A, B, C, D => E?, F?);
impl_query_mut_opt!(A => B?, C?, D?);
impl_query_mut_opt!(A, B => C?, D?, E?);
impl_query_mut_opt!(A, B, C => D?, E?, F?);