//! Maat type system.
//!
//! Implements [`Hindley-Milner`] type inference (Algorithm W), unification,
//! constant folding (via `maat_ast`), and compile-time type checking.
//! Sits between macro expansion and codegen in the compilation pipeline.
//!
//! [`Hindley-Milner`]: https://en.wikipedia.org/wiki/Hindley%E2%80%93Milner_type_system

#![forbid(unsafe_code)]

mod env;
mod resolve;
mod ty;
mod unify;

use std::rc::Rc;

pub use env::TypeEnv;
use maat_ast::*;
use maat_errors::{
    MissingTraitMethodError, TraitMethodSignatureMismatchError, TypeError, TypeErrorKind,
};
use maat_field::MODULUS as GOLDILOCKS_MODULUS;
use maat_span::Span;
pub use resolve::resolve_type_expr;
pub use ty::{
    EnumDef, FnType, ImplDef, MethodSig, StructDef, TraitDef, Type, TypeScheme, TypeVarId,
    VariantDef, VariantKind,
};
use unify::{Substitution, UnifyError};

pub struct TypeChecker {
    env: TypeEnv,
    subst: Substitution,
    errors: Vec<TypeError>,
    bitwise_ops: Vec<BitwiseOp>,
}

struct BitwiseOp {
    operator: String,
    lhs_ty: Type,
    rhs_ty: Type,
    lhs_span: Span,
    rhs_span: Span,
}

impl Default for TypeChecker {
    fn default() -> Self {
        Self::new()
    }
}

impl TypeChecker {
    pub fn new() -> Self {
        let mut env = TypeEnv::new();
        env.register_builtins();
        Self {
            env,
            subst: Substitution::new(),
            errors: Vec::new(),
            bitwise_ops: Vec::new(),
        }
    }

    pub fn env(&self) -> &TypeEnv {
        &self.env
    }

    pub fn env_mut(&mut self) -> &mut TypeEnv {
        &mut self.env
    }

    pub fn check_program(mut self, program: &mut Program) -> Vec<TypeError> {
        self.check_program_mut(program);
        self.errors
    }

    pub fn check_program_mut(&mut self, program: &mut Program) {
        for stmt in &program.statements {
            match stmt {
                Stmt::StructDecl(decl) => self.register_struct(decl),
                Stmt::EnumDecl(decl) => self.register_enum(decl),
                Stmt::TraitDecl(decl) => self.register_trait(decl),
                _ => {}
            }
        }
        for stmt in &mut program.statements {
            self.check_statement(stmt);
        }
        self.subst.resolve_inferred_literals(program);
        self.validate_bitwise_ops();
        self.validate_u64_literals(program);
    }

    pub fn errors(&self) -> &[TypeError] {
        &self.errors
    }

    pub fn drain_errors(&mut self) -> Vec<TypeError> {
        std::mem::take(&mut self.errors)
    }

    fn register_struct(&mut self, decl: &StructDecl) {
        if self.env.lookup_struct(&decl.name).is_some() {
            self.errors
                .push(TypeErrorKind::DuplicateType(decl.name.clone()).at(decl.span));
            return;
        }
        let generic_params = decl
            .generic_params
            .iter()
            .map(|g| g.name.clone())
            .collect::<Vec<String>>();
        let fields = decl
            .fields
            .iter()
            .map(|f| {
                (
                    f.name.clone(),
                    self.resolve_field_type(&f.ty, &generic_params),
                )
            })
            .collect();
        self.env.register_struct(StructDef {
            name: decl.name.clone(),
            generic_params,
            fields,
        });
    }

    fn register_enum(&mut self, decl: &EnumDecl) {
        if self.env.lookup_enum(&decl.name).is_some() {
            self.errors
                .push(TypeErrorKind::DuplicateType(decl.name.clone()).at(decl.span));
            return;
        }
        let generic_params = decl
            .generic_params
            .iter()
            .map(|g| g.name.clone())
            .collect::<Vec<String>>();
        let variants = decl
            .variants
            .iter()
            .map(|v| VariantDef {
                name: v.name.clone(),
                kind: match &v.kind {
                    EnumVariantKind::Unit => VariantKind::Unit,
                    EnumVariantKind::Tuple(tys) => VariantKind::Tuple(
                        tys.iter()
                            .map(|t| self.resolve_field_type(t, &generic_params))
                            .collect(),
                    ),
                    EnumVariantKind::Struct(fields) => VariantKind::Struct(
                        fields
                            .iter()
                            .map(|f| {
                                (
                                    f.name.clone(),
                                    self.resolve_field_type(&f.ty, &generic_params),
                                )
                            })
                            .collect(),
                    ),
                },
            })
            .collect();
        self.env.register_enum(EnumDef {
            name: decl.name.clone(),
            generic_params,
            variants,
        });
    }

    fn register_trait(&mut self, decl: &TraitDecl) {
        if self.env.lookup_trait(&decl.name).is_some() {
            self.errors
                .push(TypeErrorKind::DuplicateType(decl.name.clone()).at(decl.span));
            return;
        }
        let generic_params = decl
            .generic_params
            .iter()
            .map(|g| g.name.clone())
            .collect::<Vec<String>>();
        let methods = decl
            .methods
            .iter()
            .map(|m| {
                let takes_self = m.params.first().is_some_and(|p| p.name == "self");
                let params = m
                    .params
                    .iter()
                    .filter(|p| p.name != "self")
                    .map(|p| {
                        p.type_expr
                            .as_ref()
                            .map(|t| self.resolve_field_type(t, &generic_params))
                            .unwrap_or(Type::Var(0))
                    })
                    .collect();
                let ret = m
                    .return_type
                    .as_ref()
                    .map(|t| self.resolve_field_type(t, &generic_params))
                    .unwrap_or(Type::Unit);
                MethodSig {
                    name: m.name.clone(),
                    params,
                    ret,
                    has_default: m.default_body.is_some(),
                    takes_self,
                }
            })
            .collect();
        self.env.register_trait(TraitDef {
            name: decl.name.clone(),
            generic_params,
            methods,
        });
    }

    fn resolve_field_type(&mut self, ty: &TypeExpr, generic_params: &[String]) -> Type {
        match ty {
            TypeExpr::Named(named) if generic_params.contains(&named.name) => {
                Type::Generic(Rc::from(named.name.as_str()), vec![])
            }
            _ => self.env.resolve_type(ty),
        }
    }

    fn check_statement(&mut self, stmt: &mut Stmt) {
        match stmt {
            Stmt::Let(let_stmt) => self.check_let(let_stmt),
            Stmt::ReAssign(assign_stmt) => {
                self.infer_expression(&mut assign_stmt.value);
            }
            Stmt::Return(ret_stmt) => {
                self.infer_expression(&mut ret_stmt.value);
            }
            Stmt::Expr(expr_stmt) => {
                self.infer_expression(&mut expr_stmt.value);
            }
            Stmt::Block(block) => self.check_block(block),
            Stmt::FuncDef(fn_item) => self.check_fn_item(fn_item),
            Stmt::Loop(loop_stmt) => self.check_block(&mut loop_stmt.body),
            Stmt::While(while_stmt) => {
                let cond_ty = self.infer_expression(&mut while_stmt.condition);
                let cond_resolved = self.subst.apply(&cond_ty);
                self.require_bool(&cond_resolved, while_stmt.condition.span());
                self.check_block(&mut while_stmt.body);
            }
            Stmt::For(for_stmt) => {
                let iter_ty = self.infer_expression(&mut for_stmt.iterable);
                let resolved = self.subst.apply(&iter_ty);
                let elem_ty = match resolved {
                    Type::Vector(elem) => *elem,
                    Type::Range(elem) => *elem,
                    Type::Map(k, v) => Type::Tuple(vec![*k, *v]),
                    Type::Var(_) => self.env.fresh_var(),
                    _ => {
                        self.errors.push(
                            TypeErrorKind::Mismatch {
                                expected: "[T], Range<T>, or Map<K, V>".to_string(),
                                found: resolved.to_string(),
                            }
                            .at(for_stmt.iterable.span()),
                        );
                        self.env.fresh_var()
                    }
                };
                self.env.push_scope();
                if let Some(pattern) = &for_stmt.pattern {
                    self.check_pattern(pattern, &elem_ty);
                } else {
                    self.env.define_var(&for_stmt.ident, elem_ty);
                }
                self.check_block(&mut for_stmt.body);
                self.env.pop_scope();
            }
            // Struct/Enum/Trait declarations were registered in [self.check_program] (pass 1); nothing to do here.
            Stmt::StructDecl(_) | Stmt::EnumDecl(_) | Stmt::TraitDecl(_) => {}
            Stmt::ImplBlock(impl_block) => self.check_impl_block(impl_block),
            // Module declarations (`mod foo;` / `mod foo { ... }`) and import
            // statements (`use foo::bar;`) are resolved by the module orchestrator
            // before per-module type checking runs. No action needed here.
            Stmt::Use(_) | Stmt::Mod(_) => {}
        }
    }

