//! Typechecker for AILang (MVP).
//!
//! Monomorphic HM subset: no type variables in the body; all top-level defs
//! must be fully annotated. Effects are propagated as a set and reconciled
//! against the annotation on the function type.
//!
//! Built-in operations are resolved via [`Builtins`].

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

pub mod builtins;
pub mod diagnostic;

pub use diagnostic::{Diagnostic, Severity};

#[derive(Debug, thiserror::Error)]
pub enum CheckError {
    #[error("def `{0}`: {1}")]
    Def(String, Box<CheckError>),

    #[error("type mismatch: expected {expected}, got {got}")]
    TypeMismatch { expected: String, got: String },

    #[error("unknown identifier: `{0}`")]
    UnknownIdent(String),

    #[error("unknown effect operation: `{0}`")]
    UnknownEffectOp(String),

    #[error("`{0}` is not a function (got {1})")]
    NotAFunction(String, String),

    #[error("arity mismatch for `{name}`: expected {expected} args, got {got}")]
    ArityMismatch {
        name: String,
        expected: usize,
        got: usize,
    },

    #[error("undeclared effect `{0}` used in body")]
    UndeclaredEffect(String),

    #[error("function type required for fn `{0}`, got {1}")]
    FnTypeRequired(String, String),

    #[error("param count mismatch in `{name}`: type has {ty_count}, params has {param_count}")]
    ParamCountMismatch {
        name: String,
        ty_count: usize,
        param_count: usize,
    },

    #[error("polymorphic types not supported in MVP body of `{0}`")]
    PolymorphicNotSupported(String),

    #[error("const `{0}` may not have effects (got !{1:?})")]
    ConstHasEffects(String, Vec<String>),

    #[error("unknown type: `{0}`")]
    UnknownType(String),

    #[error("type `{ty}` has no constructor `{ctor}`")]
    UnknownCtor { ty: String, ctor: String },

    #[error("unknown constructor `{0}` in pattern")]
    UnknownCtorInPattern(String),

    #[error("constructor `{ty}/{ctor}` arity: expected {expected} fields, got {got}")]
    CtorArity {
        ty: String,
        ctor: String,
        expected: usize,
        got: usize,
    },

    #[error("non-exhaustive match on `{ty}`: missing cases {missing:?}")]
    NonExhaustive { ty: String, missing: Vec<String> },

    #[error("primitive type `{0}` requires a wildcard or variable arm in match")]
    PrimitiveNeedsWildcard(String),

    #[error("cannot match constructor pattern `{ctor}` against type `{ty}`")]
    PatternTypeMismatch { ctor: String, ty: String },

    #[error("duplicate type definition: `{0}`")]
    DuplicateType(String),

    #[error("duplicate constructor: `{ctor}` (in types `{a}` and `{b}`)")]
    DuplicateCtor { ctor: String, a: String, b: String },

    #[error("duplicate definition: `{0}`")]
    DuplicateDef(String),

    #[error("nested constructor pattern not allowed in MVP: `{0}`")]
    NestedCtorPatternNotAllowed(String),

    #[error("unknown module prefix `{module}` in qualified reference")]
    UnknownModule { module: String },

    #[error("module `{module}` has no top-level def `{name}`")]
    UnknownImport { module: String, name: String },

    #[error("invalid def name `{name}`: contains `.` (reserved for qualified refs)")]
    InvalidDefName { name: String },
}

type Result<T> = std::result::Result<T, CheckError>;

