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

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::query::ComponentIter;
use crate::uid::Uid;
use crate::util::Sortable;
use crate::{EntityComponent, World};

/// Low-level entity query structure.
#[derive(Debug)]
pub struct Query<'world, CompMetadata>
where
    CompMetadata: Sortable<Item = ComponentMetadata> + AsRef<[ComponentMetadata]>,
{
    component_storage: ReadGuard<'world, ComponentStorage>,
    comp_metadata: CompMetadata,
}

impl<'world, CompMetadata> Query<'world, CompMetadata>
where
    CompMetadata: Sortable<Item = ComponentMetadata> + AsRef<[ComponentMetadata]>,
{
    /// Iterates over the entities matching this query.
    #[must_use]
    pub fn iter<'query, OptionsT: Options>(&'query self) -> Iter<'query>
    {
        Iter {
            iter: self
                .component_storage
                .iter_archetypes_with_comps(&self.comp_metadata)
                .map(
                    (|archetype| {
                        repeat_n(archetype, archetype.entity_cnt())
                            .zip(archetype.entities())
                    }) as ComponentIterMapFn,
                )
                .flatten()
                .filter(|(_, entity)| OptionsT::entity_filter(entity.components())),
        }
    }

    pub(crate) fn new(world: &'world World, mut comp_metadata: CompMetadata) -> Self
    {
        comp_metadata.sort_by_key_b(|metadata| metadata.id);

        Self {
            component_storage: world
                .data
                .component_storage
                .read_nonblock()
                .expect("Failed to acquire read-only component storage lock"),
            comp_metadata,
        }
    }
}

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

impl<'query> Iter<'query>
{
    /// Converts this iterator into a [`ComponentIter`].
    ///
    /// Note: The matching entities of this iterator should have all of the non-[`Option`]
    /// components in `Comps`, otherwise iterating the [`ComponentIter`] will cause a
    /// panic.
    #[must_use]
    #[inline]
    pub fn into_component_iter<'world, Comps>(
        self,
        world: &'world World,
    ) -> ComponentIter<'query, 'world, Comps, Self>
    where
        Comps: ComponentRefSequence + 'world,
        'world: 'query,
    {
        ComponentIter::new(world, self)
    }
}

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]
    pub fn uid(&self) -> Uid
    {
        self.entity.uid()
    }

    #[inline]
    pub fn components(&self) -> &'query [EntityComponent]
    {
        self.entity.components()
    }

    #[inline]
    pub fn get_component_index(&self, component_uid: Uid) -> Option<usize>
    {
        self.archetype.get_index_for_component(component_uid)
    }
}

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,
>;