//! Utilities for working with generics.
//!
//! The layout for generics as expected by chalk are as follows:
//! - Parent parameters
//! - Optional Self parameter
//! - Lifetime parameters
//! - Type or Const parameters
//!
//! where parent follows the same scheme.
use std::ops;

use hir_def::{
    ConstParamId, GenericDefId, GenericParamId, ItemContainerId, LifetimeParamId, Lookup,
    TypeOrConstParamId, TypeParamId,
    db::DefDatabase,
    expr_store::ExpressionStore,
    hir::generics::{
        GenericParamDataRef, GenericParams, LifetimeParamData, LocalLifetimeParamId,
        LocalTypeOrConstParamId, TypeOrConstParamData, TypeParamProvenance, WherePredicate,
    },
};
use itertools::chain;
use triomphe::Arc;

pub fn generics(db: &dyn DefDatabase, def: GenericDefId) -> Generics {
    let parent_generics = parent_generic_def(db, def).map(|def| Box::new(generics(db, def)));
    let (params, store) = db.generic_params_and_store(def);
    let has_trait_self_param = params.trait_self_param().is_some();
    Generics { def, params, parent_generics, has_trait_self_param, store }
}
#[derive(Clone, Debug)]
pub struct Generics {
    def: GenericDefId,
    params: Arc<GenericParams>,
    store: Arc<ExpressionStore>,
    parent_generics: Option<Box<Generics>>,
    has_trait_self_param: bool,
}

impl<T> ops::Index<T> for Generics
where
    GenericParams: ops::Index<T>,
{
    type Output = <GenericParams as ops::Index<T>>::Output;
    fn index(&self, index: T) -> &Self::Output {
        &self.params[index]
    }
}

impl Generics {
    pub(crate) fn def(&self) -> GenericDefId {
        self.def
    }

    pub(crate) fn store(&self) -> &ExpressionStore {
        &self.store
    }

    pub(crate) fn where_predicates(&self) -> impl Iterator<Item = &WherePredicate> {
        self.params.where_predicates().iter()
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.params.is_empty() && self.parent_generics.as_ref().is_none_or(|g| g.params.is_empty())
    }

