// ============================================================================
// PFCL Type Checker — Inference Engine (Algorithm W)
// ============================================================================
//
// Walks the resolved AST (RExpr from the resolver) and infers types using
// Hindley-Milner type inference with:
//   - Catalog function signatures as the type environment
//   - Let-polymorphism (generalize at let bindings)
//   - Row polymorphism for record types
//   - Constructor typing for Maybe/Either/List
//   - Pattern matching with exhaustive type propagation
//
// Produces either a Typed AST (TExpr) where every node carries its inferred
// type, or a list of TypeErrors that constitute the "gap report."

use crate::types::*;
use crate::unify::*;
use crate::{Box, String, ToString, Vec, vec, BTreeMap};
use alloc::format;
use hashbrown::HashSet;

// ---------------------------------------------------------------------------
// Resolved AST (input from the resolver)
// ---------------------------------------------------------------------------
// These mirror the resolver's output. In the real compiler, these would be
// imported from the resolver crate. Here we define them for self-containment.

/// A resolved expression node (output of the resolver, input to type checker).
#[derive(Debug, Clone)]
pub struct RExpr {
    pub kind: RExprKind,
    pub span: Span,
}

/// Byte range in source text (for error reporting).
#[derive(Debug, Clone, Copy, Default)]
pub struct Span {
    pub start: usize,
    pub end: usize,
}

#[derive(Debug, Clone)]
pub enum RExprKind {
    /// A bound variable (lambda param, let binding, match pattern var).
    /// De Bruijn index + original name for error messages.
    BoundVar { index: usize, name: String },

    /// A reference to a catalog function, resolved to its content hash.
    CatalogRef { hash: String, name: String },

    /// Field access on a bound variable: `expr.field1.field2`
    FieldAccess {
        expr: Box<RExpr>,
        fields: Vec<String>,
    },

    /// Literal values
    IntLit(i64),
    FloatLit(f64),
    StringLit(String),
    BoolLit(bool),
    NothingLit,

    /// Lambda: \(params...). body
    Lambda {
        params: Vec<String>,
        body: Box<RExpr>,
    },

    /// Function application: f(arg1, arg2, ...)
    Apply {
        func: Box<RExpr>,
        args: Vec<RExpr>,
    },

    /// Let binding: let name = value in body
    Let {
        name: String,
        value: Box<RExpr>,
        body: Box<RExpr>,
    },

    /// Match expression: match scrutinee with arms
    Match {
        scrutinee: Box<RExpr>,
        arms: Vec<MatchArm>,
    },

    /// List literal: [e1, e2, ...]
    ListLit(Vec<RExpr>),

    /// Record literal: { field1: e1, field2: e2 }
    RecordLit(Vec<(String, RExpr)>),

    /// Tuple literal: (e1, e2, ...)
    TupleLit(Vec<RExpr>),

    /// Constructor: Just(e), Left(e), etc.
    Constructor {
        name: String,
        args: Vec<RExpr>,
    },
}

/// A match arm: pattern → body
#[derive(Debug, Clone)]
pub struct MatchArm {
    pub pattern: Pattern,
    pub body: RExpr,
}

/// Pattern in a match expression.
#[derive(Debug, Clone)]
pub enum Pattern {
    /// Wildcard: _
    Wildcard,
    /// Variable binding
    Var(String),
    /// Literal patterns
    IntPat(i64),
    StringPat(String),
    BoolPat(bool),
    NothingPat,
    /// Constructor pattern: Just(p), Cons(h, t)
    Constructor { name: String, args: Vec<Pattern> },
    /// Tuple pattern: (p1, p2)
    TuplePat(Vec<Pattern>),
}

// ---------------------------------------------------------------------------
// Typed AST (output of the type checker)
// ---------------------------------------------------------------------------

/// A typed expression — every node carries its inferred type.
#[derive(Debug, Clone)]
pub struct TExpr {
    pub kind: TExprKind,
    pub ty: Type,
    pub span: Span,
}

#[derive(Debug, Clone)]
pub enum TExprKind {
    BoundVar { index: usize, name: String },
    CatalogRef { hash: String, name: String },
    FieldAccess {
        expr: Box<TExpr>,
        fields: Vec<String>,
    },
    IntLit(i64),
    FloatLit(f64),
    StringLit(String),
    BoolLit(bool),
    NothingLit,
    Lambda {
        params: Vec<(String, Type)>,
        body: Box<TExpr>,
    },
    Apply {
        func: Box<TExpr>,
        args: Vec<TExpr>,
    },
    Let {
        name: String,
        value: Box<TExpr>,
        body: Box<TExpr>,
    },
    Match {
        scrutinee: Box<TExpr>,
        arms: Vec<TypedMatchArm>,
    },
    ListLit(Vec<TExpr>),
    RecordLit(Vec<(String, TExpr)>),
    TupleLit(Vec<TExpr>),
    Constructor {
        name: String,
        args: Vec<TExpr>,
    },
}

#[derive(Debug, Clone)]
pub struct TypedMatchArm {
    pub pattern: Pattern,
    pub body: TExpr,
}