impl CheckError {
    /// Stable kebab-case code for machine consumption (`ail check --json`).
    /// Passed through recursively via the `Def` wrapping — the inner error
    /// carries the actual code, the wrapper only the def context.
    pub fn code(&self) -> &'static str {
        match self {
            CheckError::Def(_, inner) => inner.code(),
            CheckError::TypeMismatch { .. } => "type-mismatch",
            CheckError::UnknownIdent(_) => "unbound-var",
            CheckError::UnknownEffectOp(_) => "unknown-effect-op",
            CheckError::NotAFunction(..) => "not-a-function",
            CheckError::ArityMismatch { .. } => "arity-mismatch",
            CheckError::UndeclaredEffect(_) => "undeclared-effect",
            CheckError::FnTypeRequired(..) => "fn-type-required",
            CheckError::ParamCountMismatch { .. } => "param-count-mismatch",
            CheckError::PolymorphicNotSupported(_) => "polymorphic-not-supported",
            CheckError::ConstHasEffects(..) => "const-has-effects",
            CheckError::UnknownType(_) => "unknown-type",
            CheckError::UnknownCtor { .. } => "unknown-ctor",
            CheckError::UnknownCtorInPattern(_) => "unknown-ctor-in-pattern",
            CheckError::CtorArity { .. } => "arity-mismatch",
            CheckError::NonExhaustive { .. } => "non-exhaustive-match",
            CheckError::PrimitiveNeedsWildcard(_) => "primitive-needs-wildcard",
            CheckError::PatternTypeMismatch { .. } => "pattern-type-mismatch",
            CheckError::DuplicateType(_) => "duplicate-type",
            CheckError::DuplicateCtor { .. } => "duplicate-ctor",
            CheckError::DuplicateDef(_) => "duplicate-def",
            CheckError::NestedCtorPatternNotAllowed(_) => "nested-ctor-pattern-not-allowed",
            CheckError::UnknownModule { .. } => "unknown-module",
            CheckError::UnknownImport { .. } => "unknown-import",
            CheckError::InvalidDefName { .. } => "invalid-def-name",
        }
    }

    /// Structured context for a diagnostic. Lands directly in the JSON
    /// under the key `ctx`. Empty object when no context is available.
    pub fn ctx(&self) -> serde_json::Value {
        match self {
            CheckError::Def(_, inner) => inner.ctx(),
            CheckError::TypeMismatch { expected, got } => {
                serde_json::json!({"expected": expected, "actual": got})
            }
            CheckError::ArityMismatch { expected, got, .. } => {
                serde_json::json!({"expected": expected, "actual": got})
            }
            CheckError::CtorArity {
                expected, got, ..
            } => serde_json::json!({"expected": expected, "actual": got}),
            CheckError::ParamCountMismatch {
                ty_count,
                param_count,
                ..
            } => serde_json::json!({"expected": ty_count, "actual": param_count}),
            CheckError::NonExhaustive { missing, .. } => {
                serde_json::json!({"missing": missing})
            }
            CheckError::UnknownModule { module } => {
                serde_json::json!({"module": module})
            }
            CheckError::UnknownImport { module, name } => {
                serde_json::json!({"module": module, "name": name})
            }
            CheckError::InvalidDefName { name } => {
                serde_json::json!({"name": name, "reason": "contains-dot"})
            }
            _ => serde_json::Value::Object(serde_json::Map::new()),
        }
    }

    /// If this error is wrapped by [`CheckError::Def`], returns the name
    /// of the affected def. Otherwise `None`.
    pub fn def(&self) -> Option<&str> {
        match self {
            CheckError::Def(n, _) => Some(n.as_str()),
            _ => None,
        }
    }

    /// Unwraps the error potentially wrapped by [`CheckError::Def`].
    pub fn inner(&self) -> &CheckError {
        match self {
            CheckError::Def(_, inner) => inner.inner(),
            other => other,
        }
    }

    /// Non-`Def`-wrapped message. Without the `def: ...` prefix.
    pub fn message(&self) -> String {
        format!("{}", self.inner())
    }

    pub fn to_diagnostic(&self) -> Diagnostic {
        let mut d = Diagnostic::error(self.code(), self.message()).with_ctx(self.ctx());
        if let Some(name) = self.def() {
            d = d.with_def(name);
        }
        d
    }
}

/// Top-level API for structured diagnostics.
///
/// Empty Vec = green. From Iter 6 onwards the body-check phase is
/// multi-diagnose: each def in each module is checked independently and
/// failures accumulate, so a single `ail check` run reports every
/// independent body error in the workspace.
///
/// Pass-1 errors (top-level symbol-table construction:
/// `invalid-def-name`, `duplicate-def`) are still fail-fast — those
/// errors corrupt the symbol table, and any further diagnostic would be
/// unreliable. Likewise, the type-def installation is fail-fast within
/// a single module, but other modules continue being checked.
///
/// Backwards compatibility: a bare `&Module` is internally lifted into a
/// trivial workspace (`modules = {m.name: m}`, `entry = m.name`) so that
/// tooling checking individual modules avoids building a `Workspace`.
/// Modules with imports on other modules not present in the trivial
/// workspace will inevitably produce `unknown-module` errors on qualified
/// references — which is correct.
pub fn check_module(m: &Module) -> Vec<Diagnostic> {
    let mut modules = BTreeMap::new();
    modules.insert(m.name.clone(), m.clone());
    let ws = Workspace {
        entry: m.name.clone(),
        modules,
        root_dir: std::path::PathBuf::from("."),
    };
    check_workspace(&ws)
}

/// Top-level API for cross-module typecheck.
///
/// Iterates over all modules of the workspace and checks each with access
/// to the top-level symbol tables of all other modules. Qualified
/// references are resolved via the import map of the respective module:
/// `Term::Var { name }` with exactly one dot in the name is interpreted
/// as `<prefix>.<def>`; `<prefix>` is an import alias (or the module name,
/// if imported without an alias).
///
/// Multi-diagnose: pass-2 collects diagnostics per def across all modules.
/// Pass-1 (symbol table) stays fail-fast — see `check_module`.
/// Module iteration order is deterministic: entry first, then the rest in
/// BTreeMap order, so output ordering is stable.
pub fn check_workspace(ws: &Workspace) -> Vec<Diagnostic> {
    // Pass 1: build per-module top-level symbol table — without checking
    // bodies. This lets module A access defs from module B even when B
    // comes later in the BTreeMap. Duplicate def names and dot-in-def
    // names are reported here immediately, because without clean symbol
    // tables all further diagnostics would be unreliable.
    let module_globals = match build_module_globals(ws) {
        Ok(g) => g,
        Err(e) => return vec![e.to_diagnostic()],
    };

    // Pass 2: body-check per module. `check_in_workspace` builds the env
    // with additional cross-module globals and an import map. Errors
    // accumulate across modules.
    let mut order: Vec<&String> = Vec::new();
    if ws.modules.contains_key(&ws.entry) {
        order.push(&ws.entry);
    }
    for name in ws.modules.keys() {
        if name != &ws.entry {
            order.push(name);
        }
    }
    let mut diagnostics: Vec<Diagnostic> = Vec::new();
    for name in order {
        let m = &ws.modules[name];
        for e in check_in_workspace(m, ws, &module_globals) {
            diagnostics.push(e.to_diagnostic());
        }
    }
    diagnostics
}