    fn check_let(&mut self, let_stmt: &mut LetStmt) {
        let inferred = self.infer_expression(&mut let_stmt.value);
        let ty = if let Some(ann) = &let_stmt.type_annotation {
            let expected = self.env.resolve_type(ann);
            self.check_literal_range(&let_stmt.value, &expected, let_stmt.span);

            let is_coercible_literal = expected.is_integer()
                && inferred.is_integer()
                && let_stmt.value.is_integer_literal();
            if is_coercible_literal {
                expected.coerce_literal(&mut let_stmt.value);
            } else if self.try_coerce_vector_to_array(
                &expected,
                &inferred,
                &mut let_stmt.value,
                let_stmt.span,
            ) {
                // Vector literal coerced to fixed-size array; no further unification needed.
            } else if let Err(e) = self.subst.unify(&expected, &inferred) {
                self.report_unify_error(e, let_stmt.span);
            }
            expected
        } else {
            inferred
        };
        if let Some(pattern) = &let_stmt.pattern {
            self.check_pattern(pattern, &ty);
        } else {
            let scheme = self.env.generalize(&ty, &self.subst);
            self.env.define_scheme(&let_stmt.ident, scheme);
        }
    }

    fn check_block(&mut self, block: &mut BlockStmt) {
        self.env.push_scope();
        for stmt in &mut block.statements {
            self.check_statement(stmt);
        }
        self.env.pop_scope();
    }

    fn check_fn_item(&mut self, fn_item: &mut FuncDef) {
        let fn_ty = self.infer_fn_body(
            &fn_item.generic_params,
            &fn_item.params,
            &fn_item.return_type,
            &mut fn_item.body,
            fn_item.span,
        );
        let scheme = self.env.generalize(&fn_ty, &self.subst);
        self.env.define_scheme(&fn_item.name, scheme);
    }

    fn infer_fn_body(
        &mut self,
        generic_params: &[GenericParam],
        params: &[TypedParam],
        return_type: &Option<TypeExpr>,
        body: &mut BlockStmt,
        span: Span,
    ) -> Type {
        self.env.push_scope();
        for gp in generic_params {
            let var = self.env.fresh_var();
            self.env.define_var(&gp.name, var);
        }
        let param_types: Vec<Type> = params
            .iter()
            .map(|p| {
                let ty = p
                    .type_expr
                    .as_ref()
                    .map(|te| self.env.resolve_type(te))
                    .unwrap_or_else(|| self.env.fresh_var());
                self.env.define_var(&p.name, ty.clone());
                ty
            })
            .collect();

        let body_ty = self.infer_block(body);
        let ret_ty = return_type
            .as_ref()
            .map(|te| self.env.resolve_type(te))
            .unwrap_or_else(|| body_ty.clone());
        if return_type.is_some()
            && let Err(e) = self.subst.unify(&ret_ty, &body_ty)
        {
            self.report_unify_error(e, span);
        }
        self.env.pop_scope();
        Type::Function(FnType {
            params: param_types,
            ret: Box::new(self.subst.apply(&ret_ty)),
        })
    }

    fn check_impl_block(&mut self, impl_block: &mut ImplBlock) {
        let self_type = self.env.resolve_type(&impl_block.self_type);
        if let Type::Generic(ref name, _) = self_type {
            self.errors
                .push(TypeErrorKind::UnknownType(name.to_string()).at(impl_block.span));
            return;
        }
        let trait_name = impl_block.trait_name.as_ref().and_then(|te| match te {
            TypeExpr::Named(n) => Some(n.name.clone()),
            TypeExpr::Generic(name, _, _) => Some(name.clone()),
            _ => None,
        });
        if let Some(ref name) = trait_name
            && self.env.lookup_trait(name).is_none()
        {
            self.errors
                .push(TypeErrorKind::UnknownTrait(name.clone()).at(impl_block.span));
        }
        let mut method_types = Vec::new();
        for method in &mut impl_block.methods {
            let has_self = method.params.first().is_some_and(|p| p.name == "self");
            let non_self_params = method
                .params
                .iter()
                .filter(|p| p.name != "self")
                .cloned()
                .collect::<Vec<TypedParam>>();
            self.env.push_scope();
            if has_self {
                self.env.define_var("self", self_type.clone());
            }
            let fn_ty = self.infer_fn_body(
                &method.generic_params,
                &non_self_params,
                &method.return_type,
                &mut method.body,
                method.span,
            );
            self.env.pop_scope();
            method_types.push((method.name.clone(), fn_ty));
        }
        if let Some(ref name) = trait_name {
            self.verify_trait_satisfaction(name, &self_type, &method_types, impl_block.span);
        }
        self.env.register_impl(ImplDef {
            self_type,
            trait_name,
            methods: method_types,
        });
    }

    fn verify_trait_satisfaction(
        &mut self,
        trait_name: &str,
        self_type: &Type,
        methods: &[(String, Type)],
        span: Span,
    ) {
        let required = self
            .env
            .required_trait_methods(trait_name)
            .into_iter()
            .cloned()
            .collect::<Vec<MethodSig>>();
        for sig in &required {
            match methods.iter().find(|(name, _)| name == &sig.name) {
                None => {
                    self.errors.push(
                        TypeErrorKind::MissingTraitMethod(Box::new(MissingTraitMethodError {
                            trait_name: trait_name.to_string(),
                            self_type: self_type.to_string(),
                            method: sig.name.clone(),
                        }))
                        .at(span),
                    );
                }
                Some((_, impl_ty)) => {
                    let expected = Type::Function(FnType {
                        params: sig.params.clone(),
                        ret: Box::new(sig.ret.clone()),
                    });
                    if self.subst.unify(&expected, impl_ty).is_err() {
                        self.errors.push(
                            TypeErrorKind::TraitMethodSignatureMismatch(Box::new(
                                TraitMethodSignatureMismatchError {
                                    trait_name: trait_name.to_string(),
                                    self_type: self_type.to_string(),
                                    method: sig.name.clone(),
                                    expected: expected.to_string(),
                                    found: impl_ty.to_string(),
                                },
                            ))
                            .at(span),
                        );
                    }
                }
            }
        }
    }

