plum

// Functions/methods are named camelCase across this project (matching plum's own
// naming convention), not Rust's idiomatic snake_case — silence the resulting lint.
#![allow(non_snake_case)]

pub mod types;
pub mod monomorphize;

use std::collections::BTreeMap;
use types::{PlumType, TypeEnv, TypeScheme, CheckError, CheckResult};
use plum_core::ast;

/// Maps a bare type name (as written in source: a param/return annotation, or
/// a method's receiver name) to its `PlumType`. Builtin primitives get their
/// dedicated variant; anything else is an unmodeled `TNamed`.
pub fn plumTypeFromName(name: &str) -> PlumType {
    match name {
        "Int" => PlumType::TInt,
        "Float" => PlumType::TFloat,
        "Bool" => PlumType::TBool,
        "Str" => PlumType::TStr,
        "Byte" => PlumType::TByte,
        "[]Byte" => PlumType::TByteSlice,
        "Unit" => PlumType::TUnit,
        other => PlumType::TNamed(other.to_string()),
    }
}

pub fn plumTypeFromAst(ty: &ast::Type) -> PlumType {
    plumTypeFromName(&ty.name)
}

/// The `ctx.methods`/`ctx.classes` receiver name for a value of type `ty`, or
/// `None` if `ty` has no methods (e.g. an unresolved `TVar` or a bare
/// function type). Builtin primitive types (`Int`/`Float`/`Bool`/`Str`) are
/// declared as `type Int = fun ...` etc in `libs/std`, exactly like a class —
/// they're just never `TNamed` at the type level, so this maps them back to
/// the same receiver name string `ctx.methods` is keyed by.
pub fn methodReceiverName(ty: &PlumType) -> Option<String> {
    match ty {
        PlumType::TNamed(name) => Some(name.clone()),
        PlumType::TInt => Some("Int".to_string()),
        PlumType::TFloat => Some("Float".to_string()),
        PlumType::TBool => Some("Bool".to_string()),
        PlumType::TStr => Some("Str".to_string()),
        PlumType::TByte => Some("Byte".to_string()),
        PlumType::TByteSlice => Some("ByteSlice".to_string()),
        _ => None,
    }
}

pub fn unify(t1: &PlumType, t2: &PlumType) -> Result<(), String> {
    match (t1, t2) {
        (PlumType::TVar(_), _) | (_, PlumType::TVar(_)) => Ok(()),
        (PlumType::TInt, PlumType::TInt) => Ok(()),
        (PlumType::TFloat, PlumType::TFloat) => Ok(()),
        (PlumType::TBool, PlumType::TBool) => Ok(()),
        (PlumType::TStr, PlumType::TStr) => Ok(()),
        (PlumType::TByte, PlumType::TByte) => Ok(()),
        (PlumType::TByteSlice, PlumType::TByteSlice) => Ok(()),
        (PlumType::TUnit, PlumType::TUnit) => Ok(()),
        (PlumType::TNamed(a), PlumType::TNamed(b)) if a == b => Ok(()),
        (PlumType::TFun(ps1, r1), PlumType::TFun(ps2, r2)) if ps1.len() == ps2.len() => {
            for (p1, p2) in ps1.iter().zip(ps2.iter()) {
                unify(p1, p2)?;
            }
            unify(r1, r2)
        }
        _ => Err(format!("type mismatch: expected {}, found {}", t1, t2)),
    }
}

pub fn lookup(env: &TypeEnv, name: &str) -> Result<PlumType, String> {
    env.get(name)
        .map(|s| *s.body.clone())
        .ok_or_else(|| format!("undefined name '{}'", name))
}

/// Field names and types for every `type ClassName = ...` declaration in the source.
pub type ClassEnv = BTreeMap<String, Vec<(String, PlumType)>>;
/// `(receiver type, method name) -> TFun` for every `name<Receiver>(...)` method.
pub type MethodEnv = BTreeMap<(String, String), PlumType>;
/// Info about one `enum` variant: which enum it belongs to, its 0-based runtime tag
/// (numbering is shared across all of that enum's variants), and its payload field
/// types (empty for a payload-free variant like `Red` or `None`).
#[derive(Debug, Clone, PartialEq)]
pub struct EnumVariantInfo {
    pub enum_name: String,
    pub tag: i32,
    pub field_types: Vec<PlumType>,
    pub values: Vec<ast::Expr>,
}
/// Enum variant name -> its info, e.g. `"True" -> { enum_name: "Bool", tag: 1, field_types: [] }`.
pub type EnumVariants = BTreeMap<String, EnumVariantInfo>;

/// Field names and types for every discriminant enum's shared params (`enum Foo(n: Int) = ...`),
/// keyed by the ENUM's name (not a variant name) — e.g. `"Step" -> [("n", TInt)]`. Field
/// access on a value of this type must load/store at offset `(field_idx + 1) * 8`, NOT
/// `field_idx * 8` like a class — slot 0 is always the variant's tag.
pub type EnumParams = BTreeMap<String, Vec<(String, PlumType)>>;

/// Shared, read-only lookup tables built once from the whole source, threaded through
/// every check/infer call alongside the (mutable, scope-local) `TypeEnv`.
pub struct CheckCtx<'a> {
    pub classes: &'a ClassEnv,
    pub methods: &'a MethodEnv,
    pub enum_variants: &'a EnumVariants,
    pub enum_params: &'a EnumParams,
}