// ---------------------------------------------------------------------------
// Catalog Type Lookup (trait for the type environment)
// ---------------------------------------------------------------------------

/// Trait for looking up type signatures from the catalog.
/// The resolver already resolved names to hashes; now we need the types.
pub trait CatalogTypes {
    /// Look up the type scheme for a catalog function by its content hash.
    /// Returns None if the hash is unknown (produces a type error).
    fn type_of(&self, hash: &str) -> Option<TypeScheme>;
}

// ---------------------------------------------------------------------------
// Type Environment
// ---------------------------------------------------------------------------

/// The typing context: maps de Bruijn indices to type schemes.
/// Index 0 is the innermost binding.
///
/// We use a Vec as a stack. Push on lambda/let/match entry, pop on exit.
/// This mirrors the resolver's scope stack exactly.
#[derive(Debug, Clone)]
struct TypeEnv {
    /// Stack of (name, type_scheme) bindings. Index 0 = outermost.
    /// We look up by de Bruijn index counting from the end.
    bindings: Vec<(String, TypeScheme)>,
}

impl TypeEnv {
    fn new() -> Self {
        TypeEnv {
            bindings: Vec::new(),
        }
    }

    /// Push a monomorphic binding (lambda params, match vars).
    fn push_mono(&mut self, name: String, ty: Type) {
        self.bindings.push((name, TypeScheme::mono(ty)));
    }

    /// Push a polymorphic binding (let bindings with generalization).
    fn push_poly(&mut self, name: String, scheme: TypeScheme) {
        self.bindings.push((name, scheme));
    }

    /// Pop the top n bindings.
    fn pop(&mut self, n: usize) {
        let new_len = self.bindings.len().saturating_sub(n);
        self.bindings.truncate(new_len);
    }

    /// Look up by de Bruijn index.
    /// Index 0 = most recently pushed binding.
    fn lookup(&self, index: usize) -> Option<&TypeScheme> {
        if index < self.bindings.len() {
            let actual = self.bindings.len() - 1 - index;
            Some(&self.bindings[actual].1)
        } else {
            None
        }
    }

    /// Collect all free type variables in the environment.
    fn free_vars(&self) -> HashSet<TypeVar> {
        let mut fv = HashSet::new();
        for (_, scheme) in &self.bindings {
            fv.extend(scheme.free_vars());
        }
        fv
    }
}

// ---------------------------------------------------------------------------
// Inference Engine
// ---------------------------------------------------------------------------

/// The type inference engine. Stateful: owns the fresh variable supply
/// and accumulates errors for gap reporting.
pub struct Inferencer<'a, C: CatalogTypes> {
    catalog: &'a C,
    fresh: FreshSupply,
    errors: Vec<TypeError>,
}

/// The result of type checking a complete expression.
pub struct TypeCheckResult {
    /// The typed AST (may have placeholder types at error sites).
    pub typed_expr: TExpr,
    /// The inferred type of the whole expression.
    pub ty: Type,
    /// Any type errors found (empty = well-typed).
    pub errors: Vec<TypeError>,
}

impl TypeCheckResult {
    pub fn is_ok(&self) -> bool {
        self.errors.is_empty()
    }

    /// Format all errors as a gap report.
    pub fn gap_report(&self) -> String {
        if self.errors.is_empty() {
            return format!("✓ Well-typed: {}", self.ty);
        }
        let mut report = String::new();
        report.push_str(&format!(
            "✗ {} type error(s) found:\n\n",
            self.errors.len()
        ));
        for (i, err) in self.errors.iter().enumerate() {
            report.push_str(&format!("  {}. {}\n", i + 1, err));
        }
        report
    }
}

impl<'a, C: CatalogTypes> Inferencer<'a, C> {
    pub fn new(catalog: &'a C) -> Self {
        Inferencer {
            catalog,
            fresh: FreshSupply::new(),
            errors: Vec::new(),
        }
    }

    /// Type-check a resolved expression. This is the main entry point.
    pub fn check(mut self, expr: &RExpr) -> TypeCheckResult {
        let mut env = TypeEnv::new();
        let (subst, typed) = self.infer(&mut env, expr);
        let final_ty = typed.ty.apply_subst(&subst);
        let typed = self.apply_subst_to_texpr(&typed, &subst);
        TypeCheckResult {
            typed_expr: typed,
            ty: final_ty,
            errors: self.errors,
        }
    }