    fn infer_expression(&mut self, expr: &mut Expr) -> Type {
        match expr {
            Expr::Number(lit) => match lit.kind {
                NumKind::I8 => Type::I8,
                NumKind::I16 => Type::I16,
                NumKind::I32 => Type::I32,
                NumKind::I64 => Type::I64,
                NumKind::I128 => Type::I128,
                NumKind::Isize => Type::Isize,
                NumKind::U8 => Type::U8,
                NumKind::U16 => Type::U16,
                NumKind::U32 => Type::U32,
                NumKind::U64 => Type::U64,
                NumKind::U128 => Type::U128,
                NumKind::Usize => Type::Usize,
                NumKind::Fe => Type::Felt,
                NumKind::Int { ref mut type_var } => {
                    let ty = self.env.fresh_int_var();
                    if let Type::IntVar(id) = ty {
                        *type_var = id;
                    }
                    ty
                }
            },
            Expr::Bool(_) => Type::Bool,
            Expr::Char(_) => Type::Char,
            Expr::Str(_) => Type::Str,
            Expr::Ident(ident) => self
                .env
                .instantiate(&ident.value, &self.subst)
                .or_else(|| self.resolve_bare_variant(&ident.value))
                .unwrap_or_else(|| self.env.fresh_var()),
            Expr::Vector(vector) => self.infer_vector(vector),
            Expr::Map(map) => self.infer_map(map),
            Expr::Index(idx) => self.infer_index(idx),
            Expr::Prefix(prefix) => self.infer_prefix(prefix),
            Expr::Infix(infix) => self.check_infix(infix),
            Expr::Cond(cond) => self.infer_conditional(cond),
            Expr::Lambda(lambda) => self.infer_fn_body(
                &lambda.generic_params,
                &lambda.params,
                &lambda.return_type,
                &mut lambda.body,
                lambda.span,
            ),
            Expr::Call(call) => self.check_call_expr(call),
            Expr::MacroCall(mc) => self.check_macro_call(mc),
            Expr::Cast(cast) => {
                self.infer_expression(&mut cast.expr);
                Type::from_cast_target(&cast.target)
            }
            Expr::Break(break_expr) => {
                if let Some(val) = &mut break_expr.value {
                    self.infer_expression(val);
                }
                Type::Never
            }
            Expr::Continue(_) => Type::Never,
            Expr::MacroLit(_) => self.env.fresh_var(),
            Expr::Tuple(tuple) => {
                let elem_types = tuple
                    .elements
                    .iter_mut()
                    .map(|e| self.infer_expression(e))
                    .collect();
                Type::Tuple(elem_types)
            }
            Expr::Try(try_expr) => self.check_try_expr(try_expr),
            Expr::Match(match_expr) => self.check_match_expr(match_expr),
            Expr::FieldAccess(field_access) => self.check_field_access(field_access),
            Expr::MethodCall(method_call) => self.check_method_call(method_call),
            Expr::StructLit(struct_lit) => self.check_struct_literal(struct_lit),
            Expr::PathExpr(path_expr) => self.check_path_expr(path_expr),
            Expr::Range(range) => self.infer_range(range),
            Expr::Array(arr) => self.infer_array(arr),
        }
    }

    fn infer_vector(&mut self, vector: &mut Vector) -> Type {
        if vector.elements.is_empty() {
            let elem = self.env.fresh_var();
            Type::Vector(Box::new(elem))
        } else {
            let first = self.infer_expression(&mut vector.elements[0]);
            for elem in &mut vector.elements[1..] {
                let elem_ty = self.infer_expression(elem);
                if let Err(e) = self.subst.unify(&first, &elem_ty) {
                    self.report_unify_error(e, elem.span());
                }
            }
            Type::Vector(Box::new(self.subst.apply(&first)))
        }
    }

    fn infer_map(&mut self, map: &mut MapLit) -> Type {
        if map.pairs.is_empty() {
            let k = self.env.fresh_var();
            let v = self.env.fresh_var();
            Type::Map(Box::new(k), Box::new(v))
        } else {
            let (first_k, first_v) = {
                let k = self.infer_expression(&mut map.pairs[0].0);
                let v = self.infer_expression(&mut map.pairs[0].1);
                (k, v)
            };
            for (k_expr, v_expr) in &mut map.pairs[1..] {
                let k = self.infer_expression(k_expr);
                let v = self.infer_expression(v_expr);
                if let Err(e) = self.subst.unify(&first_k, &k) {
                    self.report_unify_error(e, k_expr.span());
                }
                if let Err(e) = self.subst.unify(&first_v, &v) {
                    self.report_unify_error(e, v_expr.span());
                }
            }
            Type::Map(
                Box::new(self.subst.apply(&first_k)),
                Box::new(self.subst.apply(&first_v)),
            )
        }
    }

    fn infer_index(&mut self, idx: &mut IndexExpr) -> Type {
        let collection = self.infer_expression(&mut idx.expr);
        let _index_ty = self.infer_expression(&mut idx.index);
        let resolved = self.subst.apply(&collection);
        match resolved {
            Type::Array(elem, n) => {
                idx.array_len = Some(n);
                *elem
            }
            Type::Vector(elem) => *elem,
            Type::Map(_, v) => *v,
            _ => self.env.fresh_var(),
        }
    }

    fn infer_prefix(&mut self, prefix: &mut PrefixExpr) -> Type {
        let operand_ty = self.infer_expression(&mut prefix.operand);
        let resolved = self.subst.apply(&operand_ty);
        match prefix.operator.as_str() {
            "!" => {
                self.require_bool(&resolved, prefix.span);
                Type::Bool
            }
            "-" => {
                if !resolved.is_integer() && !matches!(resolved, Type::Var(_)) {
                    self.errors.push(
                        TypeErrorKind::Mismatch {
                            expected: "numeric".to_string(),
                            found: resolved.to_string(),
                        }
                        .at(prefix.span),
                    );
                }
                resolved
            }
            _ => self.env.fresh_var(),
        }
    }

    fn infer_conditional(&mut self, cond: &mut CondExpr) -> Type {
        let cond_ty = self.infer_expression(&mut cond.condition);
        let cond_resolved = self.subst.apply(&cond_ty);
        self.require_bool(&cond_resolved, cond.condition.span());

        self.env.push_scope();
        let cons_ty = self.infer_block(&mut cond.consequence);
        self.env.pop_scope();

        if let Some(alt) = &mut cond.alternative {
            self.env.push_scope();
            let alt_ty = self.infer_block(alt);
            self.env.pop_scope();
            if let Err(e) = self.subst.unify(&cons_ty, &alt_ty) {
                self.report_unify_error(e, cond.span);
            }
            self.subst.apply(&cons_ty)
        } else {
            cons_ty
        }
    }

    fn infer_range(&mut self, range: &mut RangeExpr) -> Type {
        let start_ty = self.infer_expression(&mut range.start);
        let end_ty = self.infer_expression(&mut range.end);
        let start_resolved = self.subst.apply(&start_ty);
        let end_resolved = self.subst.apply(&end_ty);
        if let Err(e) = self.subst.unify(&start_resolved, &end_resolved) {
            self.report_unify_error(e, range.span);
        }
        let elem_ty = self.subst.apply(&start_resolved);
        if !elem_ty.is_integer() && !matches!(elem_ty, Type::Var(_) | Type::IntVar(_)) {
            self.errors.push(
                TypeErrorKind::Mismatch {
                    expected: "integer".to_string(),
                    found: elem_ty.to_string(),
                }
                .at(range.span),
            );
        }
        range.kind = Some(elem_ty.to_number_kind());
        Type::Range(Box::new(elem_ty))
    }

    fn infer_array(&mut self, arr: &mut ArrayLit) -> Type {
        if arr.elements.is_empty() {
            let elem = self.env.fresh_var();
            return Type::Array(Box::new(elem), 0);
        }
        let first = self.infer_expression(&mut arr.elements[0]);
        for elem in &mut arr.elements[1..] {
            let elem_ty = self.infer_expression(elem);
            if let Err(e) = self.subst.unify(&first, &elem_ty) {
                self.report_unify_error(e, elem.span());
            }
        }
        Type::Array(Box::new(self.subst.apply(&first)), arr.elements.len())
    }

    fn try_coerce_vector_to_array(
        &mut self,
        expected: &Type,
        inferred: &Type,
        value: &mut Expr,
        span: Span,
    ) -> bool {
        let (elem_ty, n) = match expected {
            Type::Array(elem, n) => (elem.as_ref(), *n),
            _ => return false,
        };
        let vec_elem = match inferred {
            Type::Vector(elem) => elem.as_ref(),
            _ => return false,
        };
        if let Err(e) = self.subst.unify(elem_ty, vec_elem) {
            self.report_unify_error(e, span);
        }
        let actual_len = match value {
            Expr::Vector(v) => v.elements.len(),
            _ => return false,
        };
        if actual_len != n {
            self.errors.push(
                TypeErrorKind::Mismatch {
                    expected: format!("[_; {n}]"),
                    found: format!("[_; {actual_len}]"),
                }
                .at(span),
            );
        }
        // Rewrite the AST: Expr::Vector --> Expr::Array
        let expr = Expr::Bool(BoolLit {
            value: false,
            span: Span::ZERO,
        });
        if let Expr::Vector(vec_lit) = std::mem::replace(value, expr) {
            *value = Expr::Array(ArrayLit {
                elements: vec_lit.elements,
                span: vec_lit.span,
            });
        }
        true
    }