    pub(crate) fn iter_id(&self) -> impl Iterator<Item = GenericParamId> + '_ {
        self.iter_parent_id().chain(self.iter_self_id())
    }

    pub(crate) fn iter_self_id(&self) -> impl Iterator<Item = GenericParamId> + '_ {
        self.iter_self().map(|(id, _)| id)
    }

    pub(crate) fn iter_parent_id(&self) -> impl Iterator<Item = GenericParamId> + '_ {
        self.iter_parent().map(|(id, _)| id)
    }

    pub(crate) fn iter_self_type_or_consts(
        &self,
    ) -> impl DoubleEndedIterator<Item = (LocalTypeOrConstParamId, &TypeOrConstParamData)> + '_
    {
        let mut toc = self.params.iter_type_or_consts();
        let trait_self_param = self.has_trait_self_param.then(|| toc.next()).flatten();
        chain!(trait_self_param, toc)
    }

    /// Iterate over the parent params followed by self params.
    pub(crate) fn iter(
        &self,
    ) -> impl DoubleEndedIterator<Item = (GenericParamId, GenericParamDataRef<'_>)> + '_ {
        self.iter_parent().chain(self.iter_self())
    }

    pub(crate) fn iter_parents_with_store(
        &self,
    ) -> impl Iterator<Item = ((GenericParamId, GenericParamDataRef<'_>), &ExpressionStore)> + '_
    {
        self.iter_parent()
            .zip(self.parent_generics().into_iter().flat_map(|it| std::iter::repeat(&*it.store)))
    }

    /// Iterate over the params without parent params.
    pub(crate) fn iter_self(
        &self,
    ) -> impl DoubleEndedIterator<Item = (GenericParamId, GenericParamDataRef<'_>)> + '_ {
        let mut toc = self.params.iter_type_or_consts().map(from_toc_id(self));
        let trait_self_param = self.has_trait_self_param.then(|| toc.next()).flatten();
        chain!(trait_self_param, self.params.iter_lt().map(from_lt_id(self)), toc)
    }

    /// Iterator over types and const params of parent.
    pub(crate) fn iter_parent(
        &self,
    ) -> impl DoubleEndedIterator<Item = (GenericParamId, GenericParamDataRef<'_>)> + '_ {
        self.parent_generics().into_iter().flat_map(|it| {
            let mut toc = it.params.iter_type_or_consts().map(from_toc_id(it));
            let trait_self_param = it.has_trait_self_param.then(|| toc.next()).flatten();
            chain!(trait_self_param, it.params.iter_lt().map(from_lt_id(it)), toc)
        })
    }

    /// Returns total number of generic parameters in scope, including those from parent.
    pub(crate) fn len(&self) -> usize {
        let parent = self.len_parent();
        let child = self.params.len();
        parent + child
    }

    #[inline]
    pub(crate) fn len_parent(&self) -> usize {
        self.parent_generics().map_or(0, Generics::len)
    }

    /// Returns numbers of generic parameters excluding those from parent.
    pub(crate) fn len_self(&self) -> usize {
        self.params.len()
    }

    pub(crate) fn len_lifetimes_self(&self) -> usize {
        self.params.len_lifetimes()
    }

    /// (parent total, self param, type params, const params, impl trait list, lifetimes)
    pub(crate) fn provenance_split(&self) -> (usize, bool, usize, usize, usize, usize) {
        let mut self_param = false;
        let mut type_params = 0;
        let mut impl_trait_params = 0;
        let mut const_params = 0;
        self.params.iter_type_or_consts().for_each(|(_, data)| match data {
            TypeOrConstParamData::TypeParamData(p) => match p.provenance {
                TypeParamProvenance::TypeParamList => type_params += 1,
                TypeParamProvenance::TraitSelf => self_param |= true,
                TypeParamProvenance::ArgumentImplTrait => impl_trait_params += 1,
            },
            TypeOrConstParamData::ConstParamData(_) => const_params += 1,
        });

        let lifetime_params = self.params.len_lifetimes();

        let parent_len = self.parent_generics().map_or(0, Generics::len);
        (parent_len, self_param, type_params, const_params, impl_trait_params, lifetime_params)
    }

    pub(crate) fn type_or_const_param(
        &self,
        param: TypeOrConstParamId,
    ) -> Option<(usize, TypeOrConstParamData)> {
        let idx = self.find_type_or_const_param(param)?;
        self.iter().nth(idx).and_then(|p| {
            let data = match p.1 {
                GenericParamDataRef::TypeParamData(p) => p.clone().into(),
                GenericParamDataRef::ConstParamData(p) => p.clone().into(),
                _ => return None,
            };
            Some((idx, data))
        })
    }

    pub fn type_or_const_param_idx(&self, param: TypeOrConstParamId) -> Option<usize> {
        self.find_type_or_const_param(param)
    }

    fn find_type_or_const_param(&self, param: TypeOrConstParamId) -> Option<usize> {
        if param.parent == self.def {
            let idx = param.local_id.into_raw().into_u32() as usize;
            debug_assert!(
                idx <= self.params.len_type_or_consts(),
                "idx: {} len: {}",
                idx,
                self.params.len_type_or_consts()
            );
            if self.params.trait_self_param() == Some(param.local_id) {
                return Some(idx);
            }
            Some(self.parent_generics().map_or(0, |g| g.len()) + self.params.len_lifetimes() + idx)
        } else {
            debug_assert_eq!(self.parent_generics().map(|it| it.def), Some(param.parent));
            self.parent_generics().and_then(|g| g.find_type_or_const_param(param))
        }
    }

    pub fn lifetime_idx(&self, lifetime: LifetimeParamId) -> Option<usize> {
        self.find_lifetime(lifetime)
    }

    fn find_lifetime(&self, lifetime: LifetimeParamId) -> Option<usize> {
        if lifetime.parent == self.def {
            let idx = lifetime.local_id.into_raw().into_u32() as usize;
            debug_assert!(idx <= self.params.len_lifetimes());
            Some(
                self.parent_generics().map_or(0, |g| g.len())
                    + self.params.trait_self_param().is_some() as usize
                    + idx,
            )
        } else {
            debug_assert_eq!(self.parent_generics().map(|it| it.def), Some(lifetime.parent));
            self.parent_generics().and_then(|g| g.find_lifetime(lifetime))
        }
    }

    pub(crate) fn parent_generics(&self) -> Option<&Generics> {
        self.parent_generics.as_deref()
    }
}

pub(crate) fn parent_generic_def(db: &dyn DefDatabase, def: GenericDefId) -> Option<GenericDefId> {
    let container = match def {
        GenericDefId::FunctionId(it) => it.lookup(db).container,
        GenericDefId::TypeAliasId(it) => it.lookup(db).container,
        GenericDefId::ConstId(it) => it.lookup(db).container,
        GenericDefId::StaticId(_)
        | GenericDefId::AdtId(_)
        | GenericDefId::TraitId(_)
        | GenericDefId::ImplId(_) => return None,
    };

    match container {
        ItemContainerId::ImplId(it) => Some(it.into()),
        ItemContainerId::TraitId(it) => Some(it.into()),
        ItemContainerId::ModuleId(_) | ItemContainerId::ExternBlockId(_) => None,
    }
}

fn from_toc_id<'a>(
    it: &'a Generics,
) -> impl Fn(
    (LocalTypeOrConstParamId, &'a TypeOrConstParamData),
) -> (GenericParamId, GenericParamDataRef<'a>) {
    move |(local_id, p): (_, _)| {
        let id = TypeOrConstParamId { parent: it.def, local_id };
        match p {
            TypeOrConstParamData::TypeParamData(p) => (
                GenericParamId::TypeParamId(TypeParamId::from_unchecked(id)),
                GenericParamDataRef::TypeParamData(p),
            ),
            TypeOrConstParamData::ConstParamData(p) => (
                GenericParamId::ConstParamId(ConstParamId::from_unchecked(id)),
                GenericParamDataRef::ConstParamData(p),
            ),
        }
    }
}

fn from_lt_id<'a>(
    it: &'a Generics,
) -> impl Fn((LocalLifetimeParamId, &'a LifetimeParamData)) -> (GenericParamId, GenericParamDataRef<'a>)
{
    move |(local_id, p): (_, _)| {
        (
            GenericParamId::LifetimeParamId(LifetimeParamId { parent: it.def, local_id }),
            GenericParamDataRef::LifetimeParamData(p),
        )
    }
}