    /// Infer the type of an expression, returning the accumulated
    /// substitution and the typed AST node.
    fn infer(&mut self, env: &mut TypeEnv, expr: &RExpr) -> (Substitution, TExpr) {
        match &expr.kind {
            // ----- Literals -----
            RExprKind::IntLit(n) => (
                Substitution::empty(),
                TExpr {
                    kind: TExprKind::IntLit(*n),
                    ty: Type::int(),
                    span: expr.span,
                },
            ),
            RExprKind::FloatLit(f) => (
                Substitution::empty(),
                TExpr {
                    kind: TExprKind::FloatLit(*f),
                    ty: Type::float(),
                    span: expr.span,
                },
            ),
            RExprKind::StringLit(s) => (
                Substitution::empty(),
                TExpr {
                    kind: TExprKind::StringLit(s.clone()),
                    ty: Type::string(),
                    span: expr.span,
                },
            ),
            RExprKind::BoolLit(b) => (
                Substitution::empty(),
                TExpr {
                    kind: TExprKind::BoolLit(*b),
                    ty: Type::boolean(),
                    span: expr.span,
                },
            ),
            RExprKind::NothingLit => {
                // Nothing : Maybe<a> for a fresh variable
                let a = self.fresh.fresh_type();
                (
                    Substitution::empty(),
                    TExpr {
                        kind: TExprKind::NothingLit,
                        ty: Type::maybe(a),
                        span: expr.span,
                    },
                )
            }

            // ----- Bound Variables -----
            RExprKind::BoundVar { index, name } => {
                match env.lookup(*index) {
                    Some(scheme) => {
                        let ty = self.instantiate(scheme);
                        (
                            Substitution::empty(),
                            TExpr {
                                kind: TExprKind::BoundVar {
                                    index: *index,
                                    name: name.clone(),
                                },
                                ty,
                                span: expr.span,
                            },
                        )
                    }
                    None => {
                        // This shouldn't happen if the resolver is correct,
                        // but handle gracefully
                        let ty = self.fresh.fresh_type();
                        self.errors.push(TypeError {
                            kind: TypeErrorKind::Gap {
                                needed: ty.clone(),
                                description: format!(
                                    "Variable `{}` at index {} not in type environment",
                                    name, index
                                ),
                            },
                            context: vec![],
                        });
                        (
                            Substitution::empty(),
                            TExpr {
                                kind: TExprKind::BoundVar {
                                    index: *index,
                                    name: name.clone(),
                                },
                                ty,
                                span: expr.span,
                            },
                        )
                    }
                }
            }

            // ----- Catalog References -----
            RExprKind::CatalogRef { hash, name } => {
                match self.catalog.type_of(hash) {
                    Some(scheme) => {
                        let ty = self.instantiate(&scheme);
                        (
                            Substitution::empty(),
                            TExpr {
                                kind: TExprKind::CatalogRef {
                                    hash: hash.clone(),
                                    name: name.clone(),
                                },
                                ty,
                                span: expr.span,
                            },
                        )
                    }
                    None => {
                        let ty = self.fresh.fresh_type();
                        self.errors.push(TypeError::unknown_catalog(
                            hash.clone(),
                            name.clone(),
                        ));
                        (
                            Substitution::empty(),
                            TExpr {
                                kind: TExprKind::CatalogRef {
                                    hash: hash.clone(),
                                    name: name.clone(),
                                },
                                ty,
                                span: expr.span,
                            },
                        )
                    }
                }
            }

            // ----- Field Access -----
            RExprKind::FieldAccess { expr: inner, fields } => {
                self.infer_field_access(env, inner, fields, expr.span)
            }

            // ----- Lambda -----
            RExprKind::Lambda { params, body } => {
                self.infer_lambda(env, params, body, expr.span)
            }

            // ----- Application -----
            RExprKind::Apply { func, args } => {
                self.infer_apply(env, func, args, expr.span)
            }

            // ----- Let Binding -----
            RExprKind::Let { name, value, body } => {
                self.infer_let(env, name, value, body, expr.span)
            }

            // ----- Match Expression -----
            RExprKind::Match { scrutinee, arms } => {
                self.infer_match(env, scrutinee, arms, expr.span)
            }

            // ----- List Literal -----
            RExprKind::ListLit(elems) => self.infer_list(env, elems, expr.span),

            // ----- Record Literal -----
            RExprKind::RecordLit(fields) => self.infer_record(env, fields, expr.span),

            // ----- Tuple Literal -----
            RExprKind::TupleLit(elems) => self.infer_tuple(env, elems, expr.span),

            // ----- Constructor -----
            RExprKind::Constructor { name, args } => {
                self.infer_constructor(env, name, args, expr.span)
            }
        }
    }

    // -----------------------------------------------------------------------
    // Lambda inference
    // -----------------------------------------------------------------------

    fn infer_lambda(
        &mut self,
        env: &mut TypeEnv,
        params: &[String],
        body: &RExpr,
        span: Span,
    ) -> (Substitution, TExpr) {
        // Assign fresh type variables to each parameter
        let param_types: Vec<Type> = params
            .iter()
            .map(|_| self.fresh.fresh_type())
            .collect();

        // Push params onto env (in order — first param gets highest de Bruijn)
        // The resolver assigns de Bruijn 0 to the first param in a multi-param
        // lambda, then 1, 2... But our stack pushes in order so the last pushed
        // is index 0. We push in reverse so index 0 = first param.
        //
        // Actually — re-reading the resolver: \(x, y). x → x is BoundVar(0).
        // So index 0 is the *first* parameter. Our stack's top is index 0.
        // Push in reverse order: y first, then x, so x is on top (index 0).
        for (_i, (name, ty)) in params.iter().zip(param_types.iter()).rev().enumerate() {
            env.push_mono(name.clone(), ty.clone());
        }

        // Infer body type
        let (body_subst, body_typed) = self.infer(env, body);

        // Pop params
        env.pop(params.len());

        // Build the function type with substitutions applied
        let param_types_subst: Vec<Type> = param_types
            .iter()
            .map(|t| t.apply_subst(&body_subst))
            .collect();
        let fn_type = Type::Fun(param_types_subst.clone(), Box::new(body_typed.ty.clone()));

        let typed_params: Vec<(String, Type)> = params
            .iter()
            .zip(param_types_subst.iter())
            .map(|(n, t)| (n.clone(), t.clone()))
            .collect();

        (
            body_subst,
            TExpr {
                kind: TExprKind::Lambda {
                    params: typed_params,
                    body: Box::new(body_typed),
                },
                ty: fn_type,
                span,
            },
        )
    }