/// Builds the global lookup tables (function/const signatures, class fields, method
/// signatures, enum variants) from a whole source. Shared by `checkSource` and by
/// `plum-wasm-codegen`, which needs the same tables to resolve `self`, field access,
/// and method dispatch during code generation.
pub fn buildGlobalTables(source: &ast::Source) -> (TypeEnv, ClassEnv, MethodEnv, EnumVariants, EnumParams) {
    let mut global_env: TypeEnv = TypeEnv::new();
    let mut classes: ClassEnv = BTreeMap::new();
    let mut methods: MethodEnv = BTreeMap::new();
    let mut enum_variants: EnumVariants = BTreeMap::new();
    let mut enum_params: EnumParams = BTreeMap::new();
    // `Bool`'s variants are built in (see `inferExpr`'s TypeName handling) rather
    // than requiring every source file to redeclare `enum Bool = | True | False`.
    enum_variants.insert("True".to_string(), EnumVariantInfo { enum_name: "Bool".to_string(), tag: 1, field_types: vec![], values: vec![] });
    enum_variants.insert("False".to_string(), EnumVariantInfo { enum_name: "Bool".to_string(), tag: 0, field_types: vec![], values: vec![] });

    // First pass: register class fields and enum variants so later passes can
    // resolve `self.field`, `ClassName(...)`, and bare enum-tag patterns.
    for item in &source.items {
        match item {
            ast::Item::Class(c) => {
                let fields = c.fields.iter()
                    .map(|f| (f.name.clone(), plumTypeFromAst(&f.ty)))
                    .collect();
                classes.insert(c.name.clone(), fields);
            }
            // `Bool` may be re-"declared" (`enum Bool = | True | False`) purely to
            // give it a nesting site for methods (`and`/`or`/`parse`/...) — the
            // language currently has no other way to attach a method to a builtin
            // type (see `libs/std/int.plum`/`float.plum`'s own `type Int =`/`type
            // Float =` for the same pattern). Re-registering its variants here
            // would silently overwrite the hardcoded tags above with whatever
            // order this declaration happens to list them in, flipping every
            // `True`/`False` tag used throughout the rest of the codegen. Skip.
            ast::Item::Enum(e) if e.name == "Bool" => {}
            ast::Item::Enum(e) => {
                let shared_field_types: Vec<PlumType> = e.params.iter()
                    .map(|p| plumTypeFromAst(&p.ty))
                    .collect();
                if !e.params.is_empty() {
                    let params = e.params.iter()
                        .map(|p| (p.name.clone(), plumTypeFromAst(&p.ty)))
                        .collect();
                    enum_params.insert(e.name.clone(), params);
                }
                for (tag, v) in e.variants.iter().enumerate() {
                    let field_types = if !e.params.is_empty() {
                        shared_field_types.clone()
                    } else {
                        v.fields.iter()
                            .map(|f| plumTypeFromAst(&ast::Type { name: f.clone(), generics: vec![] }))
                            .collect()
                    };
                    enum_variants.insert(v.name.clone(), EnumVariantInfo {
                        enum_name: e.name.clone(),
                        tag: tag as i32,
                        field_types,
                        values: v.values.clone(),
                    });
                }
            }
            _ => {}
        }
    }

    // Second pass: register top-level function/method signatures and consts.
    // Methods (`name<Receiver>(...)`) live in `methods`, keyed by receiver type,
    // so a bare call can't accidentally resolve to some other type's method.
    for item in &source.items {
        match item {
            ast::Item::Fn(f) => {
                let param_types: Vec<PlumType> = f.params.iter().map(|p| {
                    match &p.ty {
                        ast::ParamType::Type(t) => plumTypeFromAst(t),
                        ast::ParamType::Variadic(t) => PlumType::TVariadic(Box::new(plumTypeFromAst(t))),
                        ast::ParamType::Fn(params, ret) => {
                            let param_types = params.iter().map(plumTypeFromAst).collect();
                            let ret_ty = ret.as_ref().map(|r| plumTypeFromAst(r)).unwrap_or(PlumType::TUnit);
                            PlumType::TFun(param_types, Box::new(ret_ty))
                        }
                    }
                }).collect();
                let ret = f.returns.as_ref()
                    .map(plumTypeFromAst)
                    .unwrap_or(PlumType::TUnit);
                let fn_ty = PlumType::TFun(param_types, Box::new(ret));
                if let Some(recv) = &f.type_param {
                    methods.insert((recv.clone(), f.name.clone()), fn_ty);
                } else {
                    global_env.insert(f.name.clone(), TypeScheme::mono(fn_ty));
                }
            }
            ast::Item::Const(c) => {
                // Top-level consts are always simple literals in practice; infer
                // directly from the literal kind rather than pulling in the full
                // `inferExpr` (which needs a `CheckCtx` this pass hasn't built yet).
                let ty = match &c.value {
                    ast::Expr::Int(_) => PlumType::TInt,
                    ast::Expr::Float(_) => PlumType::TFloat,
                    ast::Expr::String(_) => PlumType::TStr,
                    _ => PlumType::TVar("_".to_string()),
                };
                global_env.insert(c.name.clone(), TypeScheme::mono(ty));
            }
            _ => {}
        }
    }

    (global_env, classes, methods, enum_variants, enum_params)
}