/// Result of typechecking a module: mapping from symbol name to
/// (type, hash) — ready for `manifest` output.
#[derive(Debug, Clone)]
pub struct CheckedModule {
    pub symbols: IndexMap<String, (Type, String)>,
}

pub fn check(m: &Module) -> Result<CheckedModule> {
    // Trivial workspace: the module alone, without cross-module resolution.
    let mut modules = BTreeMap::new();
    modules.insert(m.name.clone(), m.clone());
    let ws = Workspace {
        entry: m.name.clone(),
        modules,
        root_dir: std::path::PathBuf::from("."),
    };
    let module_globals = build_module_globals(&ws)?;
    // `check` keeps single-error semantics for callers (snapshot tests,
    // legacy code). Multi-diagnose is exposed via `check_module` /
    // `check_workspace`.
    if let Some(first) = check_in_workspace(m, &ws, &module_globals).into_iter().next() {
        return Err(first);
    }
    // Collect symbols for the return value (existing semantics).
    let mut symbols = IndexMap::new();
    for def in &m.defs {
        let h = ailang_core::hash::def_hash(def);
        let ty = match def {
            Def::Fn(f) => f.ty.clone(),
            Def::Const(c) => c.ty.clone(),
            Def::Type(_) => Type::Con {
                name: def.name().to_string(),
            },
        };
        symbols.insert(def.name().to_string(), (ty, h));
    }
    Ok(CheckedModule { symbols })
}

/// Builds the top-level symbol table per module (for cross-module lookup),
/// without checking bodies. Duplicates and dot-in-def names are reported
/// here as errors immediately — they would taint all further diagnostics.
fn build_module_globals(
    ws: &Workspace,
) -> Result<BTreeMap<String, IndexMap<String, Type>>> {
    let mut out: BTreeMap<String, IndexMap<String, Type>> = BTreeMap::new();
    for (mname, m) in &ws.modules {
        let mut globals = IndexMap::new();
        for def in &m.defs {
            let def_name = def.name();
            if def_name.contains('.') {
                return Err(CheckError::Def(
                    def_name.to_string(),
                    Box::new(CheckError::InvalidDefName {
                        name: def_name.to_string(),
                    }),
                ));
            }
            if globals.contains_key(def_name) {
                return Err(CheckError::Def(
                    def_name.to_string(),
                    Box::new(CheckError::DuplicateDef(def_name.to_string())),
                ));
            }
            let ty = match def {
                Def::Fn(f) => f.ty.clone(),
                Def::Const(c) => c.ty.clone(),
                Def::Type(_) => Type::Con {
                    name: def_name.to_string(),
                },
            };
            globals.insert(def_name.to_string(), ty);
        }
        out.insert(mname.clone(), globals);
    }
    Ok(out)
}

/// Checks the bodies of a single module in the context of the workspace.
/// Assumption: `module_globals` already contains the top-level symbol
/// tables for **all** modules of the workspace (including `m`) — built
/// by `build_module_globals`.
///
/// Returns **all** errors found in this module, in def declaration order:
///
/// - The type-def setup phase is fail-fast within the module (duplicate
///   type or ctor names corrupt the env, so we abort *this* module after
///   the first such error and let the outer loop continue with others).
/// - The body-check phase is multi-diagnose: each def is checked
///   independently against the assembled env; a failure is recorded and
///   the next def is attempted.
fn check_in_workspace(
    m: &Module,
    ws: &Workspace,
    module_globals: &BTreeMap<String, IndexMap<String, Type>>,
) -> Vec<CheckError> {
    let mut env = Env::new();
    builtins::install(&mut env);
    let mut errors: Vec<CheckError> = Vec::new();

    // Register type defs (local per module; cross-module ADT sharing is
    // explicitly not part of 5b).
    for def in &m.defs {
        if let Def::Type(td) = def {
            if env.types.contains_key(&td.name) {
                errors.push(CheckError::Def(
                    td.name.clone(),
                    Box::new(CheckError::DuplicateType(td.name.clone())),
                ));
                return errors;
            }
            for c in &td.ctors {
                if let Some(prev) = env.ctor_index.get(&c.name) {
                    errors.push(CheckError::Def(
                        td.name.clone(),
                        Box::new(CheckError::DuplicateCtor {
                            ctor: c.name.clone(),
                            a: prev.type_name.clone(),
                            b: td.name.clone(),
                        }),
                    ));
                    return errors;
                }
                env.ctor_index.insert(
                    c.name.clone(),
                    CtorRef {
                        type_name: td.name.clone(),
                    },
                );
            }
            env.types.insert(td.name.clone(), td.clone());
        }
    }

    // Take local globals from the previously built table.
    if let Some(g) = module_globals.get(&m.name) {
        for (n, t) in g {
            env.globals.insert(n.clone(), t.clone());
        }
    }

    // Build import map: alias (or module name, if without alias) →
    // module name. Conflicts are not allowed in the MVP: the same `as`
    // clause twice would stand out and should surface as a duplicate
    // symbol name — currently "last wins", because Iter 5b doesn't
    // introduce a dedicated diagnostic for it; if needed later →
    // `ambiguous-import` code.
    let mut import_map: BTreeMap<String, String> = BTreeMap::new();
    for imp in &m.imports {
        let key = imp.alias.clone().unwrap_or_else(|| imp.module.clone());
        import_map.insert(key, imp.module.clone());
    }
    env.imports = import_map;
    env.module_globals = module_globals.clone();
    env.current_module = m.name.clone();
    // Workspace isn't directly needed in the env; cross-module lookup uses
    // only `module_globals`. But we keep the ws reference in the
    // comment as a reminder, in case cross-module ADTs are added later.
    let _ = ws;

    for def in &m.defs {
        if let Err(e) = check_def(def, &env) {
            errors.push(CheckError::Def(def.name().to_string(), Box::new(e)));
        }
    }
    errors
}

