//! Type checker for the Lux language

#![allow(dead_code, unused_variables)]

use std::collections::HashMap;

use crate::ast::{
    self, BinaryOp, Declaration, EffectDecl, Expr, FunctionDecl, HandlerDecl, Ident, ImplDecl,
    ImportDecl, LetDecl, Literal, LiteralKind, MatchArm, Parameter, Pattern, Program, Span,
    Statement, TraitDecl, TypeDecl, TypeExpr, UnaryOp, VariantFields,
};
use crate::diagnostics::{find_similar_names, format_did_you_mean, Diagnostic, Severity};
use crate::exhaustiveness::{check_exhaustiveness, missing_patterns_hint};
use crate::modules::ModuleLoader;
use crate::types::{
    self, unify, EffectDef, EffectOpDef, EffectSet, HandlerDef, PropertySet, TraitBoundDef,
    TraitDef, TraitImpl, TraitMethodDef, Type, TypeEnv, TypeScheme, VariantDef, VariantFieldsDef,
};

/// Type checking error
#[derive(Debug, Clone)]
pub struct TypeError {
    pub message: String,
    pub span: Span,
}

impl std::fmt::Display for TypeError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Type error at {}-{}: {}",
            self.span.start, self.span.end, self.message
        )
    }
}

impl TypeError {
    /// Convert to a rich diagnostic for Elm-style error display
    pub fn to_diagnostic(&self) -> Diagnostic {
        // Categorize the error and extract hints
        let (title, hints) = categorize_type_error(&self.message);

        Diagnostic {
            severity: Severity::Error,
            title,
            message: self.message.clone(),
            span: self.span,
            hints,
        }
    }
}

/// Categorize a type error message to provide better titles and hints
fn categorize_type_error(message: &str) -> (String, Vec<String>) {
    let message_lower = message.to_lowercase();

    if message_lower.contains("type mismatch") {
        (
            "Type Mismatch".to_string(),
            vec!["Check that the types on both sides of the expression are compatible.".to_string()],
        )
    } else if message_lower.contains("undefined variable") || message_lower.contains("not found") {
        (
            "Unknown Name".to_string(),
            vec![
                "Check the spelling of the name.".to_string(),
                "Make sure the variable is defined before use.".to_string(),
            ],
        )
    } else if message_lower.contains("cannot unify") {
        (
            "Type Mismatch".to_string(),
            vec!["The types are not compatible. Check your function arguments and return types.".to_string()],
        )
    } else if message_lower.contains("expected") && message_lower.contains("argument") {
        (
            "Wrong Number of Arguments".to_string(),
            vec!["Check the function signature and provide the correct number of arguments.".to_string()],
        )
    } else if message_lower.contains("pure") && message_lower.contains("effect") {
        (
            "Purity Violation".to_string(),
            vec![
                "Functions marked 'is pure' cannot perform effects.".to_string(),
                "Remove the 'is pure' annotation or handle the effects.".to_string(),
            ],
        )
    } else if message_lower.contains("effect") && message_lower.contains("unhandled") {
        (
            "Unhandled Effect".to_string(),
            vec!["Use a 'handle' expression to provide an implementation for this effect.".to_string()],
        )
    } else if message_lower.contains("recursive") {
        (
            "Invalid Recursion".to_string(),
            vec!["Check that recursive calls have proper base cases.".to_string()],
        )
    } else {
        ("Type Error".to_string(), vec![])
    }
}

/// Type checker
pub struct TypeChecker {
    env: TypeEnv,
    current_effects: EffectSet,
    /// Effects inferred from the current function body (for effect inference)
    inferred_effects: EffectSet,
    /// Whether we're inferring effects (no explicit declaration)
    inferring_effects: bool,
    errors: Vec<TypeError>,
    /// Type parameters in scope (maps "T" -> Type::Var(n) for generics)
    type_params: HashMap<String, Type>,
}

impl TypeChecker {
    pub fn new() -> Self {
        Self {
            env: TypeEnv::with_builtins(),
            current_effects: EffectSet::empty(),
            inferred_effects: EffectSet::empty(),
            inferring_effects: false,
            errors: Vec::new(),
            type_params: HashMap::new(),
        }
    }

    /// Look up a type scheme by name (for REPL :info command)
    pub fn lookup(&self, name: &str) -> Option<&TypeScheme> {
        self.env.bindings.get(name)
    }