pub fn checkSource(source: &ast::Source) -> CheckResult<()> {
    let source = monomorphize::monomorphizeSource(source).map_err(|e| vec![CheckError { message: e }])?;
    let mut errors: Vec<CheckError> = Vec::new();
    let (global_env, classes, methods, enum_variants, enum_params) = buildGlobalTables(&source);
    let ctx = CheckCtx { classes: &classes, methods: &methods, enum_variants: &enum_variants, enum_params: &enum_params };

    // A name that is both a class and an enum variant is ambiguous: `Name(...)`
    // could mean either construction, and downstream code (both the checker's
    // `inferExpr` and codegen) consults `enum_variants` first, so the class
    // constructor would be silently shadowed with no diagnostic. Reject it.
    for name in classes.keys() {
        if enum_variants.contains_key(name) {
            errors.push(CheckError {
                message: format!("'{}' is declared as both a class and an enum variant", name),
            });
        }
    }

    // A discriminant enum (`enum Foo(n: Int) = ...`) requires every variant to supply
    // exactly one value per declared param, unified against that param's type. An
    // ordinary (param-less) enum must NOT have variants with values — most likely
    // caused by writing `Some(5)` where the generic-payload form `Some[Int]` was meant.
    for item in &source.items {
        if let ast::Item::Enum(e) = item {
            for v in &e.variants {
                if e.params.is_empty() {
                    if !v.values.is_empty() {
                        errors.push(CheckError {
                            message: format!("enum '{}' variant '{}': has discriminant values but '{}' declares no params", e.name, v.name, e.name),
                        });
                    }
                    continue;
                }
                if v.values.len() != e.params.len() {
                    errors.push(CheckError {
                        message: format!(
                            "enum '{}' variant '{}': expected {} discriminant value(s), got {}",
                            e.name, v.name, e.params.len(), v.values.len()
                        ),
                    });
                    continue;
                }
                for (value, param) in v.values.iter().zip(e.params.iter()) {
                    match inferExpr(value, &global_env, &ctx) {
                        Ok(actual) => {
                            let expected = plumTypeFromAst(&param.ty);
                            if let Err(msg) = unify(&expected, &actual) {
                                errors.push(CheckError {
                                    message: format!("enum '{}' variant '{}': param '{}': {}", e.name, v.name, param.name, msg),
                                });
                            }
                        }
                        Err(msg) => errors.push(CheckError {
                            message: format!("enum '{}' variant '{}': param '{}': {}", e.name, v.name, param.name, msg),
                        }),
                    }
                }
            }
        }
    }

    for item in &source.items {
        if let ast::Item::Fn(f) = item {
            let mut local_errors = checkFn(f, &global_env, &ctx);
            errors.append(&mut local_errors);
        }
    }

    if errors.is_empty() { Ok(()) } else { Err(errors) }
}

fn checkFn(f: &ast::Fn, global_env: &TypeEnv, ctx: &CheckCtx) -> Vec<CheckError> {
    let mut errors = Vec::new();

    // `is_extern` and `body == FnBody::Extern` must agree — either both (a genuine
    // host-backed declaration) or neither (a normal Plum function). A mismatch here
    // is always a parser artifact of a malformed declaration, not user intent, so
    // report it plainly rather than trying to guess which side was "right".
    match (f.is_extern, &f.body) {
        (true, ast::FnBody::Extern) => {
            if f.type_param.is_some() {
                errors.push(CheckError { message: format!("fn '{}': extern functions can't be methods (no receiver)", f.name) });
            }
            return errors;
        }
        (true, _) => {
            errors.push(CheckError { message: format!("fn '{}': extern fn must not have a body", f.name) });
            return errors;
        }
        (false, ast::FnBody::Extern) => {
            errors.push(CheckError { message: format!("fn '{}': missing a body (or mark it `extern`)", f.name) });
            return errors;
        }
        (false, _) => {}
    }

    let mut env = global_env.clone();

    let variadic_positions: Vec<usize> = f.params.iter().enumerate()
        .filter(|(_, p)| matches!(p.ty, ast::ParamType::Variadic(_)))
        .map(|(i, _)| i)
        .collect();
    if variadic_positions.len() > 1 {
        errors.push(CheckError { message: format!("fn '{}': at most one variadic parameter is allowed", f.name) });
    } else if let Some(&pos) = variadic_positions.first() {
        if pos != f.params.len() - 1 {
            errors.push(CheckError { message: format!("fn '{}': a variadic parameter must be last", f.name) });
        }
    }

    // Methods (`name<Receiver>(...)`) get an implicit `self: Receiver` binding.
    if let Some(recv) = &f.type_param {
        let recv_ty = ast::Type { name: recv.clone(), generics: vec![] };
        env.insert("self".to_string(), TypeScheme::mono(plumTypeFromAst(&recv_ty)));
    }

    // Add params to env
    for p in &f.params {
        let ty = match &p.ty {
            ast::ParamType::Type(t) => plumTypeFromAst(t),
            ast::ParamType::Variadic(t) => PlumType::TVariadic(Box::new(plumTypeFromAst(t))),
            ast::ParamType::Fn(params, ret) => {
                let param_types = params.iter().map(plumTypeFromAst).collect();
                let ret_ty = ret.as_ref().map(|r| plumTypeFromAst(r)).unwrap_or(PlumType::TUnit);
                PlumType::TFun(param_types, Box::new(ret_ty))
            }
        };
        env.insert(p.name.clone(), TypeScheme::mono(ty));
    }

    let declared_ret = f.returns.as_ref()
        .map(plumTypeFromAst)
        .unwrap_or(PlumType::TUnit);

    match &f.body {
        ast::FnBody::Expr(e) => {
            match inferExpr(e, &env, ctx) {
                Ok(t) => {
                    if let Err(msg) = unify(&declared_ret, &t) {
                        errors.push(CheckError { message: format!("fn '{}': return type mismatch: {}", f.name, msg) });
                    }
                }
                Err(msg) => errors.push(CheckError { message: format!("fn '{}': {}", f.name, msg) }),
            }
        }
        ast::FnBody::Block(block) => {
            let mut block_errors = checkBlock(block, &mut env, &declared_ret, &f.name, ctx);
            errors.append(&mut block_errors);
            // Check the type of the last expression statement against the declared return type
            if let Some(ast::Stmt::Expr(last_expr)) = block.stmts.last() {
                match inferExpr(last_expr, &env, ctx) {
                    Ok(t) => {
                        if let Err(msg) = unify(&declared_ret, &t) {
                            errors.push(CheckError { message: format!("fn '{}': return type mismatch: {}", f.name, msg) });
                        }
                    }
                    Err(_) => {} // already reported by checkBlock
                }
            }
        }
        // Ruled out (and returned early) by the `is_extern`/body coherence check above.
        ast::FnBody::Extern => unreachable!("extern fn bodies are handled earlier in checkFn"),
    }
    errors
}