fn check_def(def: &Def, env: &Env) -> Result<()> {
    match def {
        Def::Fn(f) => check_fn(f, env),
        Def::Const(c) => check_const(c, env),
        Def::Type(td) => check_type_def(td, env),
    }
}

fn check_type_def(td: &TypeDef, env: &Env) -> Result<()> {
    // All fields must reference known types (or other ADTs from this
    // module; recursion is allowed).
    for c in &td.ctors {
        for f in &c.fields {
            check_type_well_formed(f, env)?;
        }
    }
    Ok(())
}

fn check_type_well_formed(t: &Type, env: &Env) -> Result<()> {
    match t {
        Type::Con { name } => {
            if matches!(name.as_str(), "Int" | "Bool" | "Unit" | "Str") {
                Ok(())
            } else if env.types.contains_key(name) {
                Ok(())
            } else {
                Err(CheckError::UnknownType(name.clone()))
            }
        }
        Type::Fn { params, ret, .. } => {
            for p in params {
                check_type_well_formed(p, env)?;
            }
            check_type_well_formed(ret, env)
        }
        Type::Var { .. } | Type::Forall { .. } => {
            // No type-level polymorphism inside ADT fields in the MVP.
            Err(CheckError::PolymorphicNotSupported(
                "type def".into(),
            ))
        }
    }
}

fn check_fn(f: &FnDef, env: &Env) -> Result<()> {
    let (param_tys, ret_ty, declared_effs) = match &f.ty {
        Type::Fn { params, ret, effects } => {
            (params.clone(), (**ret).clone(), effects.clone())
        }
        other => {
            return Err(CheckError::FnTypeRequired(
                f.name.clone(),
                ailang_core::pretty::type_to_string(other),
            ));
        }
    };

    if f.params.len() != param_tys.len() {
        return Err(CheckError::ParamCountMismatch {
            name: f.name.clone(),
            ty_count: param_tys.len(),
            param_count: f.params.len(),
        });
    }

    let mut locals = IndexMap::new();
    for (n, t) in f.params.iter().zip(param_tys.iter()) {
        locals.insert(n.clone(), t.clone());
    }
    let mut effects = BTreeSet::new();
    let body_ty = synth(&f.body, env, &mut locals, &mut effects, &f.name)?;
    expect_eq(&ret_ty, &body_ty)?;

    let declared: BTreeSet<String> = declared_effs.into_iter().collect();
    for e in &effects {
        if !declared.contains(e) {
            return Err(CheckError::UndeclaredEffect(e.clone()));
        }
    }
    Ok(())
}

fn check_const(c: &ConstDef, env: &Env) -> Result<()> {
    let mut locals = IndexMap::new();
    let mut effects = BTreeSet::new();
    let v = synth(&c.value, env, &mut locals, &mut effects, &c.name)?;
    expect_eq(&c.ty, &v)?;
    if !effects.is_empty() {
        return Err(CheckError::ConstHasEffects(
            c.name.clone(),
            effects.into_iter().collect(),
        ));
    }
    Ok(())
}

