use std::borrow::Cow;
use std::collections::HashMap;
use std::fmt;
use std::fmt::Display;

use index_vec::Idx;

use crate::ast::*;
use crate::common::TAB_INCR;
use crate::ids::IndexVec;
use crate::pretty::FmtWithCtx;

pub trait IntoFormatter {
    type C: AstFormatter;
    fn into_fmt(self) -> Self::C;
}

/// An [`AstFormatter`] contains the context required to pretty-print the ast. An ast type can then
/// be pretty-printed using the [`FmtWithCtx`] trait.
pub trait AstFormatter: Sized {
    type Reborrow<'a>: AstFormatter + 'a
    where
        Self: 'a;

    fn get_crate(&self) -> Option<&TranslatedCrate>;

    fn no_generics<'a>(&'a self) -> Self::Reborrow<'a>;
    fn set_generics<'a>(&'a self, generics: &'a GenericParams) -> Self::Reborrow<'a>;
    fn set_locals<'a>(&'a self, locals: &'a Locals) -> Self::Reborrow<'a>;
    fn push_binder<'a>(&'a self, new_params: Cow<'a, GenericParams>) -> Self::Reborrow<'a>;
    fn push_bound_regions<'a>(
        &'a self,
        regions: &'a IndexVec<RegionId, RegionParam>,
    ) -> Self::Reborrow<'a> {
        self.push_binder(Cow::Owned(GenericParams {
            regions: regions.clone(),
            ..Default::default()
        }))
    }
    /// Return the depth of binders we're under.
    fn binder_depth(&self) -> usize;

    fn increase_indent<'a>(&'a self) -> Self::Reborrow<'a>;
    fn reset_indent<'a>(&'a self) -> Self::Reborrow<'a>;
    fn indent(&self) -> String;

    fn format_local_id(&self, f: &mut fmt::Formatter<'_>, id: LocalId) -> fmt::Result;
    fn format_bound_var<Id: Idx + Display, T>(
        &self,
        f: &mut fmt::Formatter<'_>,
        var: DeBruijnVar<Id>,
        var_prefix: &str,
        fmt_var: impl Fn(&T) -> Option<String>,
    ) -> fmt::Result
    where
        GenericParams: HasIdxVecOf<Id, Output = T>;

    fn format_method_name(
        &self,
        f: &mut fmt::Formatter<'_>,
        trait_id: TraitDeclId,
        method_id: TraitMethodId,
    ) -> fmt::Result {
        if let Some(translated) = self.get_crate()
            && let Some(names) = translated.assoc_item_names.get(trait_id)
            && let Some(name) = names.methods.get(method_id).copied()
        {
            write!(f, "{name}")
        } else {
            write!(f, "{}", method_id.to_pretty_string())
        }
    }
    fn format_assoc_type_name(
        &self,
        f: &mut fmt::Formatter<'_>,
        trait_id: TraitDeclId,
        type_id: AssocTypeId,
    ) -> fmt::Result {
        if let Some(translated) = self.get_crate()
            && let Some(names) = translated.assoc_item_names.get(trait_id)
            && let Some(name) = names.types.get(type_id).copied()
        {
            write!(f, "{name}")
        } else {
            write!(f, "{}", type_id.to_pretty_string())
        }
    }
    fn format_assoc_const_name(
        &self,
        f: &mut fmt::Formatter<'_>,
        trait_id: TraitDeclId,
        const_id: AssocConstId,
    ) -> fmt::Result {
        if let Some(translated) = self.get_crate()
            && let Some(names) = translated.assoc_item_names.get(trait_id)
            && let Some(name) = names.consts.get(const_id).copied()
        {
            write!(f, "{name}")
        } else {
            write!(f, "{}", const_id.to_pretty_string())
        }
    }
    fn format_enum_variant_name(
        &self,
        f: &mut fmt::Formatter<'_>,
        type_id: TypeDeclId,
        variant_id: VariantId,
    ) -> fmt::Result {
        let variant = if let Some(translated) = self.get_crate()
            && let Some(def) = translated.type_decls.get(type_id)
            && let Some(variants) = def.kind.as_enum()
        {
            &variants.get(variant_id).unwrap().name
        } else {
            &variant_id.to_pretty_string()
        };
        write!(f, "{variant}")
    }
    fn format_enum_variant(
        &self,
        f: &mut fmt::Formatter<'_>,
        type_id: TypeDeclId,
        variant_id: VariantId,
    ) -> fmt::Result {
        write!(f, "{}::", type_id.with_ctx(self))?;
        self.format_enum_variant_name(f, type_id, variant_id)?;
        Ok(())
    }

    fn format_field_name(
        &self,
        f: &mut fmt::Formatter<'_>,
        type_id: TypeDeclId,
        opt_variant_id: Option<VariantId>,
        field_id: FieldId,
    ) -> fmt::Result {
        let field_name = if let Some(translated) = self.get_crate()
            && let Some(def) = translated.type_decls.get(type_id)
        {
            match (&def.kind, opt_variant_id) {
                (TypeDeclKind::Enum(variants), Some(variant_id)) => {
                    variants[variant_id].fields[field_id].name.as_ref()
                }
                (TypeDeclKind::Struct(fields) | TypeDeclKind::Union(fields), None) => {
                    fields[field_id].name.as_ref()
                }
                _ => None,
            }
        } else {
            None
        };
        if let Some(field_name) = field_name {
            write!(f, "{field_name}")
        } else {
            write!(f, "{field_id}")
        }
    }
}