fn checkBlock(block: &ast::Block, env: &mut TypeEnv, declared_ret: &PlumType, fn_name: &str, ctx: &CheckCtx) -> Vec<CheckError> {
    let mut errors = Vec::new();
    for stmt in &block.stmts {
        let mut stmt_errors = checkStmt(stmt, env, declared_ret, fn_name, ctx);
        errors.append(&mut stmt_errors);
    }
    errors
}

fn describeTargetObject(expr: &ast::Expr) -> String {
    match expr {
        ast::Expr::Self_ => "self".to_string(),
        ast::Expr::Var(n) => n.clone(),
        ast::Expr::Attribute(a) => {
            if let ast::AttrKind::Field(f) = &a.attr {
                format!("{}.{}", describeTargetObject(&a.object), f)
            } else {
                "<expr>".to_string()
            }
        }
        _ => "<expr>".to_string(),
    }
}

fn checkStmt(stmt: &ast::Stmt, env: &mut TypeEnv, declared_ret: &PlumType, fn_name: &str, ctx: &CheckCtx) -> Vec<CheckError> {
    let mut errors = Vec::new();
    match stmt {
        ast::Stmt::Assign(a) => {
            for (target, value) in a.targets.iter().zip(a.values.iter()) {
                match target {
                    ast::AssignTarget::Var(name) => {
                        let already_declared = lookup(env, name).is_ok();
                        if a.declare && already_declared {
                            errors.push(CheckError {
                                message: format!("fn '{}': assign '{}': already declared — use '=' to reassign, not ':='", fn_name, name),
                            });
                        } else if !a.declare && already_declared {
                            // Reassignment: the new value's type must match what
                            // `name` was already bound to — this is the actual
                            // point of `:=` existing at all (catches `x := 5`
                            // then later `x = 3.14` at compile time instead of
                            // silently miscompiling or silently changing type).
                            match inferExpr(value, env, ctx) {
                                Ok(t) => {
                                    let existing = lookup(env, name).expect("already_declared just confirmed this succeeds");
                                    if let Err(msg) = unify(&existing, &t) {
                                        errors.push(CheckError { message: format!("fn '{}': assign '{}': {}", fn_name, name, msg) });
                                    } else if matches!(existing, PlumType::TVar(_)) {
                                        // `unify` never narrows a `TVar` itself —
                                        // once a concrete type is available, record
                                        // it so later uses of `name` see it too.
                                        env.insert(name.clone(), TypeScheme::mono(t));
                                    }
                                }
                                Err(msg) => errors.push(CheckError { message: format!("fn '{}': assign '{}': {}", fn_name, name, msg) }),
                            }
                        } else {
                            // First declaration — via `:=`, or (for backward
                            // compatibility with every existing bare `x = ...`
                            // first-use) via a plain `=` on a name not yet bound.
                            match inferExpr(value, env, ctx) {
                                Ok(t) => { env.insert(name.clone(), TypeScheme::mono(t)); }
                                Err(msg) => errors.push(CheckError { message: format!("fn '{}': assign '{}': {}", fn_name, name, msg) }),
                            }
                        }
                    }
                    ast::AssignTarget::Field(object, field_name) => {
                        let label = format!("{}.{}", describeTargetObject(object), field_name);
                        match (inferExpr(object, env, ctx), inferExpr(value, env, ctx)) {
                            (Ok(PlumType::TNamed(class_name)), Ok(value_ty)) => {
                                match ctx.classes.get(&class_name).and_then(|fields| {
                                    fields.iter().find(|(n, _)| n == field_name).map(|(_, ty)| ty.clone())
                                }) {
                                    Some(field_ty) => {
                                        if let Err(msg) = unify(&field_ty, &value_ty) {
                                            errors.push(CheckError { message: format!("fn '{}': assign '{}': {}", fn_name, label, msg) });
                                        }
                                    }
                                    None => errors.push(CheckError { message: format!("fn '{}': assign '{}': no field '{}' on class '{}'", fn_name, label, field_name, class_name) }),
                                }
                            }
                            (Ok(other), Ok(_)) => errors.push(CheckError { message: format!("fn '{}': assign '{}': cannot access field on non-class type {}", fn_name, label, other) }),
                            (Err(msg), _) => errors.push(CheckError { message: format!("fn '{}': assign '{}': {}", fn_name, label, msg) }),
                            (_, Err(msg)) => errors.push(CheckError { message: format!("fn '{}': assign '{}': {}", fn_name, label, msg) }),
                        }
                    }
                }
            }
        }
        ast::Stmt::Return(Some(e)) => {
            match inferExpr(e, env, ctx) {
                Ok(t) => {
                    if let Err(msg) = unify(declared_ret, &t) {
                        errors.push(CheckError { message: format!("fn '{}': return type mismatch: {}", fn_name, msg) });
                    }
                }
                Err(msg) => errors.push(CheckError { message: format!("fn '{}': return: {}", fn_name, msg) }),
            }
        }
        ast::Stmt::Return(None) => {
            if let Err(msg) = unify(declared_ret, &PlumType::TUnit) {
                errors.push(CheckError { message: format!("fn '{}': bare return in non-Unit function: {}", fn_name, msg) });
            }
        }
        ast::Stmt::If(if_) => {
            match inferExpr(&if_.condition, env, ctx) {
                Ok(t) => {
                    if let Err(msg) = unify(&PlumType::TBool, &t) {
                        errors.push(CheckError { message: format!("fn '{}': if condition must be Bool: {}", fn_name, msg) });
                    }
                }
                Err(msg) => errors.push(CheckError { message: format!("fn '{}': if condition: {}", fn_name, msg) }),
            }
            errors.append(&mut checkBlock(&if_.body, env, declared_ret, fn_name, ctx));
            for ei in &if_.else_ifs {
                match inferExpr(&ei.condition, env, ctx) {
                    Ok(t) => {
                        if let Err(msg) = unify(&PlumType::TBool, &t) {
                            errors.push(CheckError { message: format!("fn '{}': else if condition must be Bool: {}", fn_name, msg) });
                        }
                    }
                    Err(msg) => errors.push(CheckError { message: format!("fn '{}': else if condition: {}", fn_name, msg) }),
                }
                errors.append(&mut checkBlock(&ei.body, env, declared_ret, fn_name, ctx));
            }
            if let Some(else_block) = &if_.else_ {
                errors.append(&mut checkBlock(else_block, env, declared_ret, fn_name, ctx));
            }
        }
        ast::Stmt::While(w) => {
            match inferExpr(&w.condition, env, ctx) {
                Ok(t) => {
                    if let Err(msg) = unify(&PlumType::TBool, &t) {
                        errors.push(CheckError { message: format!("fn '{}': while condition must be Bool: {}", fn_name, msg) });
                    }
                }
                Err(msg) => errors.push(CheckError { message: format!("fn '{}': while condition: {}", fn_name, msg) }),
            }
            errors.append(&mut checkBlock(&w.body, env, declared_ret, fn_name, ctx));
        }
        ast::Stmt::For(f_stmt) => {
            let iter_ty = inferExpr(&f_stmt.iter, env, ctx);
            let mut inner_env = env.clone();
            match &iter_ty {
                Ok(PlumType::TVariadic(elem)) => {
                    if f_stmt.vars.len() != 1 {
                        errors.push(CheckError { message: format!("fn '{}': for-loop over a variadic param must bind exactly one variable", fn_name) });
                    }
                    for var in &f_stmt.vars {
                        inner_env.insert(var.clone(), TypeScheme::mono((**elem).clone()));
                    }
                }
                Ok(_) => {
                    for var in &f_stmt.vars {
                        inner_env.insert(var.clone(), TypeScheme::mono(PlumType::TInt));
                    }
                }
                Err(msg) => {
                    errors.push(CheckError { message: format!("fn '{}': for iter: {}", fn_name, msg) });
                    for var in &f_stmt.vars {
                        inner_env.insert(var.clone(), TypeScheme::mono(PlumType::TInt));
                    }
                }
            }
            errors.append(&mut checkBlock(&f_stmt.body, &mut inner_env, declared_ret, fn_name, ctx));
        }
        ast::Stmt::Expr(e) => {
            if let Err(msg) = inferExpr(e, env, ctx) {
                errors.push(CheckError { message: format!("fn '{}': {}", fn_name, msg) });
            }
        }
        ast::Stmt::Assert(e) => {
            match inferExpr(e, env, ctx) {
                Ok(t) => {
                    if let Err(msg) = unify(&PlumType::TBool, &t) {
                        errors.push(CheckError { message: format!("fn '{}': assert must be Bool: {}", fn_name, msg) });
                    }
                }
                Err(msg) => errors.push(CheckError { message: format!("fn '{}': assert: {}", fn_name, msg) }),
            }
        }
        ast::Stmt::Match(m) => {
            errors.append(&mut checkMatch(m, env, declared_ret, fn_name, ctx));
        }
        ast::Stmt::Break | ast::Stmt::Continue | ast::Stmt::Todo => {}
    }
    errors
}