fn synth(
    t: &Term,
    env: &Env,
    locals: &mut IndexMap<String, Type>,
    effects: &mut BTreeSet<String>,
    in_def: &str,
) -> Result<Type> {
    match t {
        Term::Lit { lit } => Ok(match lit {
            Literal::Int { .. } => Type::int(),
            Literal::Bool { .. } => Type::bool_(),
            Literal::Str { .. } => Type::str_(),
            Literal::Unit => Type::unit(),
        }),
        Term::Var { name } => {
            // 1) Locals have highest priority — they can even shadow a
            //    qualified dotted name, if someone builds a
            //    letter-with-dot param. Not really reachable in the MVP,
            //    but harmless.
            if let Some(t) = locals.get(name) {
                return Ok(t.clone());
            }
            // 2) Local globals.
            if let Some(t) = env.globals.get(name) {
                return Ok(t.clone());
            }
            // 3) Exactly one dot → qualified cross-module reference.
            //    More than one dot is undefined in the current MVP
            //    and falls through below as `unbound-var`.
            if name.matches('.').count() == 1 {
                let (prefix, suffix) = name.split_once('.').expect("checked");
                let target_module = match env.imports.get(prefix) {
                    Some(m) => m.clone(),
                    None => {
                        // Self-reference `<self>.def` without an import entry:
                        // we deliberately disallow this — by convention,
                        // qualified references only go through imports.
                        // That way the meaning of `name` stays locally stable.
                        return Err(CheckError::UnknownModule {
                            module: prefix.to_string(),
                        });
                    }
                };
                let g = env.module_globals.get(&target_module).ok_or_else(|| {
                    // Import map points at a module not loaded in the
                    // workspace. The workspace loader should have caught
                    // this already; defensive unknown-module here.
                    CheckError::UnknownModule {
                        module: target_module.clone(),
                    }
                })?;
                return g.get(suffix).cloned().ok_or_else(|| {
                    CheckError::UnknownImport {
                        module: target_module,
                        name: suffix.to_string(),
                    }
                });
            }
            Err(CheckError::UnknownIdent(name.clone()))
        }
        Term::App { callee, args } => {
            let cty = synth(callee, env, locals, effects, in_def)?;
            let (params, ret, fx) = match &cty {
                Type::Fn { params, ret, effects: fx } => {
                    (params.clone(), (**ret).clone(), fx.clone())
                }
                Type::Forall { .. } => {
                    return Err(CheckError::PolymorphicNotSupported(in_def.to_string()));
                }
                other => {
                    return Err(CheckError::NotAFunction(
                        callee_name(callee),
                        ailang_core::pretty::type_to_string(other),
                    ));
                }
            };
            if args.len() != params.len() {
                return Err(CheckError::ArityMismatch {
                    name: callee_name(callee),
                    expected: params.len(),
                    got: args.len(),
                });
            }
            for (a, exp) in args.iter().zip(params.iter()) {
                let actual = synth(a, env, locals, effects, in_def)?;
                expect_eq(exp, &actual)?;
            }
            for e in fx {
                effects.insert(e);
            }
            Ok(ret)
        }
        Term::Let { name, value, body } => {
            let v = synth(value, env, locals, effects, in_def)?;
            let prev = locals.insert(name.clone(), v);
            let r = synth(body, env, locals, effects, in_def)?;
            match prev {
                Some(p) => {
                    locals.insert(name.clone(), p);
                }
                None => {
                    locals.shift_remove(name);
                }
            }
            Ok(r)
        }
        Term::If { cond, then, else_ } => {
            let c = synth(cond, env, locals, effects, in_def)?;
            expect_eq(&Type::bool_(), &c)?;
            let t1 = synth(then, env, locals, effects, in_def)?;
            let t2 = synth(else_, env, locals, effects, in_def)?;
            expect_eq(&t1, &t2)?;
            Ok(t1)
        }
        Term::Do { op, args } => {
            let sig = env
                .effect_ops
                .get(op)
                .ok_or_else(|| CheckError::UnknownEffectOp(op.clone()))?
                .clone();
            if args.len() != sig.params.len() {
                return Err(CheckError::ArityMismatch {
                    name: op.clone(),
                    expected: sig.params.len(),
                    got: args.len(),
                });
            }
            for (a, exp) in args.iter().zip(sig.params.iter()) {
                let actual = synth(a, env, locals, effects, in_def)?;
                expect_eq(exp, &actual)?;
            }
            effects.insert(sig.effect.clone());
            Ok(sig.ret)
        }
        Term::Ctor { type_name, ctor, args } => {
            let td = env
                .types
                .get(type_name)
                .ok_or_else(|| CheckError::UnknownType(type_name.clone()))?
                .clone();
            let cdef = td
                .ctors
                .iter()
                .find(|c| &c.name == ctor)
                .ok_or_else(|| CheckError::UnknownCtor {
                    ty: type_name.clone(),
                    ctor: ctor.clone(),
                })?
                .clone();
            if args.len() != cdef.fields.len() {
                return Err(CheckError::CtorArity {
                    ty: type_name.clone(),
                    ctor: ctor.clone(),
                    expected: cdef.fields.len(),
                    got: args.len(),
                });
            }
            for (a, exp) in args.iter().zip(cdef.fields.iter()) {
                let actual = synth(a, env, locals, effects, in_def)?;
                expect_eq(exp, &actual)?;
            }
            Ok(Type::Con {
                name: type_name.clone(),
            })
        }
        Term::Match { scrutinee, arms } => {
            let s_ty = synth(scrutinee, env, locals, effects, in_def)?;
            if arms.is_empty() {
                return Err(CheckError::NonExhaustive {
                    ty: ailang_core::pretty::type_to_string(&s_ty),
                    missing: vec!["(no arms)".into()],
                });
            }
            let mut covered_ctors: BTreeSet<String> = BTreeSet::new();
            let mut has_open_arm = false;
            let mut result_ty: Option<Type> = None;

            for arm in arms {
                // Collect local bindings and push into the env, check the
                // body, pop again — done manually because patterns can
                // produce multiple bindings.
                let bindings = type_check_pattern(&arm.pat, &s_ty, env)?;
                let mut pushed = Vec::new();
                for (n, t) in &bindings {
                    let prev = locals.insert(n.clone(), t.clone());
                    pushed.push((n.clone(), prev));
                }
                let body_ty = synth(&arm.body, env, locals, effects, in_def)?;
                // Undo bindings.
                for (n, prev) in pushed.into_iter().rev() {
                    match prev {
                        Some(p) => {
                            locals.insert(n, p);
                        }
                        None => {
                            locals.shift_remove(&n);
                        }
                    }
                }

                if let Some(rt) = &result_ty {
                    expect_eq(rt, &body_ty)?;
                } else {
                    result_ty = Some(body_ty);
                }

                match &arm.pat {
                    Pattern::Wild | Pattern::Var { .. } => {
                        has_open_arm = true;
                    }
                    Pattern::Ctor { ctor, .. } => {
                        covered_ctors.insert(ctor.clone());
                    }
                    Pattern::Lit { .. } => {
                        // Lit patterns don't structurally cover anything.
                    }
                }
            }

            // Exhaustiveness.
            if !has_open_arm {
                match &s_ty {
                    Type::Con { name } if env.types.contains_key(name) => {
                        let td = &env.types[name];
                        let missing: Vec<String> = td
                            .ctors
                            .iter()
                            .filter(|c| !covered_ctors.contains(&c.name))
                            .map(|c| c.name.clone())
                            .collect();
                        if !missing.is_empty() {
                            return Err(CheckError::NonExhaustive {
                                ty: name.clone(),
                                missing,
                            });
                        }
                    }
                    _ => {
                        return Err(CheckError::PrimitiveNeedsWildcard(
                            ailang_core::pretty::type_to_string(&s_ty),
                        ));
                    }
                }
            }

            Ok(result_ty.expect("checked arms is non-empty"))
        }
        Term::Seq { lhs, rhs } => {
            // Iter 10: `lhs ; rhs`. lhs must be Unit (its value is
            // discarded). Effects from both sides accumulate. The
            // expression's type is rhs's type.
            let lty = synth(lhs, env, locals, effects, in_def)?;
            expect_eq(&Type::unit(), &lty)?;
            synth(rhs, env, locals, effects, in_def)
        }
        Term::Lam { params, param_tys, ret_ty, effects: lam_effects, body } => {
            // Iter 8b: a lambda's type is the declared `Type::Fn`. Push
            // params as locals, check the body's type matches `ret_ty`,
            // and ensure body effects are a subset of the declared set.
            // (We trust the JSON schema for params.len() == param_tys.len();
            // serde would have rejected a malformed AST.)
            let mut pushed: Vec<(String, Option<Type>)> = Vec::new();
            for (n, t) in params.iter().zip(param_tys.iter()) {
                let prev = locals.insert(n.clone(), t.clone());
                pushed.push((n.clone(), prev));
            }
            let mut body_effects: BTreeSet<String> = BTreeSet::new();
            let body_ty = synth(body, env, locals, &mut body_effects, in_def);
            for (n, prev) in pushed.into_iter().rev() {
                match prev {
                    Some(p) => {
                        locals.insert(n, p);
                    }
                    None => {
                        locals.shift_remove(&n);
                    }
                }
            }
            let body_ty = body_ty?;
            expect_eq(ret_ty, &body_ty)?;
            // Constructing a lambda is pure — the body's effects are
            // sealed into the lambda's type, not propagated to the
            // outer effect set. Calling the lambda (Term::App) will
            // pick those effects up via Type::Fn.effects.
            // Body effects must be a subset of the declared lam_effects.
            let declared: BTreeSet<String> = lam_effects.iter().cloned().collect();
            for e in &body_effects {
                if !declared.contains(e) {
                    return Err(CheckError::UndeclaredEffect(e.clone()));
                }
            }
            Ok(Type::Fn {
                params: param_tys.clone(),
                ret: Box::new((**ret_ty).clone()),
                effects: lam_effects.clone(),
            })
        }
    }
}