    // -----------------------------------------------------------------------
    // Application inference
    // -----------------------------------------------------------------------

    fn infer_apply(
        &mut self,
        env: &mut TypeEnv,
        func: &RExpr,
        args: &[RExpr],
        span: Span,
    ) -> (Substitution, TExpr) {
        // Infer the function type
        let (s1, func_typed) = self.infer(env, func);
        self.apply_subst_to_env(env, &s1);

        // Infer each argument type
        let mut subst = s1;
        let mut arg_types = Vec::new();
        let mut args_typed = Vec::new();

        for (_i, arg) in args.iter().enumerate() {
            self.apply_subst_to_env(env, &subst);
            let (s_arg, arg_typed) = self.infer(env, arg);
            subst = s_arg.compose(&subst);
            arg_types.push(arg_typed.ty.clone());
            args_typed.push(arg_typed);
        }

        // The function type after all substitutions
        let fn_ty = func_typed.ty.apply_subst(&subst);

        // Fresh variable for the result
        let result_ty = self.fresh.fresh_type();

        // Build the expected function type from args
        let expected_fn_type = Type::Fun(arg_types.clone(), Box::new(result_ty.clone()));

        // Unify fn_ty with expected_fn_type
        let fn_name = match &func_typed.kind {
            TExprKind::CatalogRef { name, .. } => name.clone(),
            _ => "<expr>".into(),
        };

        match unify(&fn_ty, &expected_fn_type) {
            Ok(s_unify) => {
                subst = s_unify.compose(&subst);
                let final_result = result_ty.apply_subst(&subst);
                (
                    subst,
                    TExpr {
                        kind: TExprKind::Apply {
                            func: Box::new(func_typed),
                            args: args_typed,
                        },
                        ty: final_result,
                        span,
                    },
                )
            }
            Err(mut err) => {
                // Add context about where the error occurred
                err = err.with_context(ErrorContext::InApplication {
                    fn_name: fn_name.clone(),
                });

                // Handle specific error cases with better messages
                match &fn_ty {
                    Type::Fun(expected_inputs, _) if expected_inputs.len() != args.len() => {
                        self.errors.push(
                            TypeError::arity(expected_inputs.len(), args.len(), fn_ty.clone())
                                .with_context(ErrorContext::InApplication { fn_name }),
                        );
                    }
                    t if !matches!(t, Type::Fun(_, _) | Type::Var(_)) => {
                        self.errors.push(
                            TypeError::not_a_function(fn_ty.clone())
                                .with_context(ErrorContext::InApplication { fn_name }),
                        );
                    }
                    _ => {
                        self.errors.push(err);
                    }
                }

                let fallback = self.fresh.fresh_type();
                (
                    subst,
                    TExpr {
                        kind: TExprKind::Apply {
                            func: Box::new(func_typed),
                            args: args_typed,
                        },
                        ty: fallback,
                        span,
                    },
                )
            }
        }
    }

    // -----------------------------------------------------------------------
    // Let binding inference (with let-polymorphism / generalization)
    // -----------------------------------------------------------------------

    fn infer_let(
        &mut self,
        env: &mut TypeEnv,
        name: &str,
        value: &RExpr,
        body: &RExpr,
        span: Span,
    ) -> (Substitution, TExpr) {
        // Infer the value type
        let (s1, value_typed) = self.infer(env, value);
        self.apply_subst_to_env(env, &s1);

        // Generalize: quantify over type variables free in the value type
        // but not free in the environment
        let value_ty = value_typed.ty.apply_subst(&s1);
        let scheme = self.generalize(env, &value_ty);

        // Push the generalized binding
        env.push_poly(name.to_string(), scheme);

        // Infer the body type
        let (s2, body_typed) = self.infer(env, body);

        // Pop the binding
        env.pop(1);

        let subst = s2.compose(&s1);
        let body_ty = body_typed.ty.clone();
        (
            subst,
            TExpr {
                kind: TExprKind::Let {
                    name: name.to_string(),
                    value: Box::new(value_typed),
                    body: Box::new(body_typed),
                },
                ty: body_ty,
                span,
            },
        )
    }

    // -----------------------------------------------------------------------
    // Match expression inference
    // -----------------------------------------------------------------------