fn checkMatch(m: &ast::Match, env: &TypeEnv, declared_ret: &PlumType, fn_name: &str, ctx: &CheckCtx) -> Vec<CheckError> {
    let mut errors = Vec::new();

    let mut subject_types: Vec<PlumType> = Vec::new();
    for s in &m.subjects {
        match inferExpr(s, env, ctx) {
            Ok(t) => subject_types.push(t),
            Err(msg) => errors.push(CheckError { message: format!("fn '{}': match subject: {}", fn_name, msg) }),
        }
    }
    if subject_types.len() != m.subjects.len() {
        return errors; // a subject failed to type — cases can't be checked meaningfully
    }

    for case in &m.cases {
        let mut case_env = env.clone();
        if case.patterns.len() == subject_types.len() {
            for (pat, sty) in case.patterns.iter().zip(subject_types.iter()) {
                if let Err(msg) = checkPattern(pat, sty, &mut case_env, ctx) {
                    errors.push(CheckError { message: format!("fn '{}': match case: {}", fn_name, msg) });
                }
            }
        } else {
            errors.push(CheckError {
                message: format!(
                    "fn '{}': match case has {} pattern(s), expected {}",
                    fn_name, case.patterns.len(), subject_types.len()
                ),
            });
        }
        errors.append(&mut checkBlock(&case.body, &mut case_env, declared_ret, fn_name, ctx));
    }
    errors
}