/// Checks a pattern against an expected type and returns the bindings
/// introduced by the pattern.
fn type_check_pattern(
    p: &Pattern,
    expected: &Type,
    env: &Env,
) -> Result<Vec<(String, Type)>> {
    match p {
        Pattern::Wild => Ok(vec![]),
        Pattern::Var { name } => Ok(vec![(name.clone(), expected.clone())]),
        Pattern::Lit { lit } => {
            let lt = match lit {
                Literal::Int { .. } => Type::int(),
                Literal::Bool { .. } => Type::bool_(),
                Literal::Str { .. } => Type::str_(),
                Literal::Unit => Type::unit(),
            };
            expect_eq(expected, &lt)?;
            Ok(vec![])
        }
        Pattern::Ctor { ctor, fields } => {
            // MVP: sub-patterns of ctor patterns may only be `Var` or `Wild`.
            // Nested ctor or lit patterns need decision-tree lowering,
            // which we don't have in codegen yet.
            for sub in fields {
                if !matches!(sub, Pattern::Var { .. } | Pattern::Wild) {
                    return Err(CheckError::NestedCtorPatternNotAllowed(ctor.clone()));
                }
            }
            let cref = env
                .ctor_index
                .get(ctor)
                .ok_or_else(|| CheckError::UnknownCtorInPattern(ctor.clone()))?;
            // expected must be this ADT.
            match expected {
                Type::Con { name } if name == &cref.type_name => {}
                _ => {
                    return Err(CheckError::PatternTypeMismatch {
                        ctor: ctor.clone(),
                        ty: ailang_core::pretty::type_to_string(expected),
                    });
                }
            }
            let td = &env.types[&cref.type_name];
            let cdef = td
                .ctors
                .iter()
                .find(|c| &c.name == ctor)
                .expect("indexed ctor exists");
            if fields.len() != cdef.fields.len() {
                return Err(CheckError::CtorArity {
                    ty: cref.type_name.clone(),
                    ctor: ctor.clone(),
                    expected: cdef.fields.len(),
                    got: fields.len(),
                });
            }
            let mut out = Vec::new();
            for (sub, sub_ty) in fields.iter().zip(cdef.fields.iter()) {
                out.extend(type_check_pattern(sub, sub_ty, env)?);
            }
            Ok(out)
        }
    }
}