    fn check_call_expr(&mut self, call: &mut CallExpr) -> Type {
        let func_ty = self.infer_expression(&mut call.function);
        let resolved = self.subst.apply(&func_ty);
        let arg_types = call
            .arguments
            .iter_mut()
            .map(|a| self.infer_expression(a))
            .collect::<Vec<Type>>();
        match resolved {
            Type::Function(fn_ty) => {
                if fn_ty.params.len() != arg_types.len() {
                    self.errors.push(
                        TypeErrorKind::WrongArity {
                            expected: fn_ty.params.len(),
                            found: arg_types.len(),
                        }
                        .at(call.span),
                    );
                } else {
                    for (param, arg) in fn_ty.params.iter().zip(arg_types.iter()) {
                        let p = self.subst.apply(param);
                        let a = self.subst.apply(arg);
                        if let Err(e) = self.subst.unify(&p, &a) {
                            self.report_unify_error(e, call.span);
                        }
                    }
                }
                let ret = self.subst.apply(&fn_ty.ret);
                self.validate_from_conversion(call, &arg_types, &ret);
                ret
            }
            Type::Var(_) => {
                let ret = self.env.fresh_var();
                let expected_fn = Type::Function(FnType {
                    params: arg_types,
                    ret: Box::new(ret.clone()),
                });
                if let Err(e) = self.subst.unify(&resolved, &expected_fn) {
                    self.report_unify_error(e, call.span);
                }
                self.subst.apply(&ret)
            }
            _ => {
                self.errors
                    .push(TypeErrorKind::NotCallable(resolved.to_string()).at(call.span));
                self.env.fresh_var()
            }
        }
    }

    fn validate_from_conversion(&mut self, call: &CallExpr, arg_types: &[Type], ret: &Type) {
        let is_from_path = matches!(
            call.function.as_ref(),
            Expr::PathExpr(p) if p.segments.len() == 2 && p.segments[1] == "from"
        );
        if !is_from_path || arg_types.len() != 1 {
            return;
        }
        let source = self.subst.apply(&arg_types[0]);
        let target = self.subst.apply(ret);
        if source.is_integer()
            && target.is_integer()
            && !crate::ty::is_lossless_conversion(&source, &target)
        {
            self.errors.push(
                TypeErrorKind::Mismatch {
                    expected: format!("a type safely convertible to `{target}`"),
                    found: source.to_string(),
                }
                .at(call.span),
            );
        }
    }

    fn check_macro_call(&mut self, mc: &mut MacroCallExpr) -> Type {
        for arg in &mut mc.arguments {
            self.infer_expression(arg);
        }
        match mc.name.as_str() {
            "panic" | "todo" | "unimplemented" => Type::Never,
            "format" => Type::Str,
            _ => Type::Unit,
        }
    }

    fn check_struct_literal(&mut self, lit: &mut StructLitExpr) -> Type {
        let struct_def = self.env.lookup_struct(&lit.name).cloned();
        let Some(def) = struct_def else {
            self.errors
                .push(TypeErrorKind::UnknownType(lit.name.clone()).at(lit.span));
            return self.env.fresh_var();
        };
        let type_args = def
            .generic_params
            .iter()
            .map(|_| self.env.fresh_var())
            .collect::<Vec<Type>>();
        for (field_name, field_expr) in &mut lit.fields {
            let field_ty = self.infer_expression(field_expr);
            if let Some((_, expected_ty)) = def.fields.iter().find(|(n, _)| n == field_name) {
                let resolved =
                    self.instantiate_generic_type(expected_ty, &def.generic_params, &type_args);
                if let Err(e) = self.subst.unify(&resolved, &field_ty) {
                    self.report_unify_error(e, field_expr.span());
                }
            } else {
                self.errors.push(
                    TypeErrorKind::UnknownField {
                        ty: lit.name.clone(),
                        field: field_name.clone(),
                    }
                    .at(lit.span),
                );
            }
        }
        let expected_struct_ty = Type::Struct(
            Rc::from(lit.name.as_str()),
            type_args.iter().map(|t| self.subst.apply(t)).collect(),
        );
        if let Some(base_expr) = &mut lit.base {
            let base_ty = self.infer_expression(base_expr);
            if let Err(e) = self.subst.unify(&expected_struct_ty, &base_ty) {
                self.report_unify_error(e, base_expr.span());
            }
        } else {
            for (def_field_name, _) in &def.fields {
                if !lit.fields.iter().any(|(n, _)| n == def_field_name) {
                    self.errors.push(
                        TypeErrorKind::UnknownField {
                            ty: format!("missing field `{}` in `{}`", def_field_name, lit.name),
                            field: def_field_name.clone(),
                        }
                        .at(lit.span),
                    );
                }
            }
        }
        expected_struct_ty
    }

    fn check_path_expr(&mut self, path: &PathExpr) -> Type {
        if path.segments.len() == 2 {
            let type_name = &path.segments[0];
            let variant_name = &path.segments[1];
            if let Some(enum_def) = self.env.lookup_enum(type_name).cloned() {
                let type_args = enum_def
                    .generic_params
                    .iter()
                    .map(|_| self.env.fresh_var())
                    .collect::<Vec<Type>>();
                if let Some(variant) = enum_def.variants.iter().find(|v| v.name == *variant_name) {
                    match &variant.kind {
                        VariantKind::Unit => {
                            return Type::Enum(
                                Rc::from(type_name.as_str()),
                                type_args.iter().map(|t| self.subst.apply(t)).collect(),
                            );
                        }
                        VariantKind::Tuple(field_types) => {
                            let params = field_types
                                .iter()
                                .map(|t| {
                                    self.instantiate_generic_type(
                                        t,
                                        &enum_def.generic_params,
                                        &type_args,
                                    )
                                })
                                .collect();
                            let ret = Type::Enum(
                                Rc::from(type_name.as_str()),
                                type_args.iter().map(|t| self.subst.apply(t)).collect(),
                            );
                            return Type::Function(FnType {
                                params,
                                ret: Box::new(ret),
                            });
                        }
                        VariantKind::Struct(_) => {
                            return Type::Enum(
                                Rc::from(type_name.as_str()),
                                type_args.iter().map(|t| self.subst.apply(t)).collect(),
                            );
                        }
                    }
                } else {
                    self.errors.push(
                        TypeErrorKind::UnknownField {
                            ty: type_name.clone(),
                            field: variant_name.clone(),
                        }
                        .at(path.span),
                    );
                }
            }
        }
        // Try as a qualified method lookup (e.g., `Counter::new`).
        if path.segments.len() == 2 {
            let type_name = &path.segments[0];
            let method_name = &path.segments[1];
            if let Some(method_ty) = self.env.lookup_method_by_name(type_name, method_name) {
                return method_ty.clone();
            }
        }
        // Fallback: try as a variable lookup.
        let full_name = path.segments.join("::");
        self.env
            .instantiate(&full_name, &self.subst)
            .unwrap_or_else(|| {
                self.errors
                    .push(TypeErrorKind::UnknownType(full_name).at(path.span));
                self.env.fresh_var()
            })
    }

    fn check_field_access(&mut self, fa: &mut FieldAccessExpr) -> Type {
        let obj_ty = self.infer_expression(&mut fa.object);
        let resolved = self.subst.apply(&obj_ty);
        match &resolved {
            Type::Struct(name, type_args) => {
                let struct_def = self.env.lookup_struct(name).cloned();
                match struct_def {
                    Some(def) => match def.fields.iter().find(|(fname, _)| fname == &fa.field) {
                        Some((_, field_ty)) => {
                            self.instantiate_generic_type(field_ty, &def.generic_params, type_args)
                        }
                        None => {
                            self.errors.push(
                                TypeErrorKind::UnknownField {
                                    ty: resolved.to_string(),
                                    field: fa.field.clone(),
                                }
                                .at(fa.span),
                            );
                            self.env.fresh_var()
                        }
                    },
                    None => {
                        self.errors
                            .push(TypeErrorKind::UnknownType(name.to_string()).at(fa.span));
                        self.env.fresh_var()
                    }
                }
            }
            Type::Tuple(elems) => match fa.field.parse::<usize>() {
                Ok(idx) if idx < elems.len() => elems[idx].clone(),
                _ => {
                    self.errors.push(
                        TypeErrorKind::UnknownField {
                            ty: resolved.to_string(),
                            field: fa.field.clone(),
                        }
                        .at(fa.span),
                    );
                    self.env.fresh_var()
                }
            },
            Type::Var(_) => self.env.fresh_var(),
            _ => {
                self.errors.push(
                    TypeErrorKind::UnknownField {
                        ty: resolved.to_string(),
                        field: fa.field.clone(),
                    }
                    .at(fa.span),
                );
                self.env.fresh_var()
            }
        }
    }