/// Checks a single case pattern against the type of the subject it matches, binding any
/// new names it introduces into `env`. Constructor-payload sub-patterns (`Some(x)`) bind
/// against that variant's declared field types (see `EnumVariantInfo::field_types`).
fn checkPattern(pat: &ast::CasePattern, subject_ty: &PlumType, env: &mut TypeEnv, ctx: &CheckCtx) -> Result<(), String> {
    match pat {
        ast::CasePattern::Wildcard => Ok(()),
        ast::CasePattern::Int(_) => unify(subject_ty, &PlumType::TInt),
        ast::CasePattern::Float(_) => unify(subject_ty, &PlumType::TFloat),
        ast::CasePattern::String(_) => unify(subject_ty, &PlumType::TStr),
        ast::CasePattern::Name(n) => {
            let is_known_variant = n.chars().next().map(|c| c.is_uppercase()).unwrap_or(false)
                && ctx.enum_variants.contains_key(n);
            if is_known_variant {
                Ok(()) // equality check against a known enum tag, e.g. `True`
            } else {
                env.insert(n.clone(), TypeScheme::mono(subject_ty.clone()));
                Ok(())
            }
        }
        ast::CasePattern::Class { name, fields } => match ctx.enum_variants.get(name) {
            Some(info) => {
                if fields.len() != info.field_types.len() {
                    return Err(format!(
                        "constructor pattern '{}' expects {} field(s), got {}",
                        name, info.field_types.len(), fields.len()
                    ));
                }
                for (f, fty) in fields.iter().zip(info.field_types.iter()) {
                    checkPattern(f, fty, env, ctx)?;
                }
                Ok(())
            }
            // Unmodeled/builtin variant: allow, codegen will catch.
            None => {
                for f in fields {
                    checkPattern(f, &PlumType::TVar("_".to_string()), env, ctx)?;
                }
                Ok(())
            }
        },
    }
}