fn callee_name(t: &Term) -> String {
    match t {
        Term::Var { name } => name.clone(),
        _ => "<expr>".into(),
    }
}

fn expect_eq(expected: &Type, got: &Type) -> Result<()> {
    if expected == got {
        Ok(())
    } else {
        Err(CheckError::TypeMismatch {
            expected: ailang_core::pretty::type_to_string(expected),
            got: ailang_core::pretty::type_to_string(got),
        })
    }
}

#[derive(Debug, Default)]
pub struct Env {
    pub globals: IndexMap<String, Type>,
    pub effect_ops: IndexMap<String, builtins::EffectOpSig>,
    pub types: IndexMap<String, TypeDef>,
    /// Inverse index: ctor name -> reference to the owning ADT.
    pub ctor_index: IndexMap<String, CtorRef>,
    /// Import map: alias-or-module-name → actual module name.
    /// Used when `Term::Var { name }` contains a dot
    /// (qualified cross-module reference).
    pub imports: BTreeMap<String, String>,
    /// Top-level symbol table per module of the workspace.
    /// `check_in_workspace` populates this from `build_module_globals`.
    pub module_globals: BTreeMap<String, IndexMap<String, Type>>,
    /// Name of the currently checked module. Used during var lookup to
    /// treat self-references (module name == own name) as local globals,
    /// without touching the `imports` channel.
    pub current_module: String,
}

#[derive(Debug, Clone)]
pub struct CtorRef {
    pub type_name: String,
}