    /// Type check a program
    pub fn check_program(&mut self, program: &Program) -> Result<(), Vec<TypeError>> {
        // First pass: collect all declarations
        for decl in &program.declarations {
            self.collect_declaration(decl);
        }

        // Second pass: type check all declarations
        for decl in &program.declarations {
            self.check_declaration(decl);
        }

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

    /// Type check a program with module support
    pub fn check_program_with_modules(
        &mut self,
        program: &Program,
        loader: &ModuleLoader,
    ) -> Result<(), Vec<TypeError>> {
        // First, process imports
        self.load_type_imports(&program.imports, loader);

        // Then check the program normally
        self.check_program(program)
    }

    /// Load type bindings from imported modules
    fn load_type_imports(&mut self, imports: &[ImportDecl], loader: &ModuleLoader) {
        use crate::modules::ImportKind;

        let resolved = match loader.resolve_imports(imports) {
            Ok(r) => r,
            Err(e) => {
                self.errors.push(TypeError {
                    message: format!("Import error: {}", e.message),
                    span: Span { start: 0, end: 0 },
                });
                return;
            }
        };

        for (name, import) in resolved {
            match import.kind {
                ImportKind::Module => {
                    // Import as a module object - create a record type
                    let module = match loader.get_module(&import.module_path) {
                        Some(m) => m,
                        None => continue,
                    };

                    // Create a temporary checker to get types from the module
                    let mut module_checker = TypeChecker::new();
                    for decl in &module.program.declarations {
                        module_checker.collect_declaration(decl);
                    }

                    // Build a record type with all exported names
                    let mut fields = Vec::new();
                    for export_name in &module.exports {
                        if let Some(scheme) = module_checker.env.lookup(export_name) {
                            fields.push((export_name.clone(), scheme.instantiate()));
                        }
                    }

                    if !fields.is_empty() {
                        self.env.bind(&name, TypeScheme::mono(Type::Record(fields)));
                    }
                }
                ImportKind::Direct => {
                    // Import a specific name directly
                    let module = match loader.get_module(&import.module_path) {
                        Some(m) => m,
                        None => continue,
                    };

                    // Get the type for this specific export
                    let mut module_checker = TypeChecker::new();
                    for decl in &module.program.declarations {
                        module_checker.collect_declaration(decl);
                    }

                    if let Some(scheme) = module_checker.env.lookup(&import.name) {
                        self.env.bind(&name, scheme.clone());
                    }

                    // Also copy effects if the imported name is an effect
                    if let Some(effect) = module_checker.env.lookup_effect(&import.name) {
                        self.env.effects.insert(name.clone(), effect.clone());
                    }

                    // Also copy types if the imported name is a type
                    if let Some(type_def) = module_checker.env.types.get(&import.name) {
                        self.env.types.insert(name.clone(), type_def.clone());
                    }
                }
            }
        }
    }

    /// Collect type signatures from declarations (first pass)
    fn collect_declaration(&mut self, decl: &Declaration) {
        match decl {
            Declaration::Function(func) => {
                let scheme = self.function_type(func);
                self.env.bind(&func.name.name, scheme);
            }
            Declaration::Effect(effect) => {
                let effect_def = self.effect_def(effect);
                self.env
                    .effects
                    .insert(effect.name.name.clone(), effect_def);
            }
            Declaration::Type(type_decl) => {
                // Save old type params for this scope
                let old_params = std::mem::take(&mut self.type_params);

                // Bind type parameters to fresh type variables
                let mut bound_vars = Vec::new();
                for param in &type_decl.type_params {
                    let var = Type::var();
                    if let Type::Var(n) = &var {
                        bound_vars.push(*n);
                    }
                    self.type_params.insert(param.name.clone(), var);
                }

                // Build the parameterized return type
                let base_type = if type_decl.type_params.is_empty() {
                    Type::Named(type_decl.name.name.clone())
                } else {
                    Type::App {
                        constructor: Box::new(Type::Named(type_decl.name.name.clone())),
                        args: type_decl
                            .type_params
                            .iter()
                            .map(|p| self.type_params.get(&p.name).unwrap().clone())
                            .collect(),
                    }
                };

                // Register the type definition
                let type_def = self.type_def(type_decl);
                self.env.types.insert(type_decl.name.name.clone(), type_def.clone());

                // Register ADT constructors as values with polymorphic types
                if let ast::TypeDef::Enum(variants) = &type_decl.definition {
                    for variant in variants {
                        let constructor_type = match &variant.fields {
                            VariantFields::Unit => {
                                // Unit variant is just the type itself
                                base_type.clone()
                            }
                            VariantFields::Tuple(field_types) => {
                                // Tuple variant is a function from fields to the type
                                let param_types: Vec<Type> = field_types
                                    .iter()
                                    .map(|t| self.resolve_type(t))
                                    .collect();
                                Type::function(param_types, base_type.clone())
                            }
                            VariantFields::Record(fields) => {
                                // Record variant is a function from record to the type
                                let field_types: Vec<(String, Type)> = fields
                                    .iter()
                                    .map(|f| (f.name.name.clone(), self.resolve_type(&f.typ)))
                                    .collect();
                                Type::function(vec![Type::Record(field_types)], base_type.clone())
                            }
                        };
                        // Wrap in polymorphic TypeScheme for generic types
                        let scheme = TypeScheme {
                            vars: bound_vars.clone(),
                            typ: constructor_type,
                        };
                        self.env.bind(&variant.name.name, scheme);
                    }
                }

                // Restore old type params
                self.type_params = old_params;
            }
            Declaration::Handler(handler) => {
                let handler_def = self.handler_def(handler);
                self.env
                    .handlers
                    .insert(handler.name.name.clone(), handler_def);
                // Also bind the handler as a value so it can be referenced in run...with expressions
                // Handler type is the effect it handles (as an opaque type for now)
                let handler_type = Type::Named(format!("Handler<{}>", handler.effect.name));
                self.env.bind(&handler.name.name, TypeScheme::mono(handler_type));
            }
            Declaration::Let(let_decl) => {
                // Will be typed in second pass
                let typ = if let Some(ref type_expr) = let_decl.typ {
                    self.resolve_type(type_expr)
                } else {
                    Type::var()
                };
                self.env.bind(&let_decl.name.name, TypeScheme::mono(typ));
            }
            Declaration::Trait(trait_decl) => {
                let trait_def = self.trait_def(trait_decl);
                self.env.traits.insert(trait_decl.name.name.clone(), trait_def);
            }
            Declaration::Impl(impl_decl) => {
                // Will be checked in second pass
                let trait_impl = self.collect_impl(impl_decl);
                self.env.trait_impls.push(trait_impl);
            }
        }
    }

    /// Type check a declaration (second pass)
    fn check_declaration(&mut self, decl: &Declaration) {
        match decl {
            Declaration::Function(func) => {
                self.check_function(func);
            }
            Declaration::Let(let_decl) => {
                self.check_let_decl(let_decl);
            }
            Declaration::Handler(handler) => {
                self.check_handler(handler);
            }
            Declaration::Impl(impl_decl) => {
                self.check_impl(impl_decl);
            }
            // Effects, types, and traits don't need checking beyond collection
            _ => {}
        }
    }

    fn check_function(&mut self, func: &FunctionDecl) {
        // Set up the environment with parameters
        let mut local_env = self.env.clone();
        for param in &func.params {
            let param_type = self.resolve_type(&param.typ);
            local_env.bind(&param.name.name, TypeScheme::mono(param_type));
        }

        // Determine if we need to infer effects
        let explicit_effects = !func.effects.is_empty();
        let declared_effects = EffectSet::from_iter(func.effects.iter().map(|e| e.name.clone()));

        // Save old state
        let old_effects = std::mem::replace(&mut self.current_effects, declared_effects.clone());
        let old_inferring = std::mem::replace(&mut self.inferring_effects, !explicit_effects);
        let old_inferred = std::mem::replace(&mut self.inferred_effects, EffectSet::empty());

        // Type check the body
        let old_env = std::mem::replace(&mut self.env, local_env);
        let body_type = self.infer_expr(&func.body);
        self.env = old_env;

        // Get the inferred effects before restoring state
        let inferred = std::mem::replace(&mut self.inferred_effects, old_inferred);

        // Restore state
        self.current_effects = old_effects;
        self.inferring_effects = old_inferring;

        // Check that body type matches return type
        let return_type = self.resolve_type(&func.return_type);
        if let Err(e) = unify(&body_type, &return_type) {
            self.errors.push(TypeError {
                message: format!(
                    "Function '{}' body has type {}, but declared return type is {}: {}",
                    func.name.name, body_type, return_type, e
                ),
                span: func.span,
            });
        }

        // Check behavioral properties
        let properties = PropertySet::from_ast(&func.properties);

        // Pure functions cannot have effects
        let effective_effects = if explicit_effects {
            &declared_effects
        } else {
            &inferred
        };

        if properties.is_pure() && !effective_effects.is_empty() {
            let effects_str = if explicit_effects {
                func.effects
                    .iter()
                    .map(|e| e.name.as_str())
                    .collect::<Vec<_>>()
                    .join(", ")
            } else {
                format!("{} (inferred)", inferred)
            };
            self.errors.push(TypeError {
                message: format!(
                    "Function '{}' is declared as pure but has effects: {{{}}}",
                    func.name.name, effects_str
                ),
                span: func.span,
            });
        }

        // If effects were declared, verify that inferred effects are a subset
        if explicit_effects && !inferred.is_subset(&declared_effects) {
            let missing: Vec<_> = inferred
                .effects
                .iter()
                .filter(|e| !declared_effects.contains(e))
                .cloned()
                .collect();
            self.errors.push(TypeError {
                message: format!(
                    "Function '{}' uses effects {{{}}} but only declares {{{}}}",
                    func.name.name,
                    missing.join(", "),
                    declared_effects
                ),
                span: func.span,
            });
        }

        // Check where clause property constraints
        for where_clause in &func.where_clauses {
            match where_clause {
                ast::WhereClause::PropertyConstraint {
                    type_param,
                    property,
                    span,
                } => {
                    // Record the constraint for later checking when the function is called
                    // For now, we just validate that the type parameter exists
                    if !func.type_params.iter().any(|p| p.name == type_param.name)
                        && !func.params.iter().any(|p| p.name.name == type_param.name)
                    {
                        self.errors.push(TypeError {
                            message: format!(
                                "Unknown type parameter '{}' in where clause",
                                type_param.name
                            ),
                            span: *span,
                        });
                    }
                }
                ast::WhereClause::ResultRefinement { predicate, span } => {
                    // Result refinements are checked at runtime or via SMT solver
                    // For now, we just type-check the predicate expression
                    // (would need 'result' in scope, which we don't have yet)
                }
                ast::WhereClause::TraitConstraint(constraint) => {
                    // Validate that the type parameter exists
                    if !func.type_params.iter().any(|p| p.name == constraint.type_param.name)
                        && !func.params.iter().any(|p| p.name.name == constraint.type_param.name)
                    {
                        self.errors.push(TypeError {
                            message: format!(
                                "Unknown type parameter '{}' in where clause",
                                constraint.type_param.name
                            ),
                            span: constraint.span,
                        });
                    }
                    // Validate that each trait in the bounds exists
                    for bound in &constraint.bounds {
                        if !self.env.traits.contains_key(&bound.trait_name.name) {
                            // Find similar trait names for suggestion
                            let available_traits: Vec<&str> = self.env.traits.keys()
                                .map(|s| s.as_str())
                                .collect();
                            let suggestions = find_similar_names(&bound.trait_name.name, available_traits, 2);

                            let mut message = format!("Unknown trait: {}", bound.trait_name.name);
                            if let Some(hint) = format_did_you_mean(&suggestions) {
                                message.push_str(&format!(". {}", hint));
                            }

                            self.errors.push(TypeError {
                                message,
                                span: bound.span,
                            });
                        }
                    }
                }
            }
        }
    }

    fn check_let_decl(&mut self, let_decl: &LetDecl) {
        let inferred = self.infer_expr(&let_decl.value);

        if let Some(ref type_expr) = let_decl.typ {
            let declared = self.resolve_type(type_expr);
            if let Err(e) = unify(&inferred, &declared) {
                self.errors.push(TypeError {
                    message: format!(
                        "Variable '{}' has type {}, but declared type is {}: {}",
                        let_decl.name.name, inferred, declared, e
                    ),
                    span: let_decl.span,
                });
            }
        }

        // Update the binding with the inferred type
        let scheme = self.env.generalize(&inferred);
        self.env.bind(&let_decl.name.name, scheme);
    }

    fn check_handler(&mut self, handler: &HandlerDecl) {
        // Look up the effect being handled
        let effect = match self.env.lookup_effect(&handler.effect.name) {
            Some(e) => e.clone(),
            None => {
                // Find similar effect names for suggestion
                let available_effects: Vec<&str> = self.env.effects.keys()
                    .map(|s| s.as_str())
                    .collect();
                let suggestions = find_similar_names(&handler.effect.name, available_effects, 2);

                let mut message = format!("Unknown effect: {}", handler.effect.name);
                if let Some(hint) = format_did_you_mean(&suggestions) {
                    message.push_str(&format!(". {}", hint));
                }

                self.errors.push(TypeError {
                    message,
                    span: handler.effect.span,
                });
                return;
            }
        };

        // Check each implementation
        for impl_ in &handler.implementations {
            let op = match effect
                .operations
                .iter()
                .find(|o| o.name == impl_.op_name.name)
            {
                Some(o) => o.clone(),
                None => {
                    self.errors.push(TypeError {
                        message: format!(
                            "Effect '{}' has no operation '{}'",
                            handler.effect.name, impl_.op_name.name
                        ),
                        span: impl_.op_name.span,
                    });
                    continue;
                }
            };

            // Set up environment with operation parameters
            let mut local_env = self.env.clone();
            for (i, param_name) in impl_.params.iter().enumerate() {
                if i < op.params.len() {
                    local_env.bind(&param_name.name, TypeScheme::mono(op.params[i].1.clone()));
                }
            }

            // Add resume if present
            if let Some(ref resume) = impl_.resume {
                // resume has type: fn(return_type) -> result_type
                let resume_type = Type::function(vec![op.return_type.clone()], Type::var());
                local_env.bind(&resume.name, TypeScheme::mono(resume_type));
            }

            // Type check the implementation body
            let old_env = std::mem::replace(&mut self.env, local_env);
            let _body_type = self.infer_expr(&impl_.body);
            self.env = old_env;
        }
    }

    /// Infer the type of an expression
    fn infer_expr(&mut self, expr: &Expr) -> Type {
        match expr {
            Expr::Literal(lit) => self.infer_literal(lit),

            Expr::Var(ident) => match self.env.lookup(&ident.name) {
                Some(scheme) => scheme.instantiate(),
                None => {
                    // Find similar variable names for "Did you mean?" suggestion
                    let available_names: Vec<&str> = self.env.bindings.keys()
                        .map(|s| s.as_str())
                        .collect();
                    let suggestions = find_similar_names(&ident.name, available_names, 2);

                    let mut message = format!("Undefined variable: {}", ident.name);
                    if let Some(hint) = format_did_you_mean(&suggestions) {
                        message.push_str(&format!(". {}", hint));
                    }

                    self.errors.push(TypeError {
                        message,
                        span: ident.span,
                    });
                    Type::Error
                }
            },

            Expr::BinaryOp {
                op,
                left,
                right,
                span,
            } => self.infer_binary_op(*op, left, right, *span),

            Expr::UnaryOp { op, operand, span } => self.infer_unary_op(*op, operand, *span),

            Expr::Call { func, args, span } => self.infer_call(func, args, *span),

            Expr::EffectOp {
                effect,
                operation,
                args,
                span,
            } => self.infer_effect_op(effect, operation, args, *span),

            Expr::Field {
                object,
                field,
                span,
            } => self.infer_field(object, field, *span),

            Expr::Lambda {
                params,
                return_type,
                effects,
                body,
                span,
            } => self.infer_lambda(params, return_type.as_deref(), effects, body, *span),

            Expr::Let {
                name,
                typ,
                value,
                body,
                span,
            } => self.infer_let(name, typ.as_ref(), value, body, *span),

            Expr::If {
                condition,
                then_branch,
                else_branch,
                span,
            } => self.infer_if(condition, then_branch, else_branch, *span),

            Expr::Match {
                scrutinee,
                arms,
                span,
            } => self.infer_match(scrutinee, arms, *span),

            Expr::Block {
                statements,
                result,
                span,
            } => self.infer_block(statements, result, *span),

            Expr::Record { fields, span } => self.infer_record(fields, *span),

            Expr::Tuple { elements, span } => self.infer_tuple(elements, *span),

            Expr::List { elements, span } => self.infer_list(elements, *span),

            Expr::Run {
                expr,
                handlers,
                span,
            } => self.infer_run(expr, handlers, *span),

            Expr::Resume { value, span } => {
                // Resume is special - it continues the computation
                // For now, just infer the value type
                let _ = self.infer_expr(value);
                Type::var() // Resume's type depends on context
            }
        }
    }

    fn infer_literal(&self, lit: &Literal) -> Type {
        match &lit.kind {
            LiteralKind::Int(_) => Type::Int,
            LiteralKind::Float(_) => Type::Float,
            LiteralKind::String(_) => Type::String,
            LiteralKind::Char(_) => Type::Char,
            LiteralKind::Bool(_) => Type::Bool,
            LiteralKind::Unit => Type::Unit,
        }
    }

    fn infer_binary_op(&mut self, op: BinaryOp, left: &Expr, right: &Expr, span: Span) -> Type {
        let left_type = self.infer_expr(left);
        let right_type = self.infer_expr(right);

        match op {
            BinaryOp::Add => {
                // Add supports both numeric types and string concatenation
                if let Err(e) = unify(&left_type, &right_type) {
                    self.errors.push(TypeError {
                        message: format!("Operands of '{}' must have same type: {}", op, e),
                        span,
                    });
                }
                match &left_type {
                    Type::Int | Type::Float | Type::String | Type::Var(_) => left_type,
                    _ => {
                        self.errors.push(TypeError {
                            message: format!(
                                "Operator '{}' requires numeric or string operands, got {}",
                                op, left_type
                            ),
                            span,
                        });
                        Type::Error
                    }
                }
            }

            BinaryOp::Sub | BinaryOp::Mul | BinaryOp::Div | BinaryOp::Mod => {
                // Arithmetic: both operands must be same numeric type
                if let Err(e) = unify(&left_type, &right_type) {
                    self.errors.push(TypeError {
                        message: format!("Operands of '{}' must have same type: {}", op, e),
                        span,
                    });
                }
                // Check that they're numeric
                match &left_type {
                    Type::Int | Type::Float | Type::Var(_) => left_type,
                    _ => {
                        self.errors.push(TypeError {
                            message: format!(
                                "Operator '{}' requires numeric operands, got {}",
                                op, left_type
                            ),
                            span,
                        });
                        Type::Error
                    }
                }
            }

            BinaryOp::Eq | BinaryOp::Ne => {
                // Equality: operands must have same type
                if let Err(e) = unify(&left_type, &right_type) {
                    self.errors.push(TypeError {
                        message: format!("Operands of '{}' must have same type: {}", op, e),
                        span,
                    });
                }
                Type::Bool
            }

            BinaryOp::Lt | BinaryOp::Le | BinaryOp::Gt | BinaryOp::Ge => {
                // Comparison: operands must be same orderable type
                if let Err(e) = unify(&left_type, &right_type) {
                    self.errors.push(TypeError {
                        message: format!("Operands of '{}' must have same type: {}", op, e),
                        span,
                    });
                }
                Type::Bool
            }

            BinaryOp::And | BinaryOp::Or => {
                // Logical: both must be Bool
                if let Err(e) = unify(&left_type, &Type::Bool) {
                    self.errors.push(TypeError {
                        message: format!("Left operand of '{}' must be Bool: {}", op, e),
                        span: left.span(),
                    });
                }
                if let Err(e) = unify(&right_type, &Type::Bool) {
                    self.errors.push(TypeError {
                        message: format!("Right operand of '{}' must be Bool: {}", op, e),
                        span: right.span(),
                    });
                }
                Type::Bool
            }

            BinaryOp::Pipe => {
                // Pipe: a |> f means f(a)
                // right must be a function that accepts left's type
                let result_type = Type::var();
                let expected_fn = Type::function(vec![left_type.clone()], result_type.clone());
                if let Err(e) = unify(&right_type, &expected_fn) {
                    self.errors.push(TypeError {
                        message: format!(
                            "Pipe target must be a function accepting {}: {}",
                            left_type, e
                        ),
                        span,
                    });
                }
                result_type
            }
        }
    }

    fn infer_unary_op(&mut self, op: UnaryOp, operand: &Expr, span: Span) -> Type {
        let operand_type = self.infer_expr(operand);

        match op {
            UnaryOp::Neg => match &operand_type {
                Type::Int | Type::Float | Type::Var(_) => operand_type,
                _ => {
                    self.errors.push(TypeError {
                        message: format!(
                            "Operator '-' requires numeric operand, got {}",
                            operand_type
                        ),
                        span,
                    });
                    Type::Error
                }
            },
            UnaryOp::Not => {
                if let Err(e) = unify(&operand_type, &Type::Bool) {
                    self.errors.push(TypeError {
                        message: format!("Operator '!' requires Bool operand: {}", e),
                        span,
                    });
                }
                Type::Bool
            }
        }
    }

    fn infer_call(&mut self, func: &Expr, args: &[Expr], span: Span) -> Type {
        let func_type = self.infer_expr(func);
        let arg_types: Vec<Type> = args.iter().map(|a| self.infer_expr(a)).collect();

        let result_type = Type::var();
        // Include current effects in the expected function type
        // This allows calling functions that require effects when those effects are available
        let expected_fn = Type::function_with_effects(
            arg_types.clone(),
            result_type.clone(),
            self.current_effects.clone(),
        );

        match unify(&func_type, &expected_fn) {
            Ok(subst) => result_type.apply(&subst),
            Err(e) => {
                self.errors.push(TypeError {
                    message: format!("Type mismatch in function call: {}", e),
                    span,
                });
                Type::Error
            }
        }
    }

    fn infer_effect_op(
        &mut self,
        effect: &Ident,
        operation: &Ident,
        args: &[Expr],
        span: Span,
    ) -> Type {
        // Check if this is actually a module access, not an effect operation
        // This includes stdlib modules (List, String, etc.) and user-imported modules
        if let Some(scheme) = self.env.lookup(&effect.name) {
            let module_type = scheme.instantiate();
            // Get the field type - if it's a Record, treat it as a module
            if let Type::Record(fields) = &module_type {
                if let Some((_, field_type)) = fields.iter().find(|(n, _)| n == &operation.name) {
                    // It's a function call on a module field
                    let arg_types: Vec<Type> = args.iter().map(|a| self.infer_expr(a)).collect();
                    let result_type = Type::var();
                    let expected_fn = Type::function(arg_types, result_type.clone());

                    if let Err(e) = unify(field_type, &expected_fn) {
                        self.errors.push(TypeError {
                            message: format!(
                                "Type mismatch in {}.{} call: {}",
                                effect.name, operation.name, e
                            ),
                            span,
                        });
                    }
                    return result_type;
                } else {
                    self.errors.push(TypeError {
                        message: format!(
                            "Module '{}' has no member '{}'",
                            effect.name, operation.name
                        ),
                        span,
                    });
                    return Type::Error;
                }
            }
            // Fall through to normal error handling if module field not found
        }

        // Built-in effects are always available
        let builtin_effects = ["Console", "Fail", "State"];
        let is_builtin = builtin_effects.contains(&effect.name.as_str());

        // Track this effect for inference
        if self.inferring_effects {
            self.inferred_effects.insert(effect.name.clone());
        }

        // Check that we're in a context that allows this effect
        // Skip this check if we're inferring effects (no explicit declaration)
        if !self.inferring_effects && !is_builtin && !self.current_effects.contains(&effect.name) {
            self.errors.push(TypeError {
                message: format!(
                    "Effect '{}' not available in current context. Available: {{{}}}",
                    effect.name, self.current_effects
                ),
                span,
            });
        }

        // Look up the operation - clone to avoid borrow issues
        let op = self
            .env
            .lookup_effect_op(&effect.name, &operation.name)
            .cloned();

        match op {
            Some(op) => {
                // Check argument types
                let arg_types: Vec<Type> = args.iter().map(|a| self.infer_expr(a)).collect();

                if arg_types.len() != op.params.len() {
                    self.errors.push(TypeError {
                        message: format!(
                            "Effect operation '{}.{}' expects {} arguments, got {}",
                            effect.name,
                            operation.name,
                            op.params.len(),
                            arg_types.len()
                        ),
                        span,
                    });
                }

                for (i, (arg_type, (_, param_type))) in
                    arg_types.iter().zip(op.params.iter()).enumerate()
                {
                    if let Err(e) = unify(arg_type, param_type) {
                        self.errors.push(TypeError {
                            message: format!(
                                "Argument {} of '{}.{}' has type {}, expected {}: {}",
                                i + 1,
                                effect.name,
                                operation.name,
                                arg_type,
                                param_type,
                                e
                            ),
                            span,
                        });
                    }
                }

                op.return_type.clone()
            }
            None => {
                self.errors.push(TypeError {
                    message: format!(
                        "Unknown effect operation: {}.{}",
                        effect.name, operation.name
                    ),
                    span,
                });
                Type::Error
            }
        }
    }

    fn infer_field(&mut self, object: &Expr, field: &Ident, span: Span) -> Type {
        let object_type = self.infer_expr(object);

        match &object_type {
            Type::Record(fields) => match fields.iter().find(|(n, _)| n == &field.name) {
                Some((_, t)) => t.clone(),
                None => {
                    self.errors.push(TypeError {
                        message: format!("Record has no field '{}'", field.name),
                        span,
                    });
                    Type::Error
                }
            },
            Type::Var(_) => {
                // Can't infer field access on unknown type
                Type::var()
            }
            _ => {
                self.errors.push(TypeError {
                    message: format!(
                        "Cannot access field '{}' on non-record type {}",
                        field.name, object_type
                    ),
                    span,
                });
                Type::Error
            }
        }
    }

    fn infer_lambda(
        &mut self,
        params: &[Parameter],
        return_type: Option<&TypeExpr>,
        effects: &[Ident],
        body: &Expr,
        _span: Span,
    ) -> Type {
        // Set up environment with parameters
        let mut local_env = self.env.clone();
        let param_types: Vec<Type> = params
            .iter()
            .map(|p| {
                let t = self.resolve_type(&p.typ);
                local_env.bind(&p.name.name, TypeScheme::mono(t.clone()));
                t
            })
            .collect();

        // Determine if we need to infer effects for this lambda
        let explicit_effects = !effects.is_empty();
        let declared_effects = EffectSet::from_iter(effects.iter().map(|e| e.name.clone()));

        // Save old state
        let old_effects = std::mem::replace(&mut self.current_effects, declared_effects.clone());
        let old_inferring = std::mem::replace(&mut self.inferring_effects, !explicit_effects);
        let old_inferred = std::mem::replace(&mut self.inferred_effects, EffectSet::empty());

        // Type check body
        let old_env = std::mem::replace(&mut self.env, local_env);
        let body_type = self.infer_expr(body);
        self.env = old_env;

        // Get the inferred effects before restoring state
        let inferred = std::mem::replace(&mut self.inferred_effects, old_inferred);

        // Restore state
        self.current_effects = old_effects;
        self.inferring_effects = old_inferring;

        // Check return type if specified
        let ret_type = if let Some(rt) = return_type {
            let declared = self.resolve_type(rt);
            if let Err(e) = unify(&body_type, &declared) {
                self.errors.push(TypeError {
                    message: format!(
                        "Lambda body type {} doesn't match declared {}: {}",
                        body_type, declared, e
                    ),
                    span: body.span(),
                });
            }
            declared
        } else {
            body_type
        };

        // Use inferred effects if not explicitly declared
        let final_effects = if explicit_effects {
            declared_effects
        } else {
            inferred
        };

        Type::function_with_effects(param_types, ret_type, final_effects)
    }

    fn infer_let(
        &mut self,
        name: &Ident,
        typ: Option<&TypeExpr>,
        value: &Expr,
        body: &Expr,
        _span: Span,
    ) -> Type {
        let value_type = self.infer_expr(value);

        // Check declared type if present
        if let Some(type_expr) = typ {
            let declared = self.resolve_type(type_expr);
            if let Err(e) = unify(&value_type, &declared) {
                self.errors.push(TypeError {
                    message: format!(
                        "Variable '{}' has type {}, but declared type is {}: {}",
                        name.name, value_type, declared, e
                    ),
                    span: name.span,
                });
            }
        }

        // Extend environment and check body
        let scheme = self.env.generalize(&value_type);
        let old_env = self.env.clone();
        self.env.bind(&name.name, scheme);
        let body_type = self.infer_expr(body);
        self.env = old_env;

        body_type
    }

    fn infer_if(
        &mut self,
        condition: &Expr,
        then_branch: &Expr,
        else_branch: &Expr,
        span: Span,
    ) -> Type {
        let cond_type = self.infer_expr(condition);
        if let Err(e) = unify(&cond_type, &Type::Bool) {
            self.errors.push(TypeError {
                message: format!("If condition must be Bool, got {}: {}", cond_type, e),
                span: condition.span(),
            });
        }

        let then_type = self.infer_expr(then_branch);
        let else_type = self.infer_expr(else_branch);

        match unify(&then_type, &else_type) {
            Ok(subst) => then_type.apply(&subst),
            Err(e) => {
                self.errors.push(TypeError {
                    message: format!(
                        "If branches have incompatible types: then={}, else={}: {}",
                        then_type, else_type, e
                    ),
                    span,
                });
                Type::Error
            }
        }
    }

    fn infer_match(&mut self, scrutinee: &Expr, arms: &[MatchArm], span: Span) -> Type {
        let scrutinee_type = self.infer_expr(scrutinee);

        if arms.is_empty() {
            self.errors.push(TypeError {
                message: "Match expression must have at least one arm".to_string(),
                span,
            });
            return Type::Error;
        }

        let mut result_type: Option<Type> = None;

        for arm in arms {
            // Check pattern and get bindings
            let bindings = self.check_pattern(&arm.pattern, &scrutinee_type);

            // Extend environment with pattern bindings
            let old_env = self.env.clone();
            for (name, typ) in bindings {
                self.env.bind(name, TypeScheme::mono(typ));
            }

            // Check guard if present
            if let Some(ref guard) = arm.guard {
                let guard_type = self.infer_expr(guard);
                if let Err(e) = unify(&guard_type, &Type::Bool) {
                    self.errors.push(TypeError {
                        message: format!("Match guard must be Bool: {}", e),
                        span: guard.span(),
                    });
                }
            }

            // Check body
            let body_type = self.infer_expr(&arm.body);
            self.env = old_env;

            // Unify with previous arms
            match &result_type {
                None => result_type = Some(body_type),
                Some(prev) => {
                    if let Err(e) = unify(prev, &body_type) {
                        self.errors.push(TypeError {
                            message: format!(
                                "Match arm has incompatible type: expected {}, got {}: {}",
                                prev, body_type, e
                            ),
                            span: arm.span,
                        });
                    }
                }
            }
        }

        // Check exhaustiveness
        let exhaustiveness = check_exhaustiveness(&scrutinee_type, arms, &self.env);

        if !exhaustiveness.is_exhaustive {
            let hint = missing_patterns_hint(&exhaustiveness.missing_patterns);
            self.errors.push(TypeError {
                message: format!(
                    "Non-exhaustive pattern match. {}",
                    hint
                ),
                span,
            });
        }

        // Warn about redundant arms
        for idx in exhaustiveness.redundant_arms {
            self.errors.push(TypeError {
                message: format!(
                    "Redundant pattern: this arm will never be matched because previous patterns cover all cases"
                ),
                span: arms[idx].span,
            });
        }

        result_type.unwrap_or(Type::Error)
    }

    fn check_pattern(&mut self, pattern: &Pattern, expected: &Type) -> Vec<(String, Type)> {
        match pattern {
            Pattern::Wildcard(_) => Vec::new(),

            Pattern::Var(ident) => {
                vec![(ident.name.clone(), expected.clone())]
            }

            Pattern::Literal(lit) => {
                let lit_type = self.infer_literal(lit);
                if let Err(e) = unify(&lit_type, expected) {
                    self.errors.push(TypeError {
                        message: format!("Pattern literal type mismatch: {}", e),
                        span: lit.span,
                    });
                }
                Vec::new()
            }

            Pattern::Constructor { name, fields, span } => {
                // Look up constructor
                // For now, handle Option specially
                match name.name.as_str() {
                    "None" => {
                        if let Err(e) = unify(expected, &Type::Option(Box::new(Type::var()))) {
                            self.errors.push(TypeError {
                                message: format!(
                                    "None pattern doesn't match type {}: {}",
                                    expected, e
                                ),
                                span: *span,
                            });
                        }
                        Vec::new()
                    }
                    "Some" => {
                        let inner_type = Type::var();
                        if let Err(e) = unify(expected, &Type::Option(Box::new(inner_type.clone())))
                        {
                            self.errors.push(TypeError {
                                message: format!(
                                    "Some pattern doesn't match type {}: {}",
                                    expected, e
                                ),
                                span: *span,
                            });
                        }
                        if fields.len() == 1 {
                            self.check_pattern(&fields[0], &inner_type)
                        } else {
                            Vec::new()
                        }
                    }
                    _ => {
                        // Generic constructor handling
                        let mut bindings = Vec::new();
                        for field in fields {
                            bindings.extend(self.check_pattern(field, &Type::var()));
                        }
                        bindings
                    }
                }
            }

            Pattern::Tuple { elements, span } => {
                let element_types: Vec<Type> = elements.iter().map(|_| Type::var()).collect();
                if let Err(e) = unify(expected, &Type::Tuple(element_types.clone())) {
                    self.errors.push(TypeError {
                        message: format!("Tuple pattern doesn't match type {}: {}", expected, e),
                        span: *span,
                    });
                }

                let mut bindings = Vec::new();
                for (elem, typ) in elements.iter().zip(element_types.iter()) {
                    bindings.extend(self.check_pattern(elem, typ));
                }
                bindings
            }

            Pattern::Record { fields, span } => {
                let field_types: Vec<(String, Type)> = fields
                    .iter()
                    .map(|(n, _)| (n.name.clone(), Type::var()))
                    .collect();

                if let Err(e) = unify(expected, &Type::Record(field_types.clone())) {
                    self.errors.push(TypeError {
                        message: format!("Record pattern doesn't match type {}: {}", expected, e),
                        span: *span,
                    });
                }

                let mut bindings = Vec::new();
                for ((_, pattern), (_, typ)) in fields.iter().zip(field_types.iter()) {
                    bindings.extend(self.check_pattern(pattern, typ));
                }
                bindings
            }
        }
    }

    fn infer_block(&mut self, statements: &[Statement], result: &Expr, _span: Span) -> Type {
        // Process statements
        for stmt in statements {
            match stmt {
                Statement::Expr(e) => {
                    self.infer_expr(e);
                }
                Statement::Let {
                    name, typ, value, ..
                } => {
                    let value_type = self.infer_expr(value);

                    if let Some(type_expr) = typ {
                        let declared = self.resolve_type(type_expr);
                        if let Err(e) = unify(&value_type, &declared) {
                            self.errors.push(TypeError {
                                message: format!(
                                    "Variable '{}' has type {}, but declared type is {}: {}",
                                    name.name, value_type, declared, e
                                ),
                                span: name.span,
                            });
                        }
                    }

                    let scheme = self.env.generalize(&value_type);
                    self.env.bind(&name.name, scheme);
                }
            }
        }

        // Return the type of the result expression
        self.infer_expr(result)
    }

    fn infer_record(&mut self, fields: &[(Ident, Expr)], _span: Span) -> Type {
        let field_types: Vec<(String, Type)> = fields
            .iter()
            .map(|(name, expr)| (name.name.clone(), self.infer_expr(expr)))
            .collect();

        Type::Record(field_types)
    }

    fn infer_tuple(&mut self, elements: &[Expr], _span: Span) -> Type {
        let element_types: Vec<Type> = elements.iter().map(|e| self.infer_expr(e)).collect();
        Type::Tuple(element_types)
    }

    fn infer_list(&mut self, elements: &[Expr], span: Span) -> Type {
        if elements.is_empty() {
            return Type::List(Box::new(Type::var()));
        }

        let first_type = self.infer_expr(&elements[0]);
        for elem in &elements[1..] {
            let elem_type = self.infer_expr(elem);
            if let Err(e) = unify(&first_type, &elem_type) {
                self.errors.push(TypeError {
                    message: format!("List elements must have same type: {}", e),
                    span,
                });
            }
        }

        Type::List(Box::new(first_type))
    }

    fn infer_run(&mut self, expr: &Expr, handlers: &[(Ident, Expr)], _span: Span) -> Type {
        // The handlers provide implementations for effects
        // After running, those effects are no longer present

        // Add the handled effects to available effects
        let handled_effects: EffectSet =
            EffectSet::from_iter(handlers.iter().map(|(e, _)| e.name.clone()));

        // Built-in effects are always available in run blocks (they have runtime implementations)
        let builtin_effects: EffectSet =
            EffectSet::from_iter(["Console", "Fail", "State"].iter().map(|s| s.to_string()));

        // Extend current effects with handled ones and built-in effects
        let combined = self.current_effects.union(&handled_effects).union(&builtin_effects);
        let old_effects = std::mem::replace(&mut self.current_effects, combined);

        // Type check the expression
        let result_type = self.infer_expr(expr);

        // Type check handlers
        for (effect_name, handler_expr) in handlers {
            // Just check the handler expression for now
            let _ = self.infer_expr(handler_expr);

            // Verify the effect exists
            if self.env.lookup_effect(&effect_name.name).is_none() {
                // Find similar effect names for suggestion
                let available_effects: Vec<&str> = self.env.effects.keys()
                    .map(|s| s.as_str())
                    .collect();
                let suggestions = find_similar_names(&effect_name.name, available_effects, 2);

                let mut message = format!("Unknown effect: {}", effect_name.name);
                if let Some(hint) = format_did_you_mean(&suggestions) {
                    message.push_str(&format!(". {}", hint));
                }

                self.errors.push(TypeError {
                    message,
                    span: effect_name.span,
                });
            }
        }

        self.current_effects = old_effects;
        result_type
    }

    // Helper methods

    fn function_type(&mut self, func: &FunctionDecl) -> TypeScheme {
        // Save old type params and start fresh for this function's scope
        let old_params = std::mem::take(&mut self.type_params);

        // Bind type parameters to fresh type variables
        let mut bound_vars = Vec::new();
        for param in &func.type_params {
            let var = Type::var();
            if let Type::Var(n) = &var {
                bound_vars.push(*n);
            }
            self.type_params.insert(param.name.clone(), var);
        }

        // Resolve parameter and return types (will use type_params for generics)
        let param_types: Vec<Type> = func
            .params
            .iter()
            .map(|p| self.resolve_type(&p.typ))
            .collect();

        let return_type = self.resolve_type(&func.return_type);
        let effects = EffectSet::from_iter(func.effects.iter().map(|e| e.name.clone()));
        let properties = PropertySet::from_ast(&func.properties);

        // Restore old type params
        self.type_params = old_params;

        // Return polymorphic type scheme with bound variables
        TypeScheme {
            vars: bound_vars,
            typ: Type::function_with_properties(param_types, return_type, effects, properties),
        }
    }

    fn effect_def(&self, effect: &EffectDecl) -> EffectDef {
        EffectDef {
            name: effect.name.name.clone(),
            type_params: effect.type_params.iter().map(|p| p.name.clone()).collect(),
            operations: effect
                .operations
                .iter()
                .map(|op| EffectOpDef {
                    name: op.name.name.clone(),
                    params: op
                        .params
                        .iter()
                        .map(|p| (p.name.name.clone(), self.resolve_type(&p.typ)))
                        .collect(),
                    return_type: self.resolve_type(&op.return_type),
                })
                .collect(),
        }
    }

    fn type_def(&self, type_decl: &TypeDecl) -> types::TypeDef {
        match &type_decl.definition {
            ast::TypeDef::Alias(t) => types::TypeDef::Alias(self.resolve_type(t)),
            ast::TypeDef::Record(fields) => types::TypeDef::Record(
                fields
                    .iter()
                    .map(|f| (f.name.name.clone(), self.resolve_type(&f.typ)))
                    .collect(),
            ),
            ast::TypeDef::Enum(variants) => types::TypeDef::Enum(
                variants
                    .iter()
                    .map(|v| VariantDef {
                        name: v.name.name.clone(),
                        fields: match &v.fields {
                            VariantFields::Unit => VariantFieldsDef::Unit,
                            VariantFields::Tuple(types) => VariantFieldsDef::Tuple(
                                types.iter().map(|t| self.resolve_type(t)).collect(),
                            ),
                            VariantFields::Record(fields) => VariantFieldsDef::Record(
                                fields
                                    .iter()
                                    .map(|f| (f.name.name.clone(), self.resolve_type(&f.typ)))
                                    .collect(),
                            ),
                        },
                    })
                    .collect(),
            ),
        }
    }

    fn handler_def(&self, handler: &HandlerDecl) -> HandlerDef {
        HandlerDef {
            name: handler.name.name.clone(),
            effect: handler.effect.name.clone(),
            params: handler
                .params
                .iter()
                .map(|p| (p.name.name.clone(), self.resolve_type(&p.typ)))
                .collect(),
        }
    }

    fn trait_def(&self, trait_decl: &TraitDecl) -> TraitDef {
        let methods = trait_decl
            .methods
            .iter()
            .map(|m| TraitMethodDef {
                name: m.name.name.clone(),
                type_params: m.type_params.iter().map(|p| p.name.clone()).collect(),
                params: m
                    .params
                    .iter()
                    .map(|p| (p.name.name.clone(), self.resolve_type(&p.typ)))
                    .collect(),
                return_type: self.resolve_type(&m.return_type),
                has_default: m.default_impl.is_some(),
            })
            .collect();

        let super_traits = trait_decl
            .super_traits
            .iter()
            .map(|b| TraitBoundDef {
                trait_name: b.trait_name.name.clone(),
                type_args: b.type_args.iter().map(|t| self.resolve_type(t)).collect(),
            })
            .collect();

        TraitDef {
            name: trait_decl.name.name.clone(),
            type_params: trait_decl.type_params.iter().map(|p| p.name.clone()).collect(),
            super_traits,
            methods,
        }
    }

    fn collect_impl(&self, impl_decl: &ImplDecl) -> TraitImpl {
        use std::collections::HashMap;

        let methods: HashMap<String, Type> = impl_decl
            .methods
            .iter()
            .map(|m| {
                let return_type = m
                    .return_type
                    .as_ref()
                    .map(|t| self.resolve_type(t))
                    .unwrap_or_else(Type::var);
                let param_types: Vec<Type> = m
                    .params
                    .iter()
                    .map(|p| self.resolve_type(&p.typ))
                    .collect();
                let func_type = Type::function(param_types, return_type);
                (m.name.name.clone(), func_type)
            })
            .collect();

        let constraints = impl_decl
            .constraints
            .iter()
            .map(|c| {
                let bounds = c
                    .bounds
                    .iter()
                    .map(|b| TraitBoundDef {
                        trait_name: b.trait_name.name.clone(),
                        type_args: b.type_args.iter().map(|t| self.resolve_type(t)).collect(),
                    })
                    .collect();
                (c.type_param.name.clone(), bounds)
            })
            .collect();

        TraitImpl {
            trait_name: impl_decl.trait_name.name.clone(),
            trait_args: impl_decl
                .trait_args
                .iter()
                .map(|t| self.resolve_type(t))
                .collect(),
            target_type: self.resolve_type(&impl_decl.target_type),
            type_params: impl_decl.type_params.iter().map(|p| p.name.clone()).collect(),
            constraints,
            methods,
        }
    }

    fn check_impl(&mut self, impl_decl: &ImplDecl) {
        // Verify the trait exists
        let trait_name = &impl_decl.trait_name.name;
        let trait_def = match self.env.traits.get(trait_name) {
            Some(def) => def.clone(),
            None => {
                self.errors.push(TypeError {
                    message: format!("Unknown trait: {}", trait_name),
                    span: impl_decl.span,
                });
                return;
            }
        };

        // Verify all required methods are implemented
        for method_def in &trait_def.methods {
            if !method_def.has_default {
                let implemented = impl_decl.methods.iter().any(|m| m.name.name == method_def.name);
                if !implemented {
                    self.errors.push(TypeError {
                        message: format!(
                            "Missing implementation for required method '{}' of trait '{}'",
                            method_def.name, trait_name
                        ),
                        span: impl_decl.span,
                    });
                }
            }
        }

        // Type check each implemented method
        for impl_method in &impl_decl.methods {
            // Find the method signature in the trait
            let method_def = trait_def.methods.iter().find(|m| m.name == impl_method.name.name);
            if method_def.is_none() {
                self.errors.push(TypeError {
                    message: format!(
                        "Method '{}' is not defined in trait '{}'",
                        impl_method.name.name, trait_name
                    ),
                    span: impl_method.span,
                });
                continue;
            }

            // Set up local environment with parameters
            let mut local_env = self.env.clone();
            for param in &impl_method.params {
                let param_type = self.resolve_type(&param.typ);
                local_env.bind(&param.name.name, TypeScheme::mono(param_type));
            }

            // Type check the body
            let old_env = std::mem::replace(&mut self.env, local_env);
            let body_type = self.infer_expr(&impl_method.body);
            self.env = old_env;

            // Check return type matches if specified
            if let Some(ref return_type_expr) = impl_method.return_type {
                let return_type = self.resolve_type(return_type_expr);
                if let Err(e) = unify(&body_type, &return_type) {
                    self.errors.push(TypeError {
                        message: format!(
                            "Method '{}' body has type {}, but declared return type is {}: {}",
                            impl_method.name.name, body_type, return_type, e
                        ),
                        span: impl_method.span,
                    });
                }
            }
        }
    }

    fn resolve_type(&self, type_expr: &TypeExpr) -> Type {
        match type_expr {
            TypeExpr::Named(ident) => match ident.name.as_str() {
                "Int" => Type::Int,
                "Float" => Type::Float,
                "Bool" => Type::Bool,
                "String" => Type::String,
                "Char" => Type::Char,
                "Unit" => Type::Unit,
                "_" => Type::var(),
                name => {
                    // Check if it's a type parameter in scope (for generics)
                    if let Some(var) = self.type_params.get(name) {
                        return var.clone();
                    }
                    Type::Named(name.to_string())
                }
            },
            TypeExpr::App(constructor, args) => {
                let resolved_args: Vec<Type> = args.iter().map(|a| self.resolve_type(a)).collect();

                // Handle built-in generic types
                if let TypeExpr::Named(name) = constructor.as_ref() {
                    match name.name.as_str() {
                        "List" if resolved_args.len() == 1 => {
                            return Type::List(Box::new(resolved_args[0].clone()));
                        }
                        "Option" if resolved_args.len() == 1 => {
                            return Type::Option(Box::new(resolved_args[0].clone()));
                        }
                        _ => {}
                    }
                }

                Type::App {
                    constructor: Box::new(self.resolve_type(constructor)),
                    args: resolved_args,
                }
            }
            TypeExpr::Function {
                params,
                return_type,
                effects,
            } => Type::function_with_effects(
                params.iter().map(|p| self.resolve_type(p)).collect(),
                self.resolve_type(return_type),
                EffectSet::from_iter(effects.iter().map(|e| e.name.clone())),
            ),
            TypeExpr::Tuple(elements) => {
                Type::Tuple(elements.iter().map(|e| self.resolve_type(e)).collect())
            }
            TypeExpr::Record(fields) => Type::Record(
                fields
                    .iter()
                    .map(|f| (f.name.name.clone(), self.resolve_type(&f.typ)))
                    .collect(),
            ),
            TypeExpr::Unit => Type::Unit,
            TypeExpr::Versioned {
                base,
                constraint: _,
            } => {
                // For now, resolve the base type and ignore versioning
                // Full version tracking will be added in the type system
                self.resolve_type(base)
            }
        }
    }
}

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