/// Context for formatting.
#[derive(Default)]
pub struct FmtCtx<'a> {
    pub translated: Option<&'a TranslatedCrate>,
    /// Generics form a stack, where each binder introduces a new level. For DeBruijn indices to
    /// work, we keep the innermost parameters at the start of the vector.
    pub generics: BindingStack<Cow<'a, GenericParams>>,
    pub local_names: Option<IndexVec<LocalId, String>>,
    pub indent_level: usize,
}

impl<'c> AstFormatter for FmtCtx<'c> {
    type Reborrow<'a>
        = FmtCtx<'a>
    where
        Self: 'a;

    fn get_crate(&self) -> Option<&TranslatedCrate> {
        self.translated
    }

    fn no_generics<'a>(&'a self) -> Self::Reborrow<'a> {
        FmtCtx {
            generics: BindingStack::empty(),
            ..self.reborrow()
        }
    }
    fn set_generics<'a>(&'a self, generics: &'a GenericParams) -> Self::Reborrow<'a> {
        FmtCtx {
            generics: BindingStack::new(Cow::Borrowed(generics)),
            ..self.reborrow()
        }
    }
    fn set_locals<'a>(&'a self, locals: &'a Locals) -> Self::Reborrow<'a> {
        FmtCtx {
            local_names: Some(compute_local_names(locals)),
            ..self.reborrow()
        }
    }
    fn push_binder<'a>(&'a self, new_params: Cow<'a, GenericParams>) -> Self::Reborrow<'a> {
        let mut ret = self.reborrow();
        ret.generics.push(new_params);
        ret
    }
    fn binder_depth(&self) -> usize {
        self.generics.len()
    }

    fn increase_indent<'a>(&'a self) -> Self::Reborrow<'a> {
        FmtCtx {
            indent_level: self.indent_level + 1,
            ..self.reborrow()
        }
    }
    fn reset_indent<'a>(&'a self) -> Self::Reborrow<'a> {
        FmtCtx {
            indent_level: 0,
            ..self.reborrow()
        }
    }
    fn indent(&self) -> String {
        TAB_INCR.repeat(self.indent_level)
    }

    fn format_local_id(&self, f: &mut fmt::Formatter<'_>, id: LocalId) -> fmt::Result {
        if let Some(local_names) = &self.local_names {
            write!(f, "{}", local_names[id])
        } else {
            write!(f, "_{id}")
        }
    }

    fn format_bound_var<Id: Idx + Display, T>(
        &self,
        f: &mut fmt::Formatter<'_>,
        var: DeBruijnVar<Id>,
        var_prefix: &str,
        fmt_var: impl Fn(&T) -> Option<String>,
    ) -> fmt::Result
    where
        GenericParams: HasIdxVecOf<Id, Output = T>,
    {
        if self.generics.is_empty() {
            return write!(f, "{var_prefix}{var}");
        }
        match self.generics.get_var::<_, GenericParams>(var) {
            None => write!(f, "missing({var_prefix}{var})"),
            Some(v) => match fmt_var(v) {
                Some(name) => write!(f, "{name}"),
                None => {
                    write!(f, "{var_prefix}")?;
                    let (dbid, varid) = self.generics.as_bound_var(var);
                    let depth = self.generics.depth().index - dbid.index;
                    if depth == 0 {
                        write!(f, "{varid}")
                    } else {
                        write!(f, "{varid}_{depth}")
                    }
                }
            },
        }
    }
}

impl<'a> FmtCtx<'a> {
    pub fn new() -> Self {
        FmtCtx::default()
    }

    pub fn get_item(&self, id: ItemId) -> Result<ItemRef<'_>, Option<&Name>> {
        let Some(translated) = &self.translated else {
            return Err(None);
        };
        translated
            .get_item(id)
            .ok_or_else(|| Some(translated.item_short_name(id)))
    }

    /// Print the whole definition.
    pub fn format_decl_id(&self, id: impl Into<ItemId>) -> String {
        let id = id.into();
        match self.get_item(id) {
            Ok(d) => d.to_string_with_ctx(self),
            Err(opt_name) => {
                let opt_name = opt_name
                    .map(|n| format!(" ({})", n.with_ctx(self)))
                    .unwrap_or_default();
                format!("Missing decl: {id:?}{opt_name}")
            }
        }
    }

    fn reborrow<'b>(&'b self) -> FmtCtx<'b> {
        FmtCtx {
            translated: self.translated,
            generics: self.generics.clone(),
            local_names: self.local_names.clone(),
            indent_level: self.indent_level,
        }
    }
}

/// Compute a unique name for each local.
pub fn compute_local_names(locals: &Locals) -> IndexVec<LocalId, String> {
    let mut local_names = locals.locals.map_ref(|local| {
        format!(
            "{}_{}",
            local.name.as_deref().unwrap_or_default(),
            local.index
        )
    });

    let mut name_counts = HashMap::<String, usize>::new();
    for local in &locals.locals {
        *name_counts
            .entry(local_names[local.index].clone())
            .or_default() += 1;
        if let Some(name) = &local.name {
            *name_counts.entry(name.clone()).or_default() += 1;
        }
    }

    for (id, local) in locals.locals.iter_enumerated() {
        if let Some(name) = &local.name
            && !name.is_empty()
            && name_counts[name] == 1
        {
            local_names[id] = name.clone();
        }
    }
    local_names
}