impl Env {
    fn new() -> Self {
        Self::default()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ailang_core::SCHEMA;

    fn fn_def(name: &str, ty: Type, params: Vec<&str>, body: Term) -> Def {
        Def::Fn(FnDef {
            name: name.into(),
            ty,
            params: params.into_iter().map(|s| s.into()).collect(),
            body,
            doc: None,
        })
    }

    #[test]
    fn checks_simple_arithmetic_fn() {
        let m = Module {
            schema: SCHEMA.into(),
            name: "t".into(),
            imports: vec![],
            defs: vec![fn_def(
                "add",
                Type::Fn {
                    params: vec![Type::int(), Type::int()],
                    ret: Box::new(Type::int()),
                    effects: vec![],
                },
                vec!["a", "b"],
                Term::App {
                    callee: Box::new(Term::Var { name: "+".into() }),
                    args: vec![
                        Term::Var { name: "a".into() },
                        Term::Var { name: "b".into() },
                    ],
                },
            )],
        };
        check(&m).expect("should typecheck");
    }

    #[test]
    fn rejects_type_mismatch() {
        let m = Module {
            schema: SCHEMA.into(),
            name: "t".into(),
            imports: vec![],
            defs: vec![fn_def(
                "bad",
                Type::Fn {
                    params: vec![],
                    ret: Box::new(Type::int()),
                    effects: vec![],
                },
                vec![],
                Term::Lit {
                    lit: Literal::Bool { value: true },
                },
            )],
        };
        let err = check(&m).unwrap_err();
        let msg = format!("{err}");
        assert!(msg.contains("type mismatch"), "got: {msg}");
    }

    #[test]
    fn requires_effect_to_be_declared() {
        let m = Module {
            schema: SCHEMA.into(),
            name: "t".into(),
            imports: vec![],
            defs: vec![fn_def(
                "leaks",
                Type::Fn {
                    params: vec![],
                    ret: Box::new(Type::unit()),
                    effects: vec![], // !IO missing
                },
                vec![],
                Term::Do {
                    op: "io/print_int".into(),
                    args: vec![Term::Lit {
                        lit: Literal::Int { value: 1 },
                    }],
                },
            )],
        };
        let err = check(&m).unwrap_err();
        let msg = format!("{err}");
        assert!(msg.contains("undeclared effect"), "got: {msg}");
    }

    #[test]
    fn lets_local_shadow_global() {
        let m = Module {
            schema: SCHEMA.into(),
            name: "t".into(),
            imports: vec![],
            defs: vec![fn_def(
                "f",
                Type::Fn {
                    params: vec![],
                    ret: Box::new(Type::int()),
                    effects: vec![],
                },
                vec![],
                Term::Let {
                    name: "x".into(),
                    value: Box::new(Term::Lit {
                        lit: Literal::Int { value: 7 },
                    }),
                    body: Box::new(Term::Var { name: "x".into() }),
                },
            )],
        };
        check(&m).expect("should typecheck");
    }

    #[test]
    fn match_must_be_exhaustive() {
        // Type Maybe = None | Some(Int); fn f only matches None -> error.
        let m = Module {
            schema: SCHEMA.into(),
            name: "t".into(),
            imports: vec![],
            defs: vec![
                Def::Type(TypeDef {
                    name: "Maybe".into(),
                    ctors: vec![
                        Ctor { name: "None".into(), fields: vec![] },
                        Ctor {
                            name: "Some".into(),
                            fields: vec![Type::int()],
                        },
                    ],
                    doc: None,
                }),
                fn_def(
                    "f",
                    Type::Fn {
                        params: vec![Type::Con { name: "Maybe".into() }],
                        ret: Box::new(Type::int()),
                        effects: vec![],
                    },
                    vec!["m"],
                    Term::Match {
                        scrutinee: Box::new(Term::Var { name: "m".into() }),
                        arms: vec![Arm {
                            pat: Pattern::Ctor {
                                ctor: "None".into(),
                                fields: vec![],
                            },
                            body: Term::Lit {
                                lit: Literal::Int { value: 0 },
                            },
                        }],
                    },
                ),
            ],
        };
        let err = check(&m).unwrap_err();
        let msg = format!("{err}");
        assert!(msg.contains("non-exhaustive"), "got: {msg}");
        assert!(msg.contains("Some"), "got: {msg}");
    }

    #[test]
    fn match_with_wildcard_is_exhaustive() {
        let m = Module {
            schema: SCHEMA.into(),
            name: "t".into(),
            imports: vec![],
            defs: vec![
                Def::Type(TypeDef {
                    name: "Maybe".into(),
                    ctors: vec![
                        Ctor { name: "None".into(), fields: vec![] },
                        Ctor {
                            name: "Some".into(),
                            fields: vec![Type::int()],
                        },
                    ],
                    doc: None,
                }),
                fn_def(
                    "f",
                    Type::Fn {
                        params: vec![Type::Con { name: "Maybe".into() }],
                        ret: Box::new(Type::int()),
                        effects: vec![],
                    },
                    vec!["m"],
                    Term::Match {
                        scrutinee: Box::new(Term::Var { name: "m".into() }),
                        arms: vec![
                            Arm {
                                pat: Pattern::Ctor {
                                    ctor: "None".into(),
                                    fields: vec![],
                                },
                                body: Term::Lit {
                                    lit: Literal::Int { value: 0 },
                                },
                            },
                            Arm {
                                pat: Pattern::Wild,
                                body: Term::Lit {
                                    lit: Literal::Int { value: 1 },
                                },
                            },
                        ],
                    },
                ),
            ],
        };
        check(&m).expect("wildcard must satisfy exhaustiveness");
    }

    #[test]
    fn if_branches_must_match() {
        let m = Module {
            schema: SCHEMA.into(),
            name: "t".into(),
            imports: vec![],
            defs: vec![fn_def(
                "f",
                Type::Fn {
                    params: vec![],
                    ret: Box::new(Type::int()),
                    effects: vec![],
                },
                vec![],
                Term::If {
                    cond: Box::new(Term::Lit {
                        lit: Literal::Bool { value: true },
                    }),
                    then: Box::new(Term::Lit {
                        lit: Literal::Int { value: 1 },
                    }),
                    else_: Box::new(Term::Lit { lit: Literal::Unit }),
                },
            )],
        };
        let err = check(&m).unwrap_err();
        assert!(format!("{err}").contains("type mismatch"));
    }

    /// Iter 10: `seq` requires lhs to be Unit. A non-Unit lhs is a
    /// type error — the value of lhs gets discarded so a useful (non-
    /// Unit) value would silently vanish.
    #[test]
    fn seq_lhs_must_be_unit() {
        let m = Module {
            schema: SCHEMA.into(),
            name: "t".into(),
            imports: vec![],
            defs: vec![fn_def(
                "f",
                Type::Fn {
                    params: vec![],
                    ret: Box::new(Type::int()),
                    effects: vec![],
                },
                vec![],
                Term::Seq {
                    lhs: Box::new(Term::Lit {
                        lit: Literal::Int { value: 7 },
                    }),
                    rhs: Box::new(Term::Lit {
                        lit: Literal::Int { value: 1 },
                    }),
                },
            )],
        };
        let err = check(&m).unwrap_err();
        let msg = format!("{err}");
        assert!(msg.contains("type mismatch"), "got: {msg}");
    }
}