    fn infer_match(
        &mut self,
        env: &mut TypeEnv,
        scrutinee: &RExpr,
        arms: &[MatchArm],
        span: Span,
    ) -> (Substitution, TExpr) {
        // Infer scrutinee type
        let (mut subst, scrut_typed) = self.infer(env, scrutinee);
        let mut scrut_ty = scrut_typed.ty.apply_subst(&subst);

        // All arms must produce the same type
        let result_ty = self.fresh.fresh_type();
        let mut typed_arms = Vec::new();

        for (i, arm) in arms.iter().enumerate() {
            self.apply_subst_to_env(env, &subst);
            scrut_ty = scrut_ty.apply_subst(&subst);

            // Check the pattern against the scrutinee type and get bindings
            let (pat_subst, bindings) =
                self.check_pattern(&arm.pattern, &scrut_ty);
            subst = pat_subst.compose(&subst);
            scrut_ty = scrut_ty.apply_subst(&subst);

            // Push pattern bindings onto the environment
            let binding_count = bindings.len();
            for (name, ty) in bindings.iter().rev() {
                env.push_mono(name.clone(), ty.apply_subst(&subst));
            }

            // Infer the arm body type
            let (s_body, body_typed) = self.infer(env, &arm.body);
            subst = s_body.compose(&subst);

            // Pop pattern bindings
            env.pop(binding_count);

            // Unify arm body type with the expected result type
            let arm_ty = body_typed.ty.apply_subst(&subst);
            let expected = result_ty.apply_subst(&subst);
            match unify(&arm_ty, &expected) {
                Ok(s_arm) => {
                    subst = s_arm.compose(&subst);
                }
                Err(err) => {
                    self.errors.push(
                        err.with_context(ErrorContext::InMatchArm { index: i }),
                    );
                }
            }

            typed_arms.push(TypedMatchArm {
                pattern: arm.pattern.clone(),
                body: body_typed,
            });
        }

        let final_ty = result_ty.apply_subst(&subst);
        (
            subst,
            TExpr {
                kind: TExprKind::Match {
                    scrutinee: Box::new(scrut_typed),
                    arms: typed_arms,
                },
                ty: final_ty,
                span,
            },
        )
    }

    /// Check a pattern against a type, returning substitution and bindings.
    fn check_pattern(
        &mut self,
        pattern: &Pattern,
        expected: &Type,
    ) -> (Substitution, Vec<(String, Type)>) {
        match pattern {
            Pattern::Wildcard => (Substitution::empty(), vec![]),

            Pattern::Var(name) => {
                // Bind the variable to the expected type
                (
                    Substitution::empty(),
                    vec![(name.clone(), expected.clone())],
                )
            }

            Pattern::IntPat(_) => {
                let s = match unify(expected, &Type::int()) {
                    Ok(s) => s,
                    Err(err) => {
                        self.errors.push(err);
                        Substitution::empty()
                    }
                };
                (s, vec![])
            }

            Pattern::StringPat(_) => {
                let s = match unify(expected, &Type::string()) {
                    Ok(s) => s,
                    Err(err) => {
                        self.errors.push(err);
                        Substitution::empty()
                    }
                };
                (s, vec![])
            }

            Pattern::BoolPat(_) => {
                let s = match unify(expected, &Type::boolean()) {
                    Ok(s) => s,
                    Err(err) => {
                        self.errors.push(err);
                        Substitution::empty()
                    }
                };
                (s, vec![])
            }

            Pattern::NothingPat => {
                // Nothing : Maybe<a>
                let a = self.fresh.fresh_type();
                let s = match unify(expected, &Type::maybe(a)) {
                    Ok(s) => s,
                    Err(err) => {
                        self.errors.push(err);
                        Substitution::empty()
                    }
                };
                (s, vec![])
            }

            Pattern::Constructor { name, args } => {
                self.check_constructor_pattern(name, args, expected)
            }

            Pattern::TuplePat(pats) => {
                // Expected must be a tuple of the same length
                let elem_types: Vec<Type> =
                    (0..pats.len()).map(|_| self.fresh.fresh_type()).collect();
                let tuple_ty = Type::Tuple(elem_types.clone());

                let mut subst = match unify(expected, &tuple_ty) {
                    Ok(s) => s,
                    Err(err) => {
                        self.errors.push(err);
                        Substitution::empty()
                    }
                };

                let mut all_bindings = Vec::new();
                for (pat, elem_ty) in pats.iter().zip(elem_types.iter()) {
                    let ty = elem_ty.apply_subst(&subst);
                    let (s, bindings) = self.check_pattern(pat, &ty);
                    subst = s.compose(&subst);
                    all_bindings.extend(bindings);
                }

                (subst, all_bindings)
            }
        }
    }

