//! Type substitution + unification helpers for monomorphisation.
//!
//! Free functions extracted from `lib.rs` during the 18g tidy split.
//! The four-step pipeline is: `derive_substitution` walks declared
//! params against actual arg types, calling `unify_for_subst` to bind
//! `Type::Var`s; `apply_subst_to_type` / `apply_subst_to_term`
//! specialise a polymorphic def under that binding;
//! `qualify_local_types_codegen` rewrites bare ADT names into
//! `module.Type` form when a sig crosses an import boundary;
//! `descriptor_for_subst` produces the stable mangling suffix used in
//! the specialised symbol's name.

use ailang_core::ast::*;
use std::collections::{BTreeMap, BTreeSet};

use super::{CodegenError, Result};
use crate::synth::type_descriptor;

/// Iter 12b: derive a name → concrete-type substitution from the
/// declared params of a `Forall` body and the actual arg types at a
/// call site. Walks both sides in parallel; whenever a `Type::Var`
/// (rigid name) appears on the params side, binds it to the
/// corresponding concrete type. Conflicts (same var bound to two
/// different types) surface as an internal error — the typechecker
/// would already have rejected such a call.
pub(crate) fn derive_substitution(
    vars: &[String],
    params: &[Type],
    arg_tys: &[Type],
) -> Result<BTreeMap<String, Type>> {
    if params.len() != arg_tys.len() {
        return Err(CodegenError::Internal(format!(
            "derive_substitution: arity mismatch ({} params vs {} args)",
            params.len(),
            arg_tys.len(),
        )));
    }
    let var_set: BTreeSet<&str> = vars.iter().map(|s| s.as_str()).collect();
    let mut subst: BTreeMap<String, Type> = BTreeMap::new();
    for (p, a) in params.iter().zip(arg_tys.iter()) {
        unify_for_subst(p, a, &var_set, &mut subst)?;
    }
    // Any forall var not pinned by the args is left unbound. For the
    // MVP this is an error — we can't specialise without a concrete
    // type. The typechecker's body should have constrained it already
    // through return-type unification, but at the call site we only
    // see args; if needed, callers can extend this with expected-ret
    // info.
    // Iter 15a: a forall var that the args couldn't pin (e.g.
    // `is_none(Nothing) : forall a. (Maybe a) -> Bool` — `a` is
    // genuinely unobservable from the args alone) defaults to `Unit`.
    // The specialised body must not actually read an `a`-typed value,
    // or it would have failed type-checking; a dummy concrete type is
    // sound and lets monomorphisation proceed deterministically. The
    // descriptor uses the same default, so all such call sites
    // converge on a single specialisation.
    for v in vars {
        if !subst.contains_key(v) {
            subst.insert(v.clone(), Type::unit());
        }
    }
    Ok(subst)
}