pub fn inferExpr(expr: &ast::Expr, env: &TypeEnv, ctx: &CheckCtx) -> Result<PlumType, String> {
    match expr {
        ast::Expr::Int(_) => Ok(PlumType::TInt),
        ast::Expr::Float(_) => Ok(PlumType::TFloat),
        ast::Expr::String(_) => Ok(PlumType::TStr),
        ast::Expr::Var(name) => lookup(env, name),
        ast::Expr::Self_ => lookup(env, "self"),
        // A bare `NAME` lexes as a type_identifier whenever it starts with an
        // uppercase letter — which includes every SCREAMING_CASE top-level
        // const (`PI`, `MAX_VALUE`, ...), since there's no separate "constant"
        // token. Prefer a matching const's real type over the enum-variant/
        // unmodeled-type-name fallbacks below.
        ast::Expr::TypeName(n) => match n.as_str() {
            "True" | "False" => Ok(PlumType::TBool),
            _ => match lookup(env, n) {
                Ok(ty) if !matches!(ty, PlumType::TVar(_)) => Ok(ty),
                _ => match ctx.enum_variants.get(n) {
                    Some(info) => Ok(PlumType::TNamed(info.enum_name.clone())),
                    // Unmodeled/builtin type name: allow, codegen will catch.
                    None => Ok(PlumType::TNamed(n.to_string())),
                },
            },
        },
        ast::Expr::Paren(inner) => inferExpr(inner, env, ctx),
        ast::Expr::Closure(cl) => {
            let mut closure_env = env.clone();
            let param_types: Vec<PlumType> = cl.params.iter().map(|p| {
                let t = PlumType::TVar(format!("_closure_{}", p));
                closure_env.insert(p.clone(), TypeScheme::mono(t.clone()));
                t
            }).collect();
            let body_ty = match &cl.body.stmts.last() {
                Some(ast::Stmt::Expr(e)) => inferExpr(e, &closure_env, ctx)?,
                Some(ast::Stmt::Return(Some(e))) => inferExpr(e, &closure_env, ctx)?,
                _ => PlumType::TUnit,
            };
            Ok(PlumType::TFun(param_types, Box::new(body_ty)))
        }
        ast::Expr::Not(inner) => {
            let t = inferExpr(inner, env, ctx)?;
            unify(&PlumType::TBool, &t)?;
            Ok(PlumType::TBool)
        }
        ast::Expr::Unary(u) => inferExpr(&u.operand, env, ctx),
        ast::Expr::Binary(b) => {
            let lt = inferExpr(&b.left, env, ctx)?;
            let rt = inferExpr(&b.right, env, ctx)?;
            unify(&lt, &rt).map_err(|e| format!("binary op: {}", e))?;
            Ok(lt)
        }
        ast::Expr::Bool(b) => {
            let lt = inferExpr(&b.left, env, ctx)?;
            let rt = inferExpr(&b.right, env, ctx)?;
            unify(&PlumType::TBool, &lt).map_err(|e| format!("bool op left: {}", e))?;
            unify(&PlumType::TBool, &rt).map_err(|e| format!("bool op right: {}", e))?;
            Ok(PlumType::TBool)
        }
        ast::Expr::Compare(c) => {
            let lt = inferExpr(&c.left, env, ctx)?;
            let rt = inferExpr(&c.right, env, ctx)?;
            unify(&lt, &rt).map_err(|e| format!("compare op: {}", e))?;
            Ok(PlumType::TBool)
        }
        ast::Expr::Ternary(t) => {
            let ct = inferExpr(&t.condition, env, ctx)?;
            unify(&PlumType::TBool, &ct).map_err(|e| format!("ternary condition: {}", e))?;
            let tt = inferExpr(&t.then, env, ctx)?;
            let et = inferExpr(&t.else_, env, ctx)?;
            unify(&tt, &et).map_err(|e| format!("ternary branches: {}", e))?;
            Ok(tt)
        }
        ast::Expr::FnCall(call) => {
            // `Int(x)`/`Float(x)`/`Byte(x)` are builtin numeric conversions, not
            // ordinary calls — handled here so `y = Float(x)` unifies against
            // `TFloat` rather than falling through to the permissive
            // "unknown fn" case.
            if (call.name == "Int" || call.name == "Float" || call.name == "Byte") && call.args.len() == 1 {
                let arg_expr = match &call.args[0] {
                    ast::Arg::Positional(e) => e,
                    ast::Arg::Keyword { value, .. } => value,
                    ast::Arg::Pair { value, .. } => value,
                };
                let actual = inferExpr(arg_expr, env, ctx)?;
                return match (call.name.as_str(), &actual) {
                    ("Float", PlumType::TInt) => Ok(PlumType::TFloat),
                    ("Int", PlumType::TFloat) => Ok(PlumType::TInt),
                    ("Float", PlumType::TFloat) | ("Int", PlumType::TInt) => Ok(actual),
                    ("Byte", PlumType::TInt) | ("Byte", PlumType::TByte) => Ok(PlumType::TByte),
                    ("Int", PlumType::TByte) => Ok(PlumType::TInt),
                    (name, other) => Err(format!("call '{}': cannot convert {} to {}", name, other, name)),
                };
            }
            if let Some(info) = ctx.enum_variants.get(&call.name) {
                if call.args.len() != info.field_types.len() {
                    return Err(format!(
                        "variant '{}': expected {} arg(s), got {}",
                        call.name, info.field_types.len(), call.args.len()
                    ));
                }
                for (i, (arg, expected)) in call.args.iter().zip(info.field_types.iter()).enumerate() {
                    let arg_expr = match arg {
                        ast::Arg::Positional(e) => e,
                        ast::Arg::Keyword { value, .. } => value,
                        ast::Arg::Pair { value, .. } => value,
                    };
                    let actual = inferExpr(arg_expr, env, ctx)?;
                    unify(expected, &actual).map_err(|e| format!("variant '{}' arg {}: {}", call.name, i, e))?;
                }
                return Ok(PlumType::TNamed(info.enum_name.clone()));
            }
            match lookup(env, &call.name) {
                Ok(PlumType::TFun(param_types, ret)) => {
                    match param_types.last() {
                        Some(PlumType::TVariadic(elem)) => {
                            let fixed = &param_types[..param_types.len() - 1];
                            if call.args.len() < fixed.len() {
                                return Err(format!("call '{}': expected at least {} arg(s), got {}", call.name, fixed.len(), call.args.len()));
                            }
                            for (i, (arg, expected)) in call.args.iter().zip(fixed.iter()).enumerate() {
                                let arg_expr = match arg {
                                    ast::Arg::Positional(e) => e,
                                    ast::Arg::Keyword { value, .. } => value,
                                    ast::Arg::Pair { value, .. } => value,
                                };
                                let actual = inferExpr(arg_expr, env, ctx)?;
                                unify(expected, &actual).map_err(|e| format!("call '{}' arg {}: {}", call.name, i, e))?;
                            }
                            for (i, arg) in call.args.iter().enumerate().skip(fixed.len()) {
                                let arg_expr = match arg {
                                    ast::Arg::Positional(e) => e,
                                    ast::Arg::Keyword { value, .. } => value,
                                    ast::Arg::Pair { value, .. } => value,
                                };
                                let actual = inferExpr(arg_expr, env, ctx)?;
                                unify(elem, &actual).map_err(|e| format!("call '{}' variadic arg {}: {}", call.name, i, e))?;
                            }
                            Ok(*ret)
                        }
                        _ => {
                            if call.args.len() != param_types.len() {
                                return Err(format!("call '{}': expected {} args, got {}", call.name, param_types.len(), call.args.len()));
                            }
                            for (i, (arg, expected)) in call.args.iter().zip(param_types.iter()).enumerate() {
                                let arg_expr = match arg {
                                    ast::Arg::Positional(e) => e,
                                    ast::Arg::Keyword { value, .. } => value,
                                    ast::Arg::Pair { value, .. } => value,
                                };
                                let actual = inferExpr(arg_expr, env, ctx)?;
                                unify(expected, &actual).map_err(|e| format!("call '{}' arg {}: {}", call.name, i, e))?;
                            }
                            Ok(*ret)
                        }
                    }
                }
                Ok(_) => Err(format!("'{}' is not a function", call.name)),
                Err(_) => Ok(PlumType::TVar("_".to_string())), // unknown fn: allow, codegen will catch
            }
        }
        ast::Expr::ClassCall(call) => {
            match ctx.classes.get(&call.type_name) {
                Some(fields) => {
                    for fa in &call.fields {
                        match fields.iter().find(|(n, _)| n == &fa.name) {
                            Some((_, expected)) => {
                                let actual = inferExpr(&fa.value, env, ctx)?;
                                unify(expected, &actual)
                                    .map_err(|e| format!("class '{}' field '{}': {}", call.type_name, fa.name, e))?;
                            }
                            None => return Err(format!("unknown field '{}' on class '{}'", fa.name, call.type_name)),
                        }
                    }
                    Ok(PlumType::TNamed(call.type_name.clone()))
                }
                // Unmodeled (e.g. builtin/std) type: allow, codegen will catch.
                None => Ok(PlumType::TNamed(call.type_name.clone())),
            }
        }
        ast::Expr::Attribute(attr) => {
            let obj_ty = inferExpr(&attr.object, env, ctx)?;
            match &attr.attr {
                ast::AttrKind::Field(field_name) => match &obj_ty {
                    PlumType::TNamed(class_name) => match ctx.classes.get(class_name) {
                        Some(fields) => fields.iter()
                            .find(|(n, _)| n == field_name)
                            .map(|(_, t)| t.clone())
                            .ok_or_else(|| format!("no field '{}' on type '{}'", field_name, class_name)),
                        None => match ctx.enum_params.get(class_name) {
                            Some(params) => params.iter()
                                .find(|(n, _)| n == field_name)
                                .map(|(_, t)| t.clone())
                                .ok_or_else(|| format!("no field '{}' on type '{}'", field_name, class_name)),
                            // Unmodeled type: allow, codegen will catch.
                            None => Ok(PlumType::TVar("_".to_string())),
                        },
                    },
                    _ => Err(format!("cannot access field '{}' on non-class type {}", field_name, obj_ty)),
                },
                ast::AttrKind::Method(call) => match methodReceiverName(&obj_ty) {
                    Some(class_name) => match ctx.methods.get(&(class_name.clone(), call.name.clone())) {
                        Some(PlumType::TFun(param_types, ret)) => {
                            match param_types.last() {
                                Some(PlumType::TVariadic(elem)) => {
                                    let fixed = &param_types[..param_types.len() - 1];
                                    if call.args.len() < fixed.len() {
                                        return Err(format!(
                                            "method '{}.{}': expected at least {} arg(s), got {}",
                                            class_name, call.name, fixed.len(), call.args.len()
                                        ));
                                    }
                                    for (i, (arg, expected)) in call.args.iter().zip(fixed.iter()).enumerate() {
                                        let arg_expr = match arg {
                                            ast::Arg::Positional(e) => e,
                                            ast::Arg::Keyword { value, .. } => value,
                                            ast::Arg::Pair { value, .. } => value,
                                        };
                                        let actual = inferExpr(arg_expr, env, ctx)?;
                                        unify(expected, &actual)
                                            .map_err(|e| format!("method '{}.{}' arg {}: {}", class_name, call.name, i, e))?;
                                    }
                                    for (i, arg) in call.args.iter().enumerate().skip(fixed.len()) {
                                        let arg_expr = match arg {
                                            ast::Arg::Positional(e) => e,
                                            ast::Arg::Keyword { value, .. } => value,
                                            ast::Arg::Pair { value, .. } => value,
                                        };
                                        let actual = inferExpr(arg_expr, env, ctx)?;
                                        unify(elem, &actual)
                                            .map_err(|e| format!("method '{}.{}' variadic arg {}: {}", class_name, call.name, i, e))?;
                                    }
                                    Ok(*ret.clone())
                                }
                                _ => {
                                    if call.args.len() != param_types.len() {
                                        return Err(format!(
                                            "method '{}.{}': expected {} args, got {}",
                                            class_name, call.name, param_types.len(), call.args.len()
                                        ));
                                    }
                                    for (i, (arg, expected)) in call.args.iter().zip(param_types.iter()).enumerate() {
                                        let arg_expr = match arg {
                                            ast::Arg::Positional(e) => e,
                                            ast::Arg::Keyword { value, .. } => value,
                                            ast::Arg::Pair { value, .. } => value,
                                        };
                                        let actual = inferExpr(arg_expr, env, ctx)?;
                                        unify(expected, &actual)
                                            .map_err(|e| format!("method '{}.{}' arg {}: {}", class_name, call.name, i, e))?;
                                    }
                                    Ok(*ret.clone())
                                }
                            }
                        }
                        // Unmodeled method (e.g. builtin/std): allow, codegen will catch.
                        _ => Ok(PlumType::TVar("_".to_string())),
                    },
                    None => Ok(PlumType::TVar("_".to_string())),
                },
            }
        }
    }
}