    /// Check a constructor pattern (Just, Nothing, Left, Right, Cons, etc.)
    fn check_constructor_pattern(
        &mut self,
        name: &str,
        args: &[Pattern],
        expected: &Type,
    ) -> (Substitution, Vec<(String, Type)>) {
        match name {
            "Just" => {
                if args.len() != 1 {
                    self.errors.push(TypeError {
                        kind: TypeErrorKind::PatternTypeMismatch {
                            pattern_desc: format!("Just(...) with {} args", args.len()),
                            expected: expected.clone(),
                        },
                        context: vec![],
                    });
                    return (Substitution::empty(), vec![]);
                }
                let inner = self.fresh.fresh_type();
                let maybe_ty = Type::maybe(inner.clone());
                let mut subst = match unify(expected, &maybe_ty) {
                    Ok(s) => s,
                    Err(err) => {
                        self.errors.push(err);
                        Substitution::empty()
                    }
                };
                let inner_resolved = inner.apply_subst(&subst);
                let (s, bindings) = self.check_pattern(&args[0], &inner_resolved);
                subst = s.compose(&subst);
                (subst, bindings)
            }

            "Left" => {
                if args.len() != 1 {
                    self.errors.push(TypeError {
                        kind: TypeErrorKind::PatternTypeMismatch {
                            pattern_desc: format!("Left with {} args", args.len()),
                            expected: expected.clone(),
                        },
                        context: vec![],
                    });
                    return (Substitution::empty(), vec![]);
                }
                let a = self.fresh.fresh_type();
                let b = self.fresh.fresh_type();
                let either_ty = Type::either(a.clone(), b);
                let mut subst = match unify(expected, &either_ty) {
                    Ok(s) => s,
                    Err(err) => {
                        self.errors.push(err);
                        Substitution::empty()
                    }
                };
                let a_resolved = a.apply_subst(&subst);
                let (s, bindings) = self.check_pattern(&args[0], &a_resolved);
                subst = s.compose(&subst);
                (subst, bindings)
            }

            "Right" => {
                if args.len() != 1 {
                    self.errors.push(TypeError {
                        kind: TypeErrorKind::PatternTypeMismatch {
                            pattern_desc: format!("Right with {} args", args.len()),
                            expected: expected.clone(),
                        },
                        context: vec![],
                    });
                    return (Substitution::empty(), vec![]);
                }
                let a = self.fresh.fresh_type();
                let b = self.fresh.fresh_type();
                let either_ty = Type::either(a, b.clone());
                let mut subst = match unify(expected, &either_ty) {
                    Ok(s) => s,
                    Err(err) => {
                        self.errors.push(err);
                        Substitution::empty()
                    }
                };
                let b_resolved = b.apply_subst(&subst);
                let (s, bindings) = self.check_pattern(&args[0], &b_resolved);
                subst = s.compose(&subst);
                (subst, bindings)
            }

            "Cons" => {
                if args.len() != 2 {
                    self.errors.push(TypeError {
                        kind: TypeErrorKind::PatternTypeMismatch {
                            pattern_desc: format!("Cons with {} args", args.len()),
                            expected: expected.clone(),
                        },
                        context: vec![],
                    });
                    return (Substitution::empty(), vec![]);
                }
                let elem = self.fresh.fresh_type();
                let list_ty = Type::list(elem.clone());
                let mut subst = match unify(expected, &list_ty) {
                    Ok(s) => s,
                    Err(err) => {
                        self.errors.push(err);
                        Substitution::empty()
                    }
                };

                let mut all_bindings = Vec::new();

                // Head element
                let elem_resolved = elem.apply_subst(&subst);
                let (s1, bindings1) = self.check_pattern(&args[0], &elem_resolved);
                subst = s1.compose(&subst);
                all_bindings.extend(bindings1);

                // Tail: List<elem>
                let tail_ty = Type::list(elem.apply_subst(&subst));
                let (s2, bindings2) = self.check_pattern(&args[1], &tail_ty);
                subst = s2.compose(&subst);
                all_bindings.extend(bindings2);

                (subst, all_bindings)
            }

            _other => {
                // Unknown constructor — treat generically
                // All args get fresh types, the constructor itself unifies with expected
                let mut subst = Substitution::empty();
                let mut all_bindings = Vec::new();
                for arg in args {
                    let arg_ty = self.fresh.fresh_type();
                    let (s, bindings) = self.check_pattern(arg, &arg_ty);
                    subst = s.compose(&subst);
                    all_bindings.extend(bindings);
                }
                (subst, all_bindings)
            }
        }
    }

    // -----------------------------------------------------------------------
    // Field access inference
    // -----------------------------------------------------------------------

    fn infer_field_access(
        &mut self,
        env: &mut TypeEnv,
        inner: &RExpr,
        fields: &[String],
        span: Span,
    ) -> (Substitution, TExpr) {
        let (mut subst, current_typed) = self.infer(env, inner);
        let mut current_ty = current_typed.ty.apply_subst(&subst);

        // Chain field accesses
        let mut accessed_fields = Vec::new();
        for field in fields {
            let field_ty = self.fresh.fresh_type();
            let row = self.fresh.fresh();

            // Build a record type with at least this field
            let mut required_fields = BTreeMap::new();
            required_fields.insert(field.clone(), field_ty.clone());
            let record_ty = Type::Record(required_fields, Some(row));

            match unify(&current_ty, &record_ty) {
                Ok(s) => {
                    subst = s.compose(&subst);
                    current_ty = field_ty.apply_subst(&subst);
                }
                Err(err) => {
                    self.errors.push(
                        err.with_context(ErrorContext::InFieldAccess {
                            field: field.clone(),
                        }),
                    );
                    current_ty = self.fresh.fresh_type();
                }
            }

            accessed_fields.push(field.clone());
        }

        (
            subst,
            TExpr {
                kind: TExprKind::FieldAccess {
                    expr: Box::new(current_typed),
                    fields: accessed_fields,
                },
                ty: current_ty,
                span,
            },
        )
    }