    fn check_method_call(&mut self, mc: &mut MethodCallExpr) -> Type {
        let obj_ty = self.infer_expression(&mut mc.object);
        let resolved = self.subst.apply(&obj_ty);
        mc.receiver = resolved.receiver_name();
        if let Type::Array(_, n) = &resolved {
            mc.array_len = Some(*n);
            if mc.method == "len" || mc.method == "length" {
                if !mc.arguments.is_empty() {
                    self.errors.push(
                        TypeErrorKind::WrongArity {
                            expected: 0,
                            found: mc.arguments.len(),
                        }
                        .at(mc.span),
                    );
                }
                return Type::Usize;
            }
        }

        let arg_types = mc
            .arguments
            .iter_mut()
            .map(|a| self.infer_expression(a))
            .collect::<Vec<Type>>();
        let method_ty = self
            .env
            .instantiate_builtin_method(&resolved, &mc.method, &self.subst)
            .map(|(inst_self, inst_fn)| {
                if let Err(e) = self.subst.unify(&inst_self, &resolved) {
                    self.report_unify_error(e, mc.span);
                }
                self.subst.apply(&inst_fn)
            })
            .or_else(|| {
                self.env
                    .lookup_method(&resolved, &mc.method)
                    .or_else(|| match &resolved {
                        Type::Struct(name, _) | Type::Enum(name, _) => {
                            self.env.lookup_method_by_name(name, &mc.method)
                        }
                        _ => None,
                    })
                    .cloned()
                    .map(|ty| self.subst.apply(&ty))
            });
        match method_ty {
            Some(Type::Function(fn_ty)) => {
                if fn_ty.params.len() != arg_types.len() {
                    self.errors.push(
                        TypeErrorKind::WrongArity {
                            expected: fn_ty.params.len(),
                            found: arg_types.len(),
                        }
                        .at(mc.span),
                    );
                } else {
                    for (param, arg) in fn_ty.params.iter().zip(arg_types.iter()) {
                        let p = self.subst.apply(param);
                        let a = self.subst.apply(arg);
                        if let Err(e) = self.subst.unify(&p, &a) {
                            self.report_unify_error(e, mc.span);
                        }
                    }
                }
                self.subst.apply(&fn_ty.ret)
            }
            Some(_) => self.env.fresh_var(),
            None => {
                if matches!(resolved, Type::Var(_)) {
                    return self.env.fresh_var();
                }
                self.errors.push(
                    TypeErrorKind::UnknownMethod {
                        ty: resolved.to_string(),
                        method: mc.method.clone(),
                    }
                    .at(mc.span),
                );
                self.env.fresh_var()
            }
        }
    }

    fn check_try_expr(&mut self, expr: &mut TryExpr) -> Type {
        let inner_ty = self.infer_expression(&mut expr.expr);
        let resolved = self.subst.apply(&inner_ty);
        match &resolved {
            Type::Enum(name, args) if &**name == "Option" && args.len() == 1 => {
                expr.kind = TryKind::Option;
                args[0].clone()
            }
            Type::Enum(name, args) if &**name == "Result" && args.len() == 2 => {
                expr.kind = TryKind::Result;
                args[0].clone()
            }
            _ => {
                self.errors.push(
                    TypeErrorKind::Mismatch {
                        expected: "Option<T> or Result<T, E>".to_string(),
                        found: resolved.to_string(),
                    }
                    .at(expr.span),
                );
                self.env.fresh_var()
            }
        }
    }

    fn check_match_expr(&mut self, expr: &mut MatchExpr) -> Type {
        let scrutinee_ty = self.infer_expression(&mut expr.scrutinee);
        let scrutinee_resolved = self.subst.apply(&scrutinee_ty);

        let mut arm_result_ty: Option<Type> = None;

        for arm in &mut expr.arms {
            self.env.push_scope();
            self.check_pattern(&arm.pattern, &scrutinee_resolved);

            if let Some(guard) = &mut arm.guard {
                let guard_ty = self.infer_expression(guard);
                let guard_resolved = self.subst.apply(&guard_ty);
                self.require_bool(&guard_resolved, guard.span());
            }
            let body_ty = self.infer_expression(&mut arm.body);
            self.env.pop_scope();
            match &arm_result_ty {
                Some(prev) => {
                    if let Err(e) = self.subst.unify(prev, &body_ty) {
                        self.report_unify_error(e, arm.span);
                    }
                }
                None => arm_result_ty = Some(body_ty),
            }
        }
        self.check_exhaustiveness(&scrutinee_resolved, expr);
        arm_result_ty
            .map(|ty| self.subst.apply(&ty))
            .unwrap_or(Type::Never)
    }

    fn check_pattern(&mut self, p: &Pattern, scrutinee_ty: &Type) {
        match p {
            Pattern::Wildcard(_) => {}
            Pattern::Ident { name, .. } => {
                self.env.define_var(name, scrutinee_ty.clone());
            }
            Pattern::Literal(expr) => {
                let lit_ty = Type::from_literal_expr(expr);
                if self.subst.unify(&lit_ty, scrutinee_ty).is_err() {
                    self.errors.push(
                        TypeErrorKind::Mismatch {
                            expected: scrutinee_ty.to_string(),
                            found: lit_ty.to_string(),
                        }
                        .at(expr.span()),
                    );
                }
            }
            Pattern::TupleStruct { path, fields, span } => {
                self.check_tuple_struct_pattern(path, fields, scrutinee_ty, *span);
            }
            Pattern::Struct { path, fields, span } => {
                self.check_struct_pattern(path, fields, scrutinee_ty, *span);
            }
            Pattern::Tuple(fields, span) => {
                let resolved = self.subst.apply(scrutinee_ty);
                match &resolved {
                    Type::Tuple(elem_types) => {
                        if fields.len() != elem_types.len() {
                            self.errors.push(
                                TypeErrorKind::Mismatch {
                                    expected: format!("tuple of {} elements", elem_types.len()),
                                    found: format!("tuple pattern with {} elements", fields.len()),
                                }
                                .at(*span),
                            );
                        } else {
                            for (pat, ty) in fields.iter().zip(elem_types.iter()) {
                                self.check_pattern(pat, ty);
                            }
                        }
                    }
                    _ => {
                        self.errors.push(
                            TypeErrorKind::Mismatch {
                                expected: resolved.to_string(),
                                found: "tuple pattern".to_string(),
                            }
                            .at(*span),
                        );
                    }
                }
            }
            Pattern::Or(patterns, _) => {
                for pat in patterns {
                    self.check_pattern(pat, scrutinee_ty);
                }
            }
        }
    }