/// Walks `param` and `arg` in parallel, treating any `Type::Var { name }`
/// on the param side whose name is in `vars` as an unknown to be bound
/// in `subst`. Identical concrete shapes pass through; structural
/// mismatches yield an internal error.
pub(crate) fn unify_for_subst(
    param: &Type,
    arg: &Type,
    vars: &BTreeSet<&str>,
    subst: &mut BTreeMap<String, Type>,
) -> Result<()> {
    // Iter 14a fix, extended in 15g-aux: a `$u`-prefixed var is a
    // synth-only wildcard produced by `synth_arg_type` for nullary
    // ctors of a parameterised ADT (e.g. `Nil : List<$u>`). It
    // carries no real constraint — accept without binding so a
    // sibling arg can pin the type var instead. Without this,
    // `Cons(Int, Nil)` synth would unify `a = Int` (from head) and
    // then `a = $u` (from tail's recursive `List<a>` slot) and
    // falsely error.
    //
    // 15g-aux: the early-return must accept `$u` on **either** side.
    // `$u` enters in arg position from synth, but the prev-binding
    // recursion below (`unify_for_subst(&prev, arg, ...)`) can swap
    // a `$u` onto the param side when a previously-bound type is
    // unified against a fresher arg whose roles differ. Reduced
    // repro: `length [Left 1, Right 10]` — `a` first binds to
    // `Either<Int, $u>` from `Left 1`, then a recursive unification
    // against `Either<$u, Int>` from `Right 10` lands `$u` in the
    // param-pos[1] slot. Symmetric early-return is correct because
    // `$u` is a synth-only wildcard regardless of which side carries
    // it after the prev-binding swap.
    if let Type::Var { name } = arg {
        if name.starts_with("$u") {
            return Ok(());
        }
    }
    if let Type::Var { name } = param {
        if name.starts_with("$u") {
            return Ok(());
        }
    }
    match (param, arg) {
        (Type::Var { name }, _) if vars.contains(name.as_str()) => {
            if let Some(prev) = subst.get(name).cloned() {
                // Iter 15b: the previously-bound type may be more
                // concrete than `arg` (e.g. `prev = List<Int>` from a
                // sibling binding, `arg = List<$u>` from a synth-
                // wildcard nullary ctor). Use recursive unification
                // instead of strict equality so the inner `$u`
                // wildcard matches `Int`. The previous strict-
                // equality check rejected such overlaps as bogus
                // duplicate bindings.
                return unify_for_subst(&prev, arg, vars, subst);
            }
            subst.insert(name.clone(), arg.clone());
            Ok(())
        }
        (
            Type::Con { name: pn, args: pa },
            Type::Con { name: an, args: aa },
        ) if pn == an && pa.len() == aa.len() => {
            for (p, a) in pa.iter().zip(aa.iter()) {
                unify_for_subst(p, a, vars, subst)?;
            }
            Ok(())
        }
        (
            Type::Fn { params: pp, ret: pr, .. },
            Type::Fn { params: ap, ret: ar, .. },
        ) => {
            if pp.len() != ap.len() {
                return Err(CodegenError::Internal(
                    "monomorphisation: fn arity mismatch in arg".into(),
                ));
            }
            for (p, a) in pp.iter().zip(ap.iter()) {
                unify_for_subst(p, a, vars, subst)?;
            }
            unify_for_subst(pr, ar, vars, subst)
        }
        (Type::Var { name: pn }, Type::Var { name: an }) if pn == an => Ok(()),
        _ => Err(CodegenError::Internal(format!(
            "monomorphisation: cannot match param `{}` to arg `{}`",
            ailang_core::pretty::type_to_string(param),
            ailang_core::pretty::type_to_string(arg),
        ))),
    }
}

/// Iter 15a: rewrites bare `Type::Con` references that resolve against
/// `owner_local_types` into qualified `module.Type` form. Mirrors
/// `ailang_check::qualify_local_types`. Used when the codegen pulls a
/// polymorphic fn signature across the import boundary; without this
/// the substitution derived from the call site's qualified args
/// (`std_maybe.Maybe<Int>`) would fail to unify against the bare
/// signature (`Maybe<a>`).
pub(crate) fn qualify_local_types_codegen(
    t: &Type,
    owner_module: &str,
    owner_local_types: &BTreeSet<String>,
) -> Type {
    match t {
        Type::Con { name, args } => {
            let qualified = if name.contains('.') {
                name.clone()
            } else if ailang_core::primitives::is_primitive_name(name) {
                name.clone()
            } else if owner_local_types.contains(name) {
                format!("{owner_module}.{name}")
            } else {
                name.clone()
            };
            Type::Con {
                name: qualified,
                args: args
                    .iter()
                    .map(|a| qualify_local_types_codegen(a, owner_module, owner_local_types))
                    .collect(),
            }
        }
        Type::Fn { params, ret, effects, .. } => Type::Fn {
            params: params
                .iter()
                .map(|p| qualify_local_types_codegen(p, owner_module, owner_local_types))
                .collect(),
            ret: Box::new(qualify_local_types_codegen(ret, owner_module, owner_local_types)),
            effects: effects.clone(),
            param_modes: vec![],
            ret_mode: ParamMode::Implicit,
        },
        Type::Forall { vars, constraints, body } => Type::Forall {
            vars: vars.clone(),
            constraints: constraints.clone(),
            body: Box::new(qualify_local_types_codegen(body, owner_module, owner_local_types)),
        },
        Type::Var { .. } => t.clone(),
    }
}

/// Iter 12b: substitute rigid type vars in `t` according to `subst`.
/// Used to specialise the type of a polymorphic def for a given
/// instantiation.
pub(crate) fn apply_subst_to_type(t: &Type, subst: &BTreeMap<String, Type>) -> Type {
    match t {
        Type::Var { name } => subst.get(name).cloned().unwrap_or_else(|| t.clone()),
        Type::Con { name, args } => Type::Con {
            name: name.clone(),
            args: args.iter().map(|a| apply_subst_to_type(a, subst)).collect(),
        },
        Type::Fn { params, ret, effects, .. } => Type::Fn {
            params: params.iter().map(|p| apply_subst_to_type(p, subst)).collect(),
            ret: Box::new(apply_subst_to_type(ret, subst)),
            effects: effects.clone(),
            param_modes: vec![],
            ret_mode: ParamMode::Implicit,
        },
        Type::Forall { vars, constraints, body } => {
            // Inner forall shadows: don't substitute re-bound names.
            let inner: BTreeMap<String, Type> = subst
                .iter()
                .filter(|(k, _)| !vars.contains(k))
                .map(|(k, v)| (k.clone(), v.clone()))
                .collect();
            Type::Forall {
                vars: vars.clone(),
                constraints: constraints.clone(),
                body: Box::new(apply_subst_to_type(body, &inner)),
            }
        }
    }
}