    // -----------------------------------------------------------------------
    // Collection literals
    // -----------------------------------------------------------------------

    fn infer_list(
        &mut self,
        env: &mut TypeEnv,
        elems: &[RExpr],
        span: Span,
    ) -> (Substitution, TExpr) {
        let elem_ty = self.fresh.fresh_type();
        let mut subst = Substitution::empty();
        let mut typed_elems = Vec::new();

        for (i, elem) in elems.iter().enumerate() {
            self.apply_subst_to_env(env, &subst);
            let (s_elem, elem_typed) = self.infer(env, elem);
            subst = s_elem.compose(&subst);

            // Unify each element with the expected element type
            let actual = elem_typed.ty.apply_subst(&subst);
            let expected = elem_ty.apply_subst(&subst);
            match unify(&actual, &expected) {
                Ok(s) => subst = s.compose(&subst),
                Err(err) => {
                    self.errors.push(
                        err.with_context(ErrorContext::InListElement { index: i }),
                    );
                }
            }

            typed_elems.push(elem_typed);
        }

        let final_elem_ty = elem_ty.apply_subst(&subst);
        (
            subst,
            TExpr {
                kind: TExprKind::ListLit(typed_elems),
                ty: Type::list(final_elem_ty),
                span,
            },
        )
    }

    fn infer_record(
        &mut self,
        env: &mut TypeEnv,
        fields: &[(String, RExpr)],
        span: Span,
    ) -> (Substitution, TExpr) {
        let mut subst = Substitution::empty();
        let mut typed_fields = Vec::new();
        let mut field_types = BTreeMap::new();

        for (name, value) in fields {
            self.apply_subst_to_env(env, &subst);
            let (s_val, val_typed) = self.infer(env, value);
            subst = s_val.compose(&subst);

            let ty = val_typed.ty.apply_subst(&subst);
            field_types.insert(name.clone(), ty);
            typed_fields.push((name.clone(), val_typed));
        }

        (
            subst,
            TExpr {
                kind: TExprKind::RecordLit(typed_fields),
                ty: Type::Record(field_types, None),
                span,
            },
        )
    }

    fn infer_tuple(
        &mut self,
        env: &mut TypeEnv,
        elems: &[RExpr],
        span: Span,
    ) -> (Substitution, TExpr) {
        let mut subst = Substitution::empty();
        let mut typed_elems = Vec::new();
        let mut elem_types = Vec::new();

        for elem in elems {
            self.apply_subst_to_env(env, &subst);
            let (s_elem, elem_typed) = self.infer(env, elem);
            subst = s_elem.compose(&subst);

            let ty = elem_typed.ty.apply_subst(&subst);
            elem_types.push(ty);
            typed_elems.push(elem_typed);
        }

        (
            subst,
            TExpr {
                kind: TExprKind::TupleLit(typed_elems),
                ty: Type::Tuple(elem_types),
                span,
            },
        )
    }

    // -----------------------------------------------------------------------
    // Constructor inference
    // -----------------------------------------------------------------------

    fn infer_constructor(
        &mut self,
        env: &mut TypeEnv,
        name: &str,
        args: &[RExpr],
        span: Span,
    ) -> (Substitution, TExpr) {
        let mut subst = Substitution::empty();
        let mut typed_args = Vec::new();

        for arg in args {
            self.apply_subst_to_env(env, &subst);
            let (s_arg, arg_typed) = self.infer(env, arg);
            subst = s_arg.compose(&subst);
            typed_args.push(arg_typed);
        }

        let result_ty = match name {
            "Just" => {
                if typed_args.len() == 1 {
                    let inner = typed_args[0].ty.apply_subst(&subst);
                    Type::maybe(inner)
                } else {
                    self.fresh.fresh_type()
                }
            }
            "Left" => {
                if typed_args.len() == 1 {
                    let a = typed_args[0].ty.apply_subst(&subst);
                    let b = self.fresh.fresh_type();
                    Type::either(a, b)
                } else {
                    self.fresh.fresh_type()
                }
            }
            "Right" => {
                if typed_args.len() == 1 {
                    let a = self.fresh.fresh_type();
                    let b = typed_args[0].ty.apply_subst(&subst);
                    Type::either(a, b)
                } else {
                    self.fresh.fresh_type()
                }
            }
            "Cons" => {
                if typed_args.len() == 2 {
                    let elem = typed_args[0].ty.apply_subst(&subst);
                    let tail = typed_args[1].ty.apply_subst(&subst);
                    let expected_tail = Type::list(elem.clone());
                    match unify(&tail, &expected_tail) {
                        Ok(s) => {
                            subst = s.compose(&subst);
                            Type::list(elem.apply_subst(&subst))
                        }
                        Err(err) => {
                            self.errors.push(err);
                            Type::list(elem)
                        }
                    }
                } else {
                    self.fresh.fresh_type()
                }
            }
            _ => {
                // Unknown constructor — return a fresh type
                self.fresh.fresh_type()
            }
        };

        (
            subst,
            TExpr {
                kind: TExprKind::Constructor {
                    name: name.to_string(),
                    args: typed_args,
                },
                ty: result_ty,
                span,
            },
        )
    }

