//! Low-level querying.
use std::iter::{repeat_n, FlatMap, RepeatN, Zip};

use crate::component::storage::archetype::{Archetype, ArchetypeEntity, EntityIter};
use crate::component::storage::{
    ArchetypeRefIter,
    ArchetypeSearchTerms,
    Storage as ComponentStorage,
};
use crate::lock::ReadGuard;
use crate::query::Terms;
use crate::uid::Uid;
use crate::{EntityComponent, World};

/// Low-level entity query structure.
#[derive(Debug)]
pub struct Query<'world, 'terms>
{
    component_storage: ReadGuard<'world, ComponentStorage>,
    terms: Terms<'terms>,
}

impl<'world, 'terms> Query<'world, 'terms>
{
    /// Iterates over the entities matching this query.
    #[must_use]
    pub fn iter(&self) -> Iter<'_>
    {
        Iter {
            iter: self
                .component_storage
                .search_archetypes(ArchetypeSearchTerms {
                    required_components: &self.terms.required_components,
                    excluded_components: &self.terms.excluded_components,
                })
                .flat_map(
                    (|archetype| {
                        repeat_n(archetype, archetype.entity_cnt())
                            .zip(archetype.entities())
                    }) as ComponentIterMapFn,
                ),
        }
    }

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

pub struct Iter<'query>
{
    iter: QueryEntityIter<'query>,
}

impl<'query> Iterator for Iter<'query>
{
    type Item = EntityHandle<'query>;

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

        Some(EntityHandle { archetype, entity })
    }
}

pub struct EntityHandle<'query>
{
    archetype: &'query Archetype,
    entity: &'query ArchetypeEntity,
}

impl<'query> EntityHandle<'query>
{
    /// Returns the [`Uid`] of this entity.
    #[inline]
    #[must_use]
    pub fn uid(&self) -> Uid
    {
        self.entity.uid
    }

    #[inline]
    #[must_use]
    pub fn get_component(&self, component_uid: Uid) -> Option<&'query EntityComponent>
    {
        let index = self.archetype.get_index_for_component(component_uid)?;

        Some(self.entity.components.get(index).unwrap())
    }
}

type ComponentIterMapFnOutput<'a> = Zip<RepeatN<&'a Archetype>, EntityIter<'a>>;

type ComponentIterMapFn = for<'a> fn(&'a Archetype) -> ComponentIterMapFnOutput<'a>;

type QueryEntityIter<'query> = FlatMap<
    ArchetypeRefIter<'query, 'query>,
    ComponentIterMapFnOutput<'query>,
    ComponentIterMapFn,
>;