/// Iter 12b: substitute rigid type vars throughout a Term. Only
/// `Term::Lam` carries types in the AST (params/ret), so most arms
/// just recurse. `Term::Var` contains a name string only and is
/// left untouched.
pub(crate) fn apply_subst_to_term(t: &Term, subst: &BTreeMap<String, Type>) -> Term {
    match t {
        Term::Lit { .. } | Term::Var { .. } => t.clone(),
        Term::App { callee, args, tail } => Term::App {
            callee: Box::new(apply_subst_to_term(callee, subst)),
            args: args.iter().map(|a| apply_subst_to_term(a, subst)).collect(),
            tail: *tail,
        },
        Term::Let { name, value, body } => Term::Let {
            name: name.clone(),
            value: Box::new(apply_subst_to_term(value, subst)),
            body: Box::new(apply_subst_to_term(body, subst)),
        },
        Term::If { cond, then, else_ } => Term::If {
            cond: Box::new(apply_subst_to_term(cond, subst)),
            then: Box::new(apply_subst_to_term(then, subst)),
            else_: Box::new(apply_subst_to_term(else_, subst)),
        },
        Term::Do { op, args, tail } => Term::Do {
            op: op.clone(),
            args: args.iter().map(|a| apply_subst_to_term(a, subst)).collect(),
            tail: *tail,
        },
        Term::Ctor { type_name, ctor, args } => Term::Ctor {
            type_name: type_name.clone(),
            ctor: ctor.clone(),
            args: args.iter().map(|a| apply_subst_to_term(a, subst)).collect(),
        },
        Term::Match { scrutinee, arms } => Term::Match {
            scrutinee: Box::new(apply_subst_to_term(scrutinee, subst)),
            arms: arms
                .iter()
                .map(|a| Arm {
                    pat: a.pat.clone(),
                    body: apply_subst_to_term(&a.body, subst),
                })
                .collect(),
        },
        Term::Lam { params, param_tys, ret_ty, effects, body } => Term::Lam {
            params: params.clone(),
            param_tys: param_tys
                .iter()
                .map(|t| apply_subst_to_type(t, subst))
                .collect(),
            ret_ty: Box::new(apply_subst_to_type(ret_ty, subst)),
            effects: effects.clone(),
            body: Box::new(apply_subst_to_term(body, subst)),
        },
        Term::Seq { lhs, rhs } => Term::Seq {
            lhs: Box::new(apply_subst_to_term(lhs, subst)),
            rhs: Box::new(apply_subst_to_term(rhs, subst)),
        },
        Term::LetRec { .. } => {
            // Iter 16b.1: eliminated by desugar before any
            // monomorphisation pass runs.
            unreachable!("Term::LetRec eliminated by desugar")
        }
        Term::Clone { value } => Term::Clone {
            // Iter 18c.1: structural recursion through the wrapper.
            value: Box::new(apply_subst_to_term(value, subst)),
        },
        Term::ReuseAs { source, body } => Term::ReuseAs {
            // Iter 18d.1: structural recursion through both children.
            source: Box::new(apply_subst_to_term(source, subst)),
            body: Box::new(apply_subst_to_term(body, subst)),
        },
    }
}

/// Iter 12b: deterministic descriptor string for a substitution. Used
/// as the suffix in the mangled name `@ail_<m>_<def>__<descriptor>`.
/// Vars are emitted in the order given by the FnDef's forall vars
/// (so two call sites with the same instantiation map to the same
/// descriptor regardless of internal BTreeMap ordering).
pub(crate) fn descriptor_for_subst(vars: &[String], subst: &BTreeMap<String, Type>) -> String {
    let mut parts: Vec<String> = Vec::with_capacity(vars.len());
    for v in vars {
        let ty = subst.get(v).cloned().unwrap_or_else(|| Type::unit());
        parts.push(type_descriptor(&ty));
    }
    parts.join("_")
}