    // -----------------------------------------------------------------------
    // Instantiation & Generalization
    // -----------------------------------------------------------------------

    /// Instantiate a type scheme by replacing all quantified variables
    /// with fresh unification variables.
    fn instantiate(&mut self, scheme: &TypeScheme) -> Type {
        if scheme.vars.is_empty() {
            return scheme.ty.clone();
        }
        let mut subst = Substitution::empty();
        for v in &scheme.vars {
            subst.insert(*v, self.fresh.fresh_type());
        }
        scheme.ty.apply_subst(&subst)
    }

    /// Generalize a type into a type scheme by quantifying over
    /// variables free in the type but not in the environment.
    fn generalize(&self, env: &TypeEnv, ty: &Type) -> TypeScheme {
        let env_fv = env.free_vars();
        let ty_fv = ty.free_vars();
        let vars: Vec<TypeVar> = ty_fv
            .difference(&env_fv)
            .cloned()
            .collect();
        if vars.is_empty() {
            TypeScheme::mono(ty.clone())
        } else {
            TypeScheme {
                vars,
                ty: ty.clone(),
            }
        }
    }

    // -----------------------------------------------------------------------
    // Helpers
    // -----------------------------------------------------------------------

    /// Apply a substitution to all types in the environment.
    fn apply_subst_to_env(&self, env: &mut TypeEnv, subst: &Substitution) {
        if subst.is_empty() {
            return;
        }
        for (_, scheme) in &mut env.bindings {
            *scheme = scheme.apply_subst(subst);
        }
    }

    /// Recursively apply a substitution to all types in a TExpr.
    fn apply_subst_to_texpr(&self, expr: &TExpr, subst: &Substitution) -> TExpr {
        if subst.is_empty() {
            return expr.clone();
        }
        let ty = expr.ty.apply_subst(subst);
        let kind = match &expr.kind {
            TExprKind::BoundVar { index, name } => TExprKind::BoundVar {
                index: *index,
                name: name.clone(),
            },
            TExprKind::CatalogRef { hash, name } => TExprKind::CatalogRef {
                hash: hash.clone(),
                name: name.clone(),
            },
            TExprKind::FieldAccess { expr: inner, fields } => TExprKind::FieldAccess {
                expr: Box::new(self.apply_subst_to_texpr(inner, subst)),
                fields: fields.clone(),
            },
            TExprKind::IntLit(n) => TExprKind::IntLit(*n),
            TExprKind::FloatLit(f) => TExprKind::FloatLit(*f),
            TExprKind::StringLit(s) => TExprKind::StringLit(s.clone()),
            TExprKind::BoolLit(b) => TExprKind::BoolLit(*b),
            TExprKind::NothingLit => TExprKind::NothingLit,
            TExprKind::Lambda { params, body } => TExprKind::Lambda {
                params: params
                    .iter()
                    .map(|(n, t)| (n.clone(), t.apply_subst(subst)))
                    .collect(),
                body: Box::new(self.apply_subst_to_texpr(body, subst)),
            },
            TExprKind::Apply { func, args } => TExprKind::Apply {
                func: Box::new(self.apply_subst_to_texpr(func, subst)),
                args: args
                    .iter()
                    .map(|a| self.apply_subst_to_texpr(a, subst))
                    .collect(),
            },
            TExprKind::Let { name, value, body } => TExprKind::Let {
                name: name.clone(),
                value: Box::new(self.apply_subst_to_texpr(value, subst)),
                body: Box::new(self.apply_subst_to_texpr(body, subst)),
            },
            TExprKind::Match { scrutinee, arms } => TExprKind::Match {
                scrutinee: Box::new(self.apply_subst_to_texpr(scrutinee, subst)),
                arms: arms
                    .iter()
                    .map(|arm| TypedMatchArm {
                        pattern: arm.pattern.clone(),
                        body: self.apply_subst_to_texpr(&arm.body, subst),
                    })
                    .collect(),
            },
            TExprKind::ListLit(elems) => TExprKind::ListLit(
                elems
                    .iter()
                    .map(|e| self.apply_subst_to_texpr(e, subst))
                    .collect(),
            ),
            TExprKind::RecordLit(fields) => TExprKind::RecordLit(
                fields
                    .iter()
                    .map(|(n, e)| (n.clone(), self.apply_subst_to_texpr(e, subst)))
                    .collect(),
            ),
            TExprKind::TupleLit(elems) => TExprKind::TupleLit(
                elems
                    .iter()
                    .map(|e| self.apply_subst_to_texpr(e, subst))
                    .collect(),
            ),
            TExprKind::Constructor { name, args } => TExprKind::Constructor {
                name: name.clone(),
                args: args
                    .iter()
                    .map(|a| self.apply_subst_to_texpr(a, subst))
                    .collect(),
            },
        };
        TExpr {
            kind,
            ty,
            span: expr.span,
        }
    }
}