    fn check_tuple_struct_pattern(
        &mut self,
        variant_name: &str,
        fields: &[Pattern],
        scrutinee_ty: &Type,
        span: Span,
    ) {
        let enum_info = match scrutinee_ty {
            Type::Enum(name, type_args) => self
                .env
                .lookup_enum(name)
                .cloned()
                .map(|def| (def, type_args.clone())),
            _ => self.find_enum_for_variant(variant_name),
        };
        let Some((enum_def, type_args)) = enum_info else {
            // Not a known variant; skip checking but still bind identifiers.
            for field in fields {
                if let Pattern::Ident { name, .. } = field {
                    let fresh = self.env.fresh_var();
                    self.env.define_var(name, fresh);
                }
            }
            return;
        };
        let Some(variant) = enum_def.variants.iter().find(|v| v.name == variant_name) else {
            self.errors.push(
                TypeErrorKind::UnknownField {
                    ty: enum_def.name.clone(),
                    field: variant_name.to_string(),
                }
                .at(span),
            );
            return;
        };
        match &variant.kind {
            VariantKind::Tuple(payload_types) => {
                if payload_types.len() != fields.len() {
                    self.errors.push(
                        TypeErrorKind::WrongArity {
                            expected: payload_types.len(),
                            found: fields.len(),
                        }
                        .at(span),
                    );
                    return;
                }
                for (field_pat, payload_ty) in fields.iter().zip(payload_types.iter()) {
                    let resolved = self.instantiate_generic_type(
                        payload_ty,
                        &enum_def.generic_params,
                        &type_args,
                    );
                    self.check_pattern(field_pat, &resolved);
                }
            }
            VariantKind::Unit => {
                if !fields.is_empty() {
                    self.errors.push(
                        TypeErrorKind::WrongArity {
                            expected: 0,
                            found: fields.len(),
                        }
                        .at(span),
                    );
                }
            }
            VariantKind::Struct(_) => {
                self.errors.push(
                    TypeErrorKind::Mismatch {
                        expected: "struct pattern".to_string(),
                        found: "tuple pattern".to_string(),
                    }
                    .at(span),
                );
            }
        }
    }

    fn check_struct_pattern(
        &mut self,
        type_name: &str,
        fields: &[PatternField],
        scrutinee_ty: &Type,
        span: Span,
    ) {
        let struct_info = match scrutinee_ty {
            Type::Struct(name, type_args) if name.as_ref() == type_name => self
                .env
                .lookup_struct(name)
                .cloned()
                .map(|def| (def, type_args.clone())),
            _ => None,
        };
        let Some((struct_def, type_args)) = struct_info else {
            for field in fields {
                let name = field
                    .pattern
                    .as_ref()
                    .and_then(|p| match p.as_ref() {
                        Pattern::Ident { name: n, .. } => Some(n.clone()),
                        _ => None,
                    })
                    .unwrap_or_else(|| field.name.clone());
                let fresh = self.env.fresh_var();
                self.env.define_var(&name, fresh);
            }
            return;
        };
        for pf in fields {
            match struct_def
                .fields
                .iter()
                .find(|(fname, _)| fname == &pf.name)
            {
                Some((_, field_ty)) => {
                    let resolved = self.instantiate_generic_type(
                        field_ty,
                        &struct_def.generic_params,
                        &type_args,
                    );
                    match &pf.pattern {
                        Some(sub_pat) => self.check_pattern(sub_pat, &resolved),
                        None => self.env.define_var(&pf.name, resolved),
                    }
                }
                None => {
                    self.errors.push(
                        TypeErrorKind::UnknownField {
                            ty: struct_def.name.clone(),
                            field: pf.name.clone(),
                        }
                        .at(span),
                    );
                }
            }
        }
    }

    fn resolve_bare_variant(&mut self, name: &str) -> Option<Type> {
        let (enum_def, type_args) = self.find_enum_for_variant(name)?;
        let variant = enum_def.variants.iter().find(|v| v.name == name)?;
        let enum_name: Rc<str> = Rc::from(enum_def.name.as_str());
        match &variant.kind {
            VariantKind::Unit => Some(Type::Enum(
                enum_name,
                type_args.iter().map(|t| self.subst.apply(t)).collect(),
            )),
            VariantKind::Tuple(field_types) => {
                let params = field_types
                    .iter()
                    .map(|t| self.instantiate_generic_type(t, &enum_def.generic_params, &type_args))
                    .collect();
                let ret = Type::Enum(
                    enum_name,
                    type_args.iter().map(|t| self.subst.apply(t)).collect(),
                );
                Some(Type::Function(FnType {
                    params,
                    ret: Box::new(ret),
                }))
            }
            VariantKind::Struct(_) => Some(Type::Enum(
                enum_name,
                type_args.iter().map(|t| self.subst.apply(t)).collect(),
            )),
        }
    }

    fn find_enum_for_variant(&mut self, variant_name: &str) -> Option<(EnumDef, Vec<Type>)> {
        let def = self
            .env
            .all_enums()
            .find(|def| def.variants.iter().any(|v| v.name == variant_name))
            .cloned()?;
        let type_args = def
            .generic_params
            .iter()
            .map(|_| self.env.fresh_var())
            .collect();
        Some((def, type_args))
    }

    fn instantiate_generic_type(
        &self,
        ty: &Type,
        generic_params: &[String],
        type_args: &[Type],
    ) -> Type {
        match ty {
            Type::Generic(name, _) => generic_params
                .iter()
                .position(|g| g.as_str() == name.as_ref())
                .and_then(|i| type_args.get(i))
                .cloned()
                .unwrap_or_else(|| ty.clone()),
            Type::Vector(elem) => Type::Vector(Box::new(self.instantiate_generic_type(
                elem,
                generic_params,
                type_args,
            ))),
            Type::Range(elem) => Type::Range(Box::new(self.instantiate_generic_type(
                elem,
                generic_params,
                type_args,
            ))),
            Type::Map(k, v) => Type::Map(
                Box::new(self.instantiate_generic_type(k, generic_params, type_args)),
                Box::new(self.instantiate_generic_type(v, generic_params, type_args)),
            ),
            Type::Function(fn_ty) => Type::Function(FnType {
                params: fn_ty
                    .params
                    .iter()
                    .map(|p| self.instantiate_generic_type(p, generic_params, type_args))
                    .collect(),
                ret: Box::new(self.instantiate_generic_type(&fn_ty.ret, generic_params, type_args)),
            }),
            Type::Struct(name, args) => Type::Struct(
                name.clone(),
                args.iter()
                    .map(|a| self.instantiate_generic_type(a, generic_params, type_args))
                    .collect(),
            ),
            Type::Enum(name, args) => Type::Enum(
                name.clone(),
                args.iter()
                    .map(|a| self.instantiate_generic_type(a, generic_params, type_args))
                    .collect(),
            ),
            _ => ty.clone(),
        }
    }

    fn check_exhaustiveness(&mut self, scrutinee_ty: &Type, match_expr: &MatchExpr) {
        let enum_def = match scrutinee_ty {
            Type::Enum(name, _) => self.env.lookup_enum(name).cloned(),
            _ => None,
        };
        let has_wildcard = match_expr
            .arms
            .iter()
            .any(|arm| arm.guard.is_none() && self.pattern_is_catch_all(&arm.pattern, &enum_def));
        if has_wildcard {
            return;
        }
        match scrutinee_ty {
            Type::Enum(_, _) => {
                if let Some(def) = &enum_def {
                    let covered = match_expr
                        .arms
                        .iter()
                        .filter_map(|arm| self.extract_variant_name(&arm.pattern, def))
                        .collect::<std::collections::HashSet<&str>>();
                    let missing = def
                        .variants
                        .iter()
                        .map(|v| v.name.as_str())
                        .filter(|name| !covered.contains(name))
                        .collect::<Vec<&str>>();
                    if !missing.is_empty() {
                        self.errors.push(
                            TypeErrorKind::NonExhaustiveMatch {
                                missing: missing.join(", "),
                            }
                            .at(match_expr.span),
                        );
                    }
                }
            }
            Type::Bool => {
                let has_true = match_expr.arms.iter().any(|arm| {
                    matches!(&arm.pattern, Pattern::Literal(e) if matches!(e.as_ref(), Expr::Bool(BoolLit { value: true, .. })))
                });
                let has_false = match_expr.arms.iter().any(|arm| {
                    matches!(&arm.pattern, Pattern::Literal(e) if matches!(e.as_ref(), Expr::Bool(BoolLit { value: false, .. })))
                });
                if !has_true || !has_false {
                    self.errors.push(
                        TypeErrorKind::NonExhaustiveMatch {
                            missing: "not all boolean values covered".to_string(),
                        }
                        .at(match_expr.span),
                    );
                }
            }
            _ => {
                // For non-enum, non-bool types, a catch-all is required.
                self.errors.push(
                    TypeErrorKind::NonExhaustiveMatch {
                        missing: "missing wildcard `_` or catch-all pattern".to_string(),
                    }
                    .at(match_expr.span),
                );
            }
        }
    }

