use std::iter::{repeat_n, Filter, Flatten, Map, RepeatN, Zip};
use std::marker::PhantomData;

use crate::component::storage::{
    Archetype,
    ArchetypeEntity,
    ArchetypeRefIter,
    EntityIter,
    Storage as ComponentStorage,
};
use crate::component::{
    Metadata as ComponentMetadata,
    RefSequence as ComponentRefSequence,
};
use crate::lock::ReadGuard;
use crate::query::options::Options;
use crate::system::{Param as SystemParam, System};
use crate::uid::Uid;
use crate::World;

pub mod options;

#[derive(Debug)]
pub struct Query<'world, Comps, OptionsT = ()>
where
    Comps: ComponentRefSequence,
{
    world: &'world World,
    component_storage: ReadGuard<'world, ComponentStorage>,
    _pd: PhantomData<(Comps, OptionsT)>,
}

impl<'world, Comps, OptionsT> Query<'world, Comps, OptionsT>
where
    Comps: ComponentRefSequence,
    OptionsT: Options,
{
    /// Iterates over the entities matching this query.
    #[must_use]
    pub fn iter<'query>(
        &'query self,
    ) -> ComponentIter<'query, 'world, Comps, QueryEntityIter<'query>>
    {
        #[cfg(feature = "debug")]
        tracing::debug!("Searching for {}", std::any::type_name::<Comps>());

        #[allow(clippy::map_flatten)]
        ComponentIter {
            world: self.world,
            entities: self
                .component_storage
                .iter_archetypes_with_comps(Comps::metadata())
                .map(
                    (|archetype| {
                        repeat_n(archetype, archetype.entity_cnt())
                            .zip(archetype.entities())
                    }) as ComponentIterMapFn,
                )
                .flatten()
                .filter(|(_, entity)| OptionsT::entity_filter(entity.components())),
            comps_pd: PhantomData,
        }
    }

    /// Iterates over the entities matching this query and has the provided extra
    /// component.
    #[must_use]
    pub fn iter_with_extra_comps<'query>(
        &'query self,
        extra_components: impl IntoIterator<Item = ComponentMetadata>,
    ) -> ComponentIter<'query, 'world, Comps, QueryEntityIter<'query>>
    {
        #[cfg(feature = "debug")]
        tracing::debug!(
            "Searching for {} + extra components",
            std::any::type_name::<Comps>()
        );

        #[allow(clippy::map_flatten)]
        ComponentIter {
            world: self.world,
            entities: self
                .component_storage
                .iter_archetypes_with_comps(
                    Comps::metadata()
                        .into_iter()
                        .chain(extra_components)
                        .collect::<Vec<_>>(),
                )
                .map(
                    (|archetype| {
                        repeat_n(archetype, archetype.entity_cnt())
                            .zip(archetype.entities())
                    }) as ComponentIterMapFn,
                )
                .flatten()
                .filter(|(_, entity)| OptionsT::entity_filter(entity.components())),
            comps_pd: PhantomData,
        }
    }

    /// Returns the UID of the entity at the given query iteration index.
    #[must_use]
    pub fn get_entity_uid(&self, entity_index: usize) -> Option<Uid>
    {
        Some(
            self.component_storage
                .iter_archetypes_with_comps(Comps::metadata())
                .flat_map(|archetype| archetype.entities())
                .filter(|entity| OptionsT::entity_filter(entity.components()))
                .nth(entity_index)?
                .uid(),
        )
    }

    pub(crate) fn new(world: &'world World) -> Self
    {
        Self {
            world,
            component_storage: world
                .data
                .component_storage
                .read_nonblock()
                .expect("Failed to acquire read-only component storage lock"),
            _pd: PhantomData,
        }
    }
}

impl<'query, 'world, Comps, OptionsT> IntoIterator
    for &'query Query<'world, Comps, OptionsT>
where
    Comps: ComponentRefSequence + 'world,
    OptionsT: Options,
{
    type IntoIter = ComponentIter<'query, 'world, Comps, QueryEntityIter<'query>>;
    type Item = Comps::Handles<'query>;

    fn into_iter(self) -> Self::IntoIter
    {
        self.iter()
    }
}

impl<'world, Comps, OptionsT> SystemParam<'world> for Query<'world, Comps, OptionsT>
where
    Comps: ComponentRefSequence,
    OptionsT: Options,
{
    type Input = ();

    fn initialize<SystemImpl>(
        _system: &mut impl System<'world, SystemImpl>,
        _input: Self::Input,
    )
    {
    }

    fn new<SystemImpl>(
        _system: &'world impl System<'world, SystemImpl>,
        world: &'world World,
    ) -> Self
    {
        Self::new(world)
    }
}

type ComponentIterMapFn =
    for<'a> fn(&'a Archetype) -> Zip<RepeatN<&'a Archetype>, EntityIter<'a>>;

type ComponentIterFilterFn =
    for<'a, 'b> fn(&'a (&'b Archetype, &'b ArchetypeEntity)) -> bool;

type QueryEntityIter<'query> = Filter<
    Flatten<Map<ArchetypeRefIter<'query>, ComponentIterMapFn>>,
    ComponentIterFilterFn,
>;

pub struct ComponentIter<'query, 'world, Comps, EntityIter>
where
    EntityIter: Iterator<Item = (&'query Archetype, &'query ArchetypeEntity)>,
{
    world: &'world World,
    entities: EntityIter,
    comps_pd: PhantomData<Comps>,
}

impl<'query, 'world, Comps, EntityIter> Iterator
    for ComponentIter<'query, 'world, Comps, EntityIter>
where
    Comps: ComponentRefSequence + 'world,
    EntityIter: Iterator<Item = (&'query Archetype, &'query ArchetypeEntity)>,
    'world: 'query,
{
    type Item = Comps::Handles<'query>;

    fn next(&mut self) -> Option<Self::Item>
    {
        let (archetype, entity) = self.entities.next()?;

        Some(Comps::from_components(
            entity.components(),
            |component_uid| archetype.get_index_for_component(component_uid),
            self.world,
        ))
    }
}