    fn pattern_is_catch_all(&self, p: &Pattern, enum_def: &Option<EnumDef>) -> bool {
        match p {
            Pattern::Wildcard(_) => true,
            Pattern::Ident { name, .. } => !enum_def
                .as_ref()
                .is_some_and(|def| def.variants.iter().any(|v| v.name == *name)),
            _ => false,
        }
    }

    fn extract_variant_name<'a>(&self, p: &'a Pattern, enum_def: &EnumDef) -> Option<&'a str> {
        match p {
            Pattern::TupleStruct { path, .. } => Some(path.as_str()),
            Pattern::Struct { path, .. } => Some(path.as_str()),
            Pattern::Ident { name, .. } => {
                // Check if this identifier is actually a unit variant name.
                if enum_def.variants.iter().any(|v| v.name == *name) {
                    Some(name.as_str())
                } else {
                    None
                }
            }
            Pattern::Or(pats, _) => pats
                .iter()
                .find_map(|p| self.extract_variant_name(p, enum_def)),
            _ => None,
        }
    }

    fn check_infix(&mut self, infix: &mut InfixExpr) -> Type {
        if infix.operator == "&&" || infix.operator == "||" {
            let lhs_ty = self.infer_expression(&mut infix.lhs);
            let lhs_resolved = self.subst.apply(&lhs_ty);
            self.require_bool(&lhs_resolved, infix.lhs.span());

            let rhs_ty = self.infer_expression(&mut infix.rhs);
            let rhs_resolved = self.subst.apply(&rhs_ty);
            self.require_bool(&rhs_resolved, infix.rhs.span());

            return Type::Bool;
        }

        let lhs_ty = self.infer_expression(&mut infix.lhs);
        let rhs_ty = self.infer_expression(&mut infix.rhs);
        let lhs_resolved = self.subst.apply(&lhs_ty);
        let rhs_resolved = self.subst.apply(&rhs_ty);

        let is_comparison = matches!(
            infix.operator.as_str(),
            "<" | ">" | "<=" | ">=" | "==" | "!="
        );

        if matches!(infix.operator.as_str(), "&" | "|" | "^" | "<<" | ">>") {
            self.bitwise_ops.push(BitwiseOp {
                operator: infix.operator.clone(),
                lhs_ty: lhs_resolved.clone(),
                rhs_ty: rhs_resolved.clone(),
                lhs_span: infix.lhs.span(),
                rhs_span: infix.rhs.span(),
            });
        }

        if matches!(infix.operator.as_str(), "==" | "!=")
            && let (Type::Array(_, n_lhs), Type::Array(_, n_rhs)) = (&lhs_resolved, &rhs_resolved)
            && n_lhs == n_rhs
        {
            infix.array_eq_len = Some(*n_lhs);
        }

        if lhs_resolved == Type::Str && rhs_resolved == Type::Str {
            if infix.operator == "+" {
                return Type::Str;
            }
            if is_comparison {
                return Type::Bool;
            }
        }

        if (infix.operator == "==" || infix.operator == "!=")
            && lhs_resolved == Type::Bool
            && rhs_resolved == Type::Bool
        {
            return Type::Bool;
        }

        if lhs_resolved == Type::Char && rhs_resolved == Type::Char && is_comparison {
            return Type::Bool;
        }

        if lhs_resolved == Type::Felt && rhs_resolved == Type::Felt {
            match infix.operator.as_str() {
                "+" | "-" | "*" | "/" => {
                    infix.op_class = BinOpClass::Felt;
                    return Type::Felt;
                }
                "==" | "!=" => {
                    infix.op_class = BinOpClass::Felt;
                    return Type::Bool;
                }
                "<" | ">" | "<=" | ">=" => {
                    self.errors.push(
                        TypeErrorKind::Unsupported(format!(
                            "ordering operator `{}` is not defined on `Felt` (prime-field \
                             elements have no canonical order)",
                            infix.operator
                        ))
                        .at(infix.span),
                    );
                    return Type::Felt;
                }
                _ => {}
            }
        }

        if lhs_resolved.is_integer() && rhs_resolved.is_integer() {
            if lhs_resolved != rhs_resolved {
                if matches!(lhs_resolved, Type::IntVar(_))
                    || matches!(rhs_resolved, Type::IntVar(_))
                {
                    if let Err(e) = self.subst.unify(&lhs_resolved, &rhs_resolved) {
                        self.report_unify_error(e, infix.span);
                    }
                } else {
                    self.errors.push(
                        TypeErrorKind::NumericMismatch {
                            expected: lhs_resolved.to_string(),
                            found: rhs_resolved.to_string(),
                        }
                        .at(infix.span),
                    );
                }
            }
            return if is_comparison {
                Type::Bool
            } else {
                self.subst.apply(&lhs_resolved)
            };
        }
        if matches!(lhs_resolved, Type::Var(_)) || matches!(rhs_resolved, Type::Var(_)) {
            if let Err(e) = self.subst.unify(&lhs_resolved, &rhs_resolved) {
                self.report_unify_error(e, infix.span);
            }
            if is_comparison {
                Type::Bool
            } else {
                self.subst.apply(&lhs_resolved)
            }
        } else {
            if let Err(e) = self.subst.unify(&lhs_resolved, &rhs_resolved) {
                self.report_unify_error(e, infix.span);
            }
            if is_comparison {
                Type::Bool
            } else {
                lhs_resolved
            }
        }
    }

    fn validate_bitwise_ops(&mut self) {
        let ops = std::mem::take(&mut self.bitwise_ops);
        for op in ops {
            let lhs = self.subst.resolve_int_vars(&op.lhs_ty);
            let rhs = self.subst.resolve_int_vars(&op.rhs_ty);
            self.check_bitwise_operand(&op.operator, &lhs, op.lhs_span);
            self.check_bitwise_operand(&op.operator, &rhs, op.rhs_span);
        }
    }

    fn check_bitwise_operand(&mut self, operator: &str, ty: &Type, span: Span) {
        let is_shift = matches!(operator, "<<" | ">>");
        if is_shift {
            if ty.is_shift_safe_unsigned() {
                return;
            }
            self.errors.push(
                TypeErrorKind::Unsupported(format!(
                    "bitwise operator `{operator}` requires an unsigned integer operand \
                     in the range `u64`/`usize`, found `{ty}`; cast the operand via `as u64`."
                ))
                .at(span),
            );
            return;
        }
        if ty.is_bitwise_safe_unsigned() {
            return;
        }
        self.errors.push(
            TypeErrorKind::Unsupported(format!(
                "bitwise operator `{operator}` requires an unsigned integer operand \
                 in the range `u8`/`u16`/`u32`/`u64`/`usize`, found `{ty}`; \
                 cast the operand via `as u64` before applying `{operator}`."
            ))
            .at(span),
        );
    }

    fn validate_u64_literals(&mut self, program: &mut Program) {
        for stmt in &program.statements {
            self.validate_u64_literals_in_stmt(stmt);
        }
    }

    fn validate_u64_literals_in_stmt(&mut self, stmt: &Stmt) {
        match stmt {
            Stmt::Let(let_stmt) => self.validate_u64_literals_in_expr(&let_stmt.value),
            Stmt::ReAssign(assign) => self.validate_u64_literals_in_expr(&assign.value),
            Stmt::Return(ret) => self.validate_u64_literals_in_expr(&ret.value),
            Stmt::Expr(expr_stmt) => self.validate_u64_literals_in_expr(&expr_stmt.value),
            Stmt::Block(block) => {
                for s in &block.statements {
                    self.validate_u64_literals_in_stmt(s);
                }
            }
            Stmt::FuncDef(fn_item) => {
                for s in &fn_item.body.statements {
                    self.validate_u64_literals_in_stmt(s);
                }
            }
            Stmt::Loop(loop_stmt) => {
                for s in &loop_stmt.body.statements {
                    self.validate_u64_literals_in_stmt(s);
                }
            }
            Stmt::While(while_stmt) => {
                self.validate_u64_literals_in_expr(&while_stmt.condition);
                for s in &while_stmt.body.statements {
                    self.validate_u64_literals_in_stmt(s);
                }
            }
            Stmt::For(for_stmt) => {
                self.validate_u64_literals_in_expr(&for_stmt.iterable);
                for s in &for_stmt.body.statements {
                    self.validate_u64_literals_in_stmt(s);
                }
            }
            Stmt::Mod(mod_stmt) => {
                if let Some(body) = &mod_stmt.body {
                    for s in body {
                        self.validate_u64_literals_in_stmt(s);
                    }
                }
            }
            Stmt::ImplBlock(impl_block) => {
                for method in &impl_block.methods {
                    for s in &method.body.statements {
                        self.validate_u64_literals_in_stmt(s);
                    }
                }
            }
            Stmt::StructDecl(_) | Stmt::EnumDecl(_) | Stmt::TraitDecl(_) | Stmt::Use(_) => {}
        }
    }

    fn validate_u64_literals_in_expr(&mut self, expr: &Expr) {
        match expr {
            Expr::Number(lit) => self.validate_u64_literal_value(lit),
            Expr::Prefix(prefix) => self.validate_u64_literals_in_expr(&prefix.operand),
            Expr::Infix(infix) => {
                self.validate_u64_literals_in_expr(&infix.lhs);
                self.validate_u64_literals_in_expr(&infix.rhs);
            }
            Expr::Call(call) => {
                self.validate_u64_literals_in_expr(&call.function);
                for arg in &call.arguments {
                    self.validate_u64_literals_in_expr(arg);
                }
            }
            Expr::MacroCall(mc) => {
                for arg in &mc.arguments {
                    self.validate_u64_literals_in_expr(arg);
                }
            }
            Expr::MethodCall(mc) => {
                self.validate_u64_literals_in_expr(&mc.object);
                for arg in &mc.arguments {
                    self.validate_u64_literals_in_expr(arg);
                }
            }
            Expr::Index(idx) => {
                self.validate_u64_literals_in_expr(&idx.expr);
                self.validate_u64_literals_in_expr(&idx.index);
            }
            Expr::FieldAccess(fa) => self.validate_u64_literals_in_expr(&fa.object),
            Expr::Cast(cast) => self.validate_u64_literals_in_expr(&cast.expr),
            Expr::Try(try_expr) => self.validate_u64_literals_in_expr(&try_expr.expr),
            Expr::Vector(vec) => {
                for elem in &vec.elements {
                    self.validate_u64_literals_in_expr(elem);
                }
            }
            Expr::Map(map) => {
                for (k, v) in &map.pairs {
                    self.validate_u64_literals_in_expr(k);
                    self.validate_u64_literals_in_expr(v);
                }
            }
            Expr::Tuple(tuple) => {
                for elem in &tuple.elements {
                    self.validate_u64_literals_in_expr(elem);
                }
            }
            Expr::Range(range) => {
                self.validate_u64_literals_in_expr(&range.start);
                self.validate_u64_literals_in_expr(&range.end);
            }
            Expr::Array(arr) => {
                for elem in &arr.elements {
                    self.validate_u64_literals_in_expr(elem);
                }
            }
            Expr::Cond(cond) => {
                self.validate_u64_literals_in_expr(&cond.condition);
                for s in &cond.consequence.statements {
                    self.validate_u64_literals_in_stmt(s);
                }
                if let Some(alt) = &cond.alternative {
                    for s in &alt.statements {
                        self.validate_u64_literals_in_stmt(s);
                    }
                }
            }
            Expr::Lambda(lambda) => {
                for s in &lambda.body.statements {
                    self.validate_u64_literals_in_stmt(s);
                }
            }
            Expr::Match(match_expr) => {
                self.validate_u64_literals_in_expr(&match_expr.scrutinee);
                for arm in &match_expr.arms {
                    self.validate_u64_literals_in_pattern(&arm.pattern);
                    if let Some(guard) = &arm.guard {
                        self.validate_u64_literals_in_expr(guard);
                    }
                    self.validate_u64_literals_in_expr(&arm.body);
                }
            }
            Expr::Break(brk) => {
                if let Some(value) = &brk.value {
                    self.validate_u64_literals_in_expr(value);
                }
            }
            Expr::StructLit(struct_lit) => {
                for (_, value) in &struct_lit.fields {
                    self.validate_u64_literals_in_expr(value);
                }
                if let Some(base) = &struct_lit.base {
                    self.validate_u64_literals_in_expr(base);
                }
            }
            Expr::Bool(_)
            | Expr::Str(_)
            | Expr::Char(_)
            | Expr::Ident(_)
            | Expr::Continue(_)
            | Expr::PathExpr(_)
            | Expr::MacroLit(_) => {}
        }
    }

    fn validate_u64_literals_in_pattern(&mut self, pat: &Pattern) {
        match pat {
            Pattern::Literal(expr) => self.validate_u64_literals_in_expr(expr),
            Pattern::TupleStruct { fields, .. } => {
                for sub in fields {
                    self.validate_u64_literals_in_pattern(sub);
                }
            }
            Pattern::Struct { fields, .. } => {
                for f in fields {
                    if let Some(sub) = &f.pattern {
                        self.validate_u64_literals_in_pattern(sub);
                    }
                }
            }
            Pattern::Tuple(items, _) | Pattern::Or(items, _) => {
                for sub in items {
                    self.validate_u64_literals_in_pattern(sub);
                }
            }
            Pattern::Wildcard(_) | Pattern::Ident { .. } => {}
        }
    }

    fn validate_u64_literal_value(&mut self, lit: &Number) {
        if !matches!(lit.kind, NumKind::U64 | NumKind::Usize) {
            return;
        }
        let Ok(value) = u64::try_from(lit.value) else {
            return;
        };
        if value >= GOLDILOCKS_MODULUS {
            self.errors.push(
                TypeErrorKind::Unsupported(format!(
                    "u64 literal `{value}` lies in the Goldilocks off-canonical range \
                     `[p, 2^64)` (where `p = 2^64 - 2^32 + 1`); use a value less than \
                     `0xFFFF_FFFF_0000_0001`, or construct the field element \
                     directly via `Felt::new(...)`."
                ))
                .at(lit.span),
            );
        }
    }

    fn infer_block(&mut self, block: &mut BlockStmt) -> Type {
        let mut last = Type::Unit;
        for stmt in &mut block.statements {
            match stmt {
                Stmt::Expr(es) => {
                    last = self.infer_expression(&mut es.value);
                }
                Stmt::Return(ret) => {
                    self.infer_expression(&mut ret.value);
                    last = Type::Never;
                }
                _ => {
                    self.check_statement(stmt);
                    last = Type::Unit;
                }
            }
        }
        last
    }

    fn check_literal_range(&mut self, expr: &Expr, expected: &Type, span: Span) {
        let Some(val) = expr.extract_integer_value() else {
            return;
        };

        macro_rules! check_int_range {
            ($target:ty, $target_name:expr) => {
                if <$target>::try_from(val).is_err() {
                    self.errors.push(
                        TypeErrorKind::NumericOverflow {
                            value: val.to_string(),
                            target: $target_name.to_string(),
                        }
                        .at(span),
                    );
                }
            };
        }

        match expected {
            Type::I8 => check_int_range!(i8, "i8"),
            Type::I16 => check_int_range!(i16, "i16"),
            Type::I32 => check_int_range!(i32, "i32"),
            Type::I64 => check_int_range!(i64, "i64"),
            Type::I128 => {}
            Type::Isize => check_int_range!(i64, "isize"),
            Type::U8 => check_int_range!(u8, "u8"),
            Type::U16 => check_int_range!(u16, "u16"),
            Type::U32 => check_int_range!(u32, "u32"),
            Type::U64 => check_int_range!(u64, "u64"),
            Type::U128 => check_int_range!(u128, "u128"),
            Type::Usize => check_int_range!(u64, "usize"),
            _ => {}
        }
    }

    fn require_bool(&mut self, resolved: &Type, span: Span) {
        if !matches!(resolved, Type::Bool | Type::Var(_)) {
            self.errors.push(
                TypeErrorKind::Mismatch {
                    expected: "bool".to_string(),
                    found: resolved.to_string(),
                }
                .at(span),
            );
        }
        if let Type::Var(_) = resolved
            && let Err(e) = self.subst.unify(resolved, &Type::Bool)
        {
            self.report_unify_error(e, span);
        }
    }

    fn report_unify_error(&mut self, err: UnifyError, span: Span) {
        self.errors.push(err.to_type_error().at(span));
    }
}
