/// Type checker — Hindley-Milner with signal-awareness and effect tracking.
///
/// Walks the DreamStack AST and:
/// 1. Infers types for all expressions via HM unification
/// 2. Classifies signals using the signal graph (Source vs Derived)
/// 3. Tracks effect perform/handle pairs with scoped handlers
/// 4. Ensures views only contain valid UI expressions

use std::collections::HashMap;

use ds_parser::{Program, Declaration, LetDecl, ViewDecl, Expr, BinOp, UnaryOp, TypeExpr, TypeAliasDecl, Pattern};
use crate::types::{Type, TypeVar, EffectType, Predicate, PredicateExpr};
use crate::errors::{TypeError, TypeErrorKind};

/// Signal classification from the analyzer (mirrors ds_analyzer::SignalKind).
#[derive(Debug, Clone, PartialEq)]
pub enum SignalClass {
    Source,
    Derived,
}

/// Minimal signal graph info for type checking.
#[derive(Debug, Clone, Default)]
pub struct SignalInfo {
    pub signals: HashMap<String, SignalClass>,
}

/// The type checker.
pub struct TypeChecker {
    /// Type environment: variable name → type.
    env: HashMap<String, Type>,
    /// Accumulated errors (error-recovery: keep going after first error).
    errors: Vec<TypeError>,
    /// Effect stack: currently handled effects.
    effect_handlers: Vec<Vec<EffectType>>,
    /// Next type variable ID.
    next_tv: u32,
    /// Type substitutions from unification.
    substitutions: HashMap<TypeVar, Type>,
    /// Currently inside a view block?
    in_view: bool,
    /// Type alias registry: name → resolved Type.
    type_aliases: HashMap<String, Type>,
    /// Enum variant registry: enum name → list of variant names.
    enum_variants: HashMap<String, Vec<String>>,
}

impl TypeChecker {
    pub fn new() -> Self {
        TypeChecker {
            env: HashMap::new(),
            errors: Vec::new(),
            effect_handlers: Vec::new(),
            next_tv: 0,
            substitutions: HashMap::new(),
            in_view: false,
            type_aliases: HashMap::new(),
            enum_variants: HashMap::new(),
        }
    }

    /// Generate a fresh type variable.
    fn fresh_tv(&mut self) -> Type {
        let tv = TypeVar(self.next_tv);
        self.next_tv += 1;
        Type::Var(tv)
    }

    /// Record a type error.
    fn error(&mut self, kind: TypeErrorKind) {
        self.errors.push(TypeError::new(kind));
    }

    // ── Hindley-Milner Unification ──────────────────────────────────

    /// Apply substitutions to a type, chasing type variables to their bindings.
    pub fn apply_subst(&self, ty: &Type) -> Type {
        match ty {
            Type::Var(tv) => {
                if let Some(bound) = self.substitutions.get(tv) {
                    self.apply_subst(bound)
                } else {
                    ty.clone()
                }
            }
            Type::Signal(inner) => Type::Signal(Box::new(self.apply_subst(inner))),
            Type::Derived(inner) => Type::Derived(Box::new(self.apply_subst(inner))),
            Type::Array(inner) => Type::Array(Box::new(self.apply_subst(inner))),
            Type::Stream(inner) => Type::Stream(Box::new(self.apply_subst(inner))),
            Type::Spring(inner) => Type::Spring(Box::new(self.apply_subst(inner))),
            Type::Fn { params, ret, effects } => Type::Fn {
                params: params.iter().map(|p| self.apply_subst(p)).collect(),
                ret: Box::new(self.apply_subst(ret)),
                effects: effects.clone(),
            },
            Type::Record(fields) => Type::Record(
                fields.iter().map(|(k, v)| (k.clone(), self.apply_subst(v))).collect(),
            ),
            Type::Refined { base, predicate } => Type::Refined {
                base: Box::new(self.apply_subst(base)),
                predicate: predicate.clone(),
            },
            _ => ty.clone(),
        }
    }

    /// Occurs check: does `tv` appear free in `ty`?
    fn occurs_check(&self, tv: TypeVar, ty: &Type) -> bool {
        let resolved = self.apply_subst(ty);
        match &resolved {
            Type::Var(other) => *other == tv,
            Type::Signal(inner) | Type::Derived(inner) | Type::Array(inner)
            | Type::Stream(inner) | Type::Spring(inner) => self.occurs_check(tv, inner),
            Type::Fn { params, ret, .. } => {
                params.iter().any(|p| self.occurs_check(tv, p)) || self.occurs_check(tv, ret)
            }
            Type::Record(fields) => fields.values().any(|v| self.occurs_check(tv, v)),
            Type::Refined { base, .. } => self.occurs_check(tv, base),
            _ => false,
        }
    }

    /// Unify two types, recording substitutions. Returns Ok or an error kind.
    fn unify(&mut self, t1: &Type, t2: &Type) -> Result<(), TypeErrorKind> {
        let a = self.apply_subst(t1);
        let b = self.apply_subst(t2);

        if a == b {
            return Ok(());
        }

        match (&a, &b) {
            // Type variable binds to anything (with occurs check)
            (Type::Var(tv), _) => {
                if self.occurs_check(*tv, &b) {
                    return Err(TypeErrorKind::InfiniteType {
                        var: format!("?{}", tv.0),
                        ty: b.clone(),
                    });
                }
                self.substitutions.insert(*tv, b);
                Ok(())
            }
            (_, Type::Var(tv)) => {
                if self.occurs_check(*tv, &a) {
                    return Err(TypeErrorKind::InfiniteType {
                        var: format!("?{}", tv.0),
                        ty: a.clone(),
                    });
                }
                self.substitutions.insert(*tv, a);
                Ok(())
            }

            // Numeric coercion: Int unifies with Float
            (Type::Int, Type::Float) | (Type::Float, Type::Int) => Ok(()),

            // Structural: unify inner types
            (Type::Signal(a_inner), Type::Signal(b_inner)) => self.unify(a_inner, b_inner),
            (Type::Derived(a_inner), Type::Derived(b_inner)) => self.unify(a_inner, b_inner),
            (Type::Array(a_inner), Type::Array(b_inner)) => self.unify(a_inner, b_inner),
            (Type::Stream(a_inner), Type::Stream(b_inner)) => self.unify(a_inner, b_inner),
            (Type::Spring(a_inner), Type::Spring(b_inner)) => self.unify(a_inner, b_inner),

            // Signal<T> can unify with Derived<T> (both are reactive wrappers)
            (Type::Signal(a_inner), Type::Derived(b_inner))
            | (Type::Derived(a_inner), Type::Signal(b_inner)) => self.unify(a_inner, b_inner),

            // Functions: unify params and return, effects must match
            (Type::Fn { params: p1, ret: r1, .. }, Type::Fn { params: p2, ret: r2, .. }) => {
                if p1.len() != p2.len() {
                    return Err(TypeErrorKind::Mismatch {
                        expected: a.clone(),
                        found: b.clone(),
                        context: format!("Function arity: expected {} params, found {}", p1.len(), p2.len()),
                    });
                }
                for (pa, pb) in p1.iter().zip(p2.iter()) {
                    self.unify(pa, pb)?;
                }
                self.unify(r1, r2)
            }

            // Records: unify shared fields
            (Type::Record(f1), Type::Record(f2)) => {
                for (key, ty1) in f1 {
                    if let Some(ty2) = f2.get(key) {
                        self.unify(ty1, ty2)?;
                    }
                }
                Ok(())
            }

            // Refinement: unify base types
            (Type::Refined { base: b1, .. }, _) => self.unify(b1, &b),
            (_, Type::Refined { base: b2, .. }) => self.unify(&a, b2),

            // Error recovery
            (Type::Error, _) | (_, Type::Error) => Ok(()),

            // Mismatch
            _ => Err(TypeErrorKind::Mismatch {
                expected: a.clone(),
                found: b.clone(),
                context: "Type unification failed".to_string(),
            }),
        }
    }

    /// Push an effect handler scope.
    fn push_effect_scope(&mut self, effects: Vec<EffectType>) {
        self.effect_handlers.push(effects);
    }

    /// Pop an effect handler scope.
    fn pop_effect_scope(&mut self) {
        self.effect_handlers.pop();
    }

    /// Check an entire program. Optionally accepts signal graph info for accurate classification.
    pub fn check_program(&mut self, program: &Program) {
        self.check_program_with_signals(program, None);
    }

    /// Check a program with signal graph classification data.
    pub fn check_program_with_signals(&mut self, program: &Program, signal_info: Option<&SignalInfo>) {
        // Pass 0: register type aliases (with cycle detection)
        for decl in &program.declarations {
            if let Declaration::TypeAlias(alias) = decl {
                // Check for direct self-reference cycle
                if Self::type_expr_references(&alias.definition, &alias.name) {
                    self.error(TypeErrorKind::TypeAliasCycle {
                        name: alias.name.clone(),
                    });
                    continue;
                }
                let resolved = self.resolve_type_expr(&alias.definition);
                self.type_aliases.insert(alias.name.clone(), resolved);
            }
        }

        // Pass 0b: register enum declarations
        for decl in &program.declarations {
            if let Declaration::Enum(enum_decl) = decl {
                // Register each variant as a constructor function in the environment
                for variant in &enum_decl.variants {
                    if variant.fields.is_empty() {
                        // Unit variant: register as a constant of the enum type
                        self.env.insert(
                            format!("{}.{}", enum_decl.name, variant.name),
                            Type::String, // unit variants are string tags at runtime
                        );
                    } else {
                        // Data variant: register as a constructor function
                        let param_types: Vec<Type> = variant.fields.iter()
                            .map(|_| self.fresh_tv())
                            .collect();
                        let ret_ty = self.fresh_tv();
                        let key = format!("{}.{}", enum_decl.name, variant.name);
                        self.env.insert(
                            key,
                            Type::Fn {
                                params: param_types,
                                ret: Box::new(ret_ty),
                                effects: vec![EffectType::Pure],
                            },
                        );
                    }
                }
                // Register the enum itself as a record of its variants
                let enum_ty = self.fresh_tv();
                self.type_aliases.insert(enum_decl.name.clone(), enum_ty);
                // Register variant names for exhaustiveness checking
                let variant_names: Vec<String> = enum_decl.variants.iter()
                    .map(|v| v.name.clone())
                    .collect();
                self.enum_variants.insert(enum_decl.name.clone(), variant_names);
            }
        }

        // First pass: register all let declarations
        for decl in &program.declarations {
            if let Declaration::Let(let_decl) = decl {
                let inferred_ty = self.infer_expr(&let_decl.value);

                // If there's a type annotation, resolve and check it
                if let Some(ref type_ann) = let_decl.type_annotation {
                    let declared_ty = self.resolve_type_expr(type_ann);

                    // Check base type compatibility
                    let base_declared = match &declared_ty {
                        Type::Refined { base, .. } => base.as_ref(),
                        other => other,
                    };

                    let base_inferred = inferred_ty.unwrap_reactive();
                    if *base_declared != *base_inferred
                        && !matches!(base_declared, Type::Var(_))
                        && !matches!(base_inferred, Type::Var(_))
                        && *base_declared != Type::Error
                        && *base_inferred != Type::Error
                        // Allow Int/Float coercion
                        && !(matches!(base_declared, Type::Float) && matches!(base_inferred, Type::Int))
                    {
                        self.error(TypeErrorKind::Mismatch {
                            expected: base_declared.clone(),
                            found: base_inferred.clone(),
                            context: format!("in declaration of `{}`", let_decl.name),
                        });
                    }

                    // Check refinement predicate if present
                    if let Type::Refined { predicate, .. } = &declared_ty {
                        self.check_refinement(
                            predicate,
                            &let_decl.value,
                            &let_decl.name,
                            type_ann,
                        );
                    }
                }

                // Use signal graph classification if available, else heuristic
                let final_ty = if let Some(info) = signal_info {
                    match info.signals.get(&let_decl.name) {
                        Some(SignalClass::Source) => Type::Signal(Box::new(inferred_ty)),
                        Some(SignalClass::Derived) => Type::Derived(Box::new(inferred_ty)),
                        None => inferred_ty, // Not a signal — plain value
                    }
                } else {
                    // Heuristic fallback: literal init = Source, expression = Derived
                    let is_source = matches!(
                        let_decl.value,
                        Expr::IntLit(_) | Expr::FloatLit(_) | Expr::StringLit(_) | Expr::BoolLit(_)
                        | Expr::List(_) | Expr::Record(_)
                    );
                    if is_source {
                        Type::Signal(Box::new(inferred_ty))
                    } else {
                        Type::Derived(Box::new(inferred_ty))
                    }
                };
                self.env.insert(let_decl.name.clone(), final_ty);
            }
        }

        // Register effect declarations
        for decl in &program.declarations {
            if let Declaration::Effect(eff) = decl {
                let fn_type = Type::Fn {
                    params: eff.params.iter().map(|_| self.fresh_tv()).collect(),
                    ret: Box::new(self.fresh_tv()),
                    effects: vec![EffectType::Custom(eff.name.clone())],
                };
                self.env.insert(eff.name.clone(), fn_type);
            }
        }

        // Register handlers
        for decl in &program.declarations {
            if let Declaration::OnHandler(handler) = decl {
                let fn_type = Type::Fn {
                    params: vec![],
                    ret: Box::new(Type::Unit),
                    effects: vec![EffectType::Dom],
                };
                self.env.insert(handler.event.clone(), fn_type);
            }
        }

        // Second pass: check views
        for decl in &program.declarations {
            if let Declaration::View(view) = decl {
                self.check_view(view);
            }
        }
    }

    /// Resolve a TypeExpr from the AST into a semantic Type.
    fn resolve_type_expr(&self, type_expr: &TypeExpr) -> Type {
        match type_expr {
            TypeExpr::Named(name) => {
                // Check type aliases first
                if let Some(resolved) = self.type_aliases.get(name) {
                    return resolved.clone();
                }
                // Built-in type names
                match name.as_str() {
                    "Int" => Type::Int,
                    "Float" => Type::Float,
                    "String" => Type::String,
                    "Bool" => Type::Bool,
                    "View" => Type::View,
                    _ => Type::Named(name.clone()),
                }
            }
            TypeExpr::Generic(name, params) => {
                let resolved_params: Vec<Type> = params.iter()
                    .map(|p| self.resolve_type_expr(p))
                    .collect();
                match name.as_str() {
                    "Signal" if resolved_params.len() == 1 => {
                        Type::Signal(Box::new(resolved_params.into_iter().next().unwrap()))
                    }
                    "Array" if resolved_params.len() == 1 => {
                        Type::Array(Box::new(resolved_params.into_iter().next().unwrap()))
                    }
                    "Stream" if resolved_params.len() == 1 => {
                        Type::Stream(Box::new(resolved_params.into_iter().next().unwrap()))
                    }
                    _ => Type::Named(name.clone()),
                }
            }
            TypeExpr::Refined { base, predicate } => {
                let base_type = self.resolve_type_expr(base);
                let pred = Self::ast_to_predicate(predicate);
                Type::Refined {
                    base: Box::new(base_type),
                    predicate: pred,
                }
            }
        }
    }

    /// Convert an AST expression (from `where` clause) into a semantic Predicate.
    fn ast_to_predicate(expr: &Expr) -> Predicate {
        match expr {
            Expr::BinOp(left, op, right) => {
                match op {
                    BinOp::And => {
                        let l = Self::ast_to_predicate(left);
                        let r = Self::ast_to_predicate(right);
                        Predicate::And(Box::new(l), Box::new(r))
                    }
                    BinOp::Or => {
                        let l = Self::ast_to_predicate(left);
                        let r = Self::ast_to_predicate(right);
                        Predicate::Or(Box::new(l), Box::new(r))
                    }
                    BinOp::Gt => Predicate::Gt(
                        Box::new(Self::ast_to_pred_expr(left)),
                        Box::new(Self::ast_to_pred_expr(right)),
                    ),
                    BinOp::Gte => Predicate::Gte(
                        Box::new(Self::ast_to_pred_expr(left)),
                        Box::new(Self::ast_to_pred_expr(right)),
                    ),
                    BinOp::Lt => Predicate::Lt(
                        Box::new(Self::ast_to_pred_expr(left)),
                        Box::new(Self::ast_to_pred_expr(right)),
                    ),
                    BinOp::Lte => Predicate::Lte(
                        Box::new(Self::ast_to_pred_expr(left)),
                        Box::new(Self::ast_to_pred_expr(right)),
                    ),
                    BinOp::Eq => Predicate::Eq(
                        Box::new(Self::ast_to_pred_expr(left)),
                        Box::new(Self::ast_to_pred_expr(right)),
                    ),
                    BinOp::Neq => Predicate::Neq(
                        Box::new(Self::ast_to_pred_expr(left)),
                        Box::new(Self::ast_to_pred_expr(right)),
                    ),
                    _ => Predicate::Expr(format!("{:?}", expr)),
                }
            }
            Expr::UnaryOp(UnaryOp::Not, inner) => {
                Predicate::Not(Box::new(Self::ast_to_predicate(inner)))
            }
            _ => Predicate::Expr(format!("{:?}", expr)),
        }
    }

    /// Convert an AST expression into a predicate sub-expression.
    fn ast_to_pred_expr(expr: &Expr) -> PredicateExpr {
        match expr {
            Expr::Ident(name) if name == "value" => PredicateExpr::Value,
            Expr::Ident(name) => {
                // Non-value ident in a predicate — treat as an unresolvable reference.
                // This allows predicates like `value > min_threshold` where
                // min_threshold won't be statically evaluated.
                PredicateExpr::Call(name.clone(), vec![]) // model as zero-arg "call"
            }
            Expr::IntLit(n) => PredicateExpr::IntLit(*n),
            Expr::FloatLit(f) => PredicateExpr::FloatLit(*f),
            Expr::StringLit(s) => {
                let text: String = s.segments.iter().map(|seg| match seg {
                    ds_parser::StringSegment::Literal(l) => l.clone(),
                    _ => String::new(),
                }).collect();
                PredicateExpr::StringLit(text)
            }
            Expr::BoolLit(b) => PredicateExpr::BoolLit(*b),
            Expr::Call(name, args) => {
                let pred_args: Vec<PredicateExpr> = args.iter()
                    .map(|a| Self::ast_to_pred_expr(a))
                    .collect();
                PredicateExpr::Call(name.clone(), pred_args)
            }
            _ => {
                // Unrecognized expression — model as opaque. This will prevent
                // static evaluation (returns None) rather than silently succeeding.
                PredicateExpr::Call("<expr>".to_string(), vec![])
            }
        }
    }

    /// Check a refinement predicate against a value expression.
    /// For literals: evaluate statically and report error if violated.
    /// For dynamic expressions: accept (runtime guard will be emitted by codegen).
    fn check_refinement(
        &mut self,
        predicate: &Predicate,
        value_expr: &Expr,
        _var_name: &str,
        type_ann: &TypeExpr,
    ) {
        // Try to get a static value from the expression
        let static_val = match value_expr {
            Expr::IntLit(n) => Some(PredicateExpr::IntLit(*n)),
            Expr::FloatLit(f) => Some(PredicateExpr::FloatLit(*f)),
            Expr::BoolLit(b) => Some(PredicateExpr::BoolLit(*b)),
            Expr::StringLit(s) => {
                let text: String = s.segments.iter().map(|seg| match seg {
                    ds_parser::StringSegment::Literal(l) => l.clone(),
                    _ => String::new(),
                }).collect();
                Some(PredicateExpr::StringLit(text))
            }
            _ => None,
        };

        if let Some(ref val) = static_val {
            if let Some(result) = predicate.evaluate_static(val) {
                if !result {
                    // Static violation — compile-time error
                    let type_name = match type_ann {
                        TypeExpr::Named(n) => n.clone(),
                        TypeExpr::Refined { base, .. } => match base.as_ref() {
                            TypeExpr::Named(n) => format!("{} where ...", n),
                            _ => "<refined>".to_string(),
                        },
                        _ => "<unknown>".to_string(),
                    };
                    let value_display = match value_expr {
                        Expr::IntLit(n) => n.to_string(),
                        Expr::FloatLit(f) => f.to_string(),
                        Expr::BoolLit(b) => b.to_string(),
                        Expr::StringLit(_) => match val {
                            PredicateExpr::StringLit(s) => format!("\"{}\"", s),
                            _ => format!("{:?}", value_expr),
                        },
                        _ => format!("{:?}", value_expr),
                    };
                    self.error(TypeErrorKind::RefinementViolation {
                        type_name,
                        predicate: predicate.display(),
                        value: value_display,
                    });
                }
            }
        }
        // Dynamic values: accepted (codegen emits runtime guard)
    }

    /// Check if a TypeExpr references a given name (for cycle detection).
    fn type_expr_references(type_expr: &TypeExpr, name: &str) -> bool {
        match type_expr {
            TypeExpr::Named(n) => n == name,
            TypeExpr::Generic(n, params) => {
                n == name || params.iter().any(|p| Self::type_expr_references(p, name))
            }
            TypeExpr::Refined { base, .. } => Self::type_expr_references(base, name),
        }
    }

    /// Check a view declaration.
    fn check_view(&mut self, view: &ViewDecl) {
        self.in_view = true;
        // Dom effects are automatically handled inside view blocks
        self.push_effect_scope(vec![EffectType::Dom]);
        self.check_view_expr(&view.body);
        self.pop_effect_scope();
        self.in_view = false;
    }

    /// Check that an expression is valid inside a view.
    fn check_view_expr(&mut self, expr: &Expr) {
        match expr {
            Expr::Element(el) => {
                // Check child arguments and props
                for arg in &el.args {
                    self.check_view_expr(arg);
                }
                for (_, val) in &el.props {
                    self.check_view_expr(val);
                }
            }
            Expr::Container(container) => {
                for child in &container.children {
                    self.check_view_expr(child);
                }
            }
            Expr::StringLit(_) => { /* always valid */ }
            Expr::Ident(name) => {
                if !self.env.contains_key(name) {
                    self.error(TypeErrorKind::UnboundVariable {
                        name: name.clone(),
                    });
                }
            }
            Expr::When(condition, body, else_body) => {
                let cond_type = self.infer_expr(condition);
                let inner = cond_type.unwrap_reactive().clone();
                if inner != Type::Bool && !matches!(inner, Type::Var(_)) && inner != Type::Error {
                    self.error(TypeErrorKind::Mismatch {
                        expected: Type::Bool,
                        found: cond_type,
                        context: "Condition in `when` must be a boolean".to_string(),
                    });
                }
                self.check_view_expr(body);
            }
            Expr::If(condition, then_branch, else_branch) => {
                let cond_type = self.infer_expr(condition);
                let inner = cond_type.unwrap_reactive().clone();
                if inner != Type::Bool && !matches!(inner, Type::Var(_)) && inner != Type::Error {
                    self.error(TypeErrorKind::Mismatch {
                        expected: Type::Bool,
                        found: cond_type,
                        context: "Condition in if must be a boolean".to_string(),
                    });
                }
                self.check_view_expr(then_branch);
                self.check_view_expr(else_branch);
            }
            Expr::Block(exprs) => {
                for e in exprs {
                    self.check_view_expr(e);
                }
            }
            _ => {
                let _ = self.infer_expr(expr);
            }
        }
    }

    /// Infer the type of an expression.
    fn infer_expr(&mut self, expr: &Expr) -> Type {
        match expr {
            Expr::IntLit(_) => Type::Int,
            Expr::FloatLit(_) => Type::Float,
            Expr::StringLit(_) => Type::String,
            Expr::BoolLit(_) => Type::Bool,

            Expr::Ident(name) => {
                if let Some(ty) = self.env.get(name) {
                    ty.clone()
                } else {
                    self.error(TypeErrorKind::UnboundVariable { name: name.clone() });
                    Type::Error
                }
            }

            Expr::DotAccess(obj, field) => {
                let obj_ty = self.infer_expr(obj);
                // Unwrap reactive wrappers (Signal<Record> -> Record)
                let inner_ty = match &obj_ty {
                    Type::Signal(inner) | Type::Derived(inner) => inner.as_ref(),
                    other => other,
                };
                let resolved = self.apply_subst(inner_ty);
                match &resolved {
                    Type::Record(fields) => {
                        if let Some(ty) = fields.get(field) {
                            ty.clone()
                        } else {
                            self.error(TypeErrorKind::MissingField {
                                field: field.clone(),
                                record_type: resolved.clone(),
                            });
                            Type::Error
                        }
                    }
                    // Stream fields (e.g., remote.count) — always dynamic
                    Type::Stream(_) => self.fresh_tv(),
                    _ => self.fresh_tv(),
                }
            }

            Expr::BinOp(left, op, right) => {
                let left_ty = self.infer_expr(left);
                let right_ty = self.infer_expr(right);

                let left_inner = self.apply_subst(left_ty.unwrap_reactive());
                let right_inner = self.apply_subst(right_ty.unwrap_reactive());

                match op {
                    BinOp::Add | BinOp::Sub | BinOp::Mul | BinOp::Div | BinOp::Mod => {
                        // Try to unify operands for numeric ops
                        if let Err(_) = self.unify(&left_inner, &right_inner) {
                            // Check for string concat
                            if left_inner == Type::String && matches!(op, BinOp::Add) {
                                return Type::String;
                            }
                            if left_inner != Type::Error && right_inner != Type::Error {
                                self.error(TypeErrorKind::Mismatch {
                                    expected: left_inner.clone(),
                                    found: right_inner.clone(),
                                    context: format!("Both sides of {:?} must be the same numeric type", op),
                                });
                            }
                            return Type::Error;
                        }
                        // Result type: Int*Int=Int, anything with Float=Float
                        let resolved = self.apply_subst(&left_inner);
                        if resolved == Type::String && matches!(op, BinOp::Add) {
                            Type::String
                        } else if matches!(resolved, Type::Float) || matches!(self.apply_subst(&right_inner), Type::Float) {
                            Type::Float
                        } else if matches!(resolved, Type::Int) {
                            Type::Int
                        } else {
                            resolved
                        }
                    }
                    BinOp::Eq | BinOp::Neq | BinOp::Lt | BinOp::Gt | BinOp::Lte | BinOp::Gte => {
                        // Unify operands (comparisons require same type)
                        let _ = self.unify(&left_inner, &right_inner);
                        Type::Bool
                    }
                    BinOp::And | BinOp::Or => {
                        let _ = self.unify(&left_inner, &Type::Bool);
                        let _ = self.unify(&right_inner, &Type::Bool);
                        Type::Bool
                    }
                }
            }

            Expr::UnaryOp(op, operand) => {
                let ty = self.infer_expr(operand);
                let inner = self.apply_subst(ty.unwrap_reactive());
                match op {
                    UnaryOp::Neg => {
                        // Unify with Int first, fall back to Float
                        if self.unify(&inner, &Type::Int).is_ok() {
                            Type::Int
                        } else if self.unify(&inner, &Type::Float).is_ok() {
                            Type::Float
                        } else if inner == Type::Error {
                            Type::Error
                        } else {
                            self.error(TypeErrorKind::Mismatch {
                                expected: Type::Int,
                                found: inner.clone(),
                                context: "Unary `-` requires a numeric type".to_string(),
                            });
                            Type::Error
                        }
                    }
                    UnaryOp::Not => {
                        let _ = self.unify(&inner, &Type::Bool);
                        Type::Bool
                    }
                }
            }

            Expr::Call(name, args) => {
                if let Some(fn_ty) = self.env.get(name).cloned() {
                    match &fn_ty {
                        Type::Fn { params, ret, effects } => {
                            if args.len() != params.len() {
                                self.error(TypeErrorKind::ArityMismatch {
                                    function: name.clone(),
                                    expected: params.len(),
                                    found: args.len(),
                                });
                            } else {
                                // Unify each argument with its parameter type
                                for (arg, param) in args.iter().zip(params.iter()) {
                                    let arg_ty = self.infer_expr(arg);
                                    let arg_inner = self.apply_subst(arg_ty.unwrap_reactive());
                                    if let Err(e) = self.unify(&arg_inner, param) {
                                        self.error(e);
                                    }
                                }
                            }
                            // Check effects are handled
                            for eff in effects {
                                if *eff != EffectType::Pure && !self.is_effect_handled(eff) {
                                    self.error(TypeErrorKind::UnhandledEffect {
                                        effect: format!("{:?}", eff),
                                        function: name.clone(),
                                    });
                                }
                            }
                            self.apply_subst(ret)
                        }
                        _ => {
                            for arg in args {
                                self.infer_expr(arg);
                            }
                            self.fresh_tv()
                        }
                    }
                } else {
                    // Infer args even for unknown functions (useful for builtins)
                    for arg in args {
                        self.infer_expr(arg);
                    }
                    self.fresh_tv()
                }
            }

            Expr::Lambda(params, body) => {
                let mut param_types = Vec::new();
                for p in params {
                    let tv = self.fresh_tv();
                    self.env.insert(p.clone(), tv.clone());
                    param_types.push(tv);
                }
                let ret = self.infer_expr(body);
                Type::Fn {
                    params: param_types,
                    ret: Box::new(ret),
                    effects: vec![EffectType::Pure],
                }
            }

            Expr::Element(_) | Expr::Container(_) => {
                if !self.in_view {
                    self.error(TypeErrorKind::ViewOutsideBlock {
                        expr: "UI element".to_string(),
                    });
                }
                Type::View
            }

            Expr::Perform(name, _args) => {
                // An effect performance — check it's handled
                if !self.is_effect_handled(&EffectType::Custom(name.clone())) {
                    self.error(TypeErrorKind::UnhandledEffect {
                        effect: name.clone(),
                        function: "<perform>".to_string(),
                    });
                }
                self.fresh_tv()
            }

            Expr::StreamFrom { source, .. } => {
                Type::Stream(Box::new(self.fresh_tv()))
            }

            Expr::Spring(_props) => {
                Type::Spring(Box::new(Type::Float))
            }

            Expr::Record(fields) => {
                let mut field_types = HashMap::new();
                for (name, expr) in fields {
                    field_types.insert(name.clone(), self.infer_expr(expr));
                }
                Type::Record(field_types)
            }

            Expr::List(items) => {
                if items.is_empty() {
                    Type::Array(Box::new(self.fresh_tv()))
                } else {
                    let first_ty = self.infer_expr(&items[0]);
                    // Unify all element types for consistency
                    for item in &items[1..] {
                        let item_ty = self.infer_expr(item);
                        let _ = self.unify(&first_ty, &item_ty);
                    }
                    Type::Array(Box::new(self.apply_subst(&first_ty)))
                }
            }

            Expr::If(cond, then_br, else_br) => {
                let cond_ty = self.infer_expr(cond);
                let cond_inner = self.apply_subst(cond_ty.unwrap_reactive());
                let _ = self.unify(&cond_inner, &Type::Bool);
                let then_ty = self.infer_expr(then_br);
                let else_ty = self.infer_expr(else_br);
                // Unify both branches — they should return the same type
                let _ = self.unify(&then_ty, &else_ty);
                self.apply_subst(&then_ty)
            }

            Expr::When(cond, body, else_body) => {
                let cond_ty = self.infer_expr(cond);
                let cond_inner = self.apply_subst(cond_ty.unwrap_reactive());
                let _ = self.unify(&cond_inner, &Type::Bool);
                let body_ty = self.infer_expr(body);
                if let Some(eb) = else_body {
                    let else_ty = self.infer_expr(eb);
                    let _ = self.unify(&body_ty, &else_ty);
                }
                self.apply_subst(&body_ty)
            }

            Expr::Block(exprs) => {
                let mut last_ty = Type::Unit;
                for e in exprs {
                    last_ty = self.infer_expr(e);
                }
                last_ty
            }

            Expr::Pipe(left, right) => {
                let _input_ty = self.infer_expr(left);
                self.infer_expr(right)
            }

            Expr::Match(expr, arms) => {
                let _ = self.infer_expr(expr);

                // Exhaustiveness check: if scrutinee is a known variable with enum type
                if let Expr::Ident(name) = expr.as_ref() {
                    // Check if name matches a known enum type variable
                    for (enum_name, variants) in &self.enum_variants {
                        if name.to_lowercase() == enum_name.to_lowercase()
                            || self.env.get(name).map_or(false, |_| true)
                        {
                            // Collect matched patterns
                            let mut matched: std::collections::HashSet<String> = std::collections::HashSet::new();
                            let mut has_wildcard = false;
                            for arm in arms {
                                match &arm.pattern {
                                    Pattern::Wildcard => { has_wildcard = true; }
                                    Pattern::Ident(p) => { matched.insert(p.clone()); }
                                    Pattern::Literal(_) => {}
                                    Pattern::Constructor(p, _) => { matched.insert(p.clone()); }
                                    Pattern::IntLiteral(_) => {}
                                    Pattern::BoolLiteral(_) => {}
                                    Pattern::Tuple(_) => {}
                                }
                            }

                            if !has_wildcard {
                                let missing: Vec<String> = variants.iter()
                                    .filter(|v| !matched.contains(*v))
                                    .cloned()
                                    .collect();
                                if !missing.is_empty() {
                                    self.error(TypeErrorKind::NonExhaustiveMatch {
                                        enum_name: enum_name.clone(),
                                        missing_variants: missing,
                                    });
                                }
                            }
                            break;
                        }
                    }
                }

                if arms.is_empty() {
                    Type::Unit
                } else {
                    let first_ty = self.infer_expr(&arms[0].body);
                    // Unify all arm types — match must return consistent type
                    for arm in &arms[1..] {
                        let arm_ty = self.infer_expr(&arm.body);
                        let _ = self.unify(&first_ty, &arm_ty);
                    }
                    self.apply_subst(&first_ty)
                }
            }

            Expr::Assign(target, _op, value) => {
                let val_ty = self.infer_expr(value);
                // Check that assigned type is compatible with target type
                if let Expr::Ident(name) = target.as_ref() {
                    if let Some(var_ty) = self.env.get(name).cloned() {
                        let var_inner = self.apply_subst(var_ty.unwrap_reactive());
                        let val_inner = self.apply_subst(val_ty.unwrap_reactive());
                        let _ = self.unify(&var_inner, &val_inner);
                    }
                }
                Type::Unit
            }

            Expr::ForIn { item, index, iter, body } => {
                let iter_ty = self.infer_expr(iter);
                let iter_inner = self.apply_subst(iter_ty.unwrap_reactive());
                // Extract element type from Array<T> and bind iteration variable
                let elem_ty = match &iter_inner {
                    Type::Array(elem) => *elem.clone(),
                    _ => self.fresh_tv(),
                };
                self.env.insert(item.clone(), elem_ty);
                if let Some(idx) = index {
                    self.env.insert(idx.clone(), Type::Int);
                }
                self.infer_expr(body)
            }

            Expr::Each(_, list, body) => {
                let _ = self.infer_expr(list);
                self.infer_expr(body)
            }

            Expr::ComponentUse { props, children, .. } => {
                for (_, val) in props {
                    self.infer_expr(val);
                }
                for child in children {
                    self.infer_expr(child);
                }
                Type::View
            }

            Expr::Index(base, index) => {
                let base_ty = self.infer_expr(base);
                let _ = self.infer_expr(index);
                match base_ty {
                    Type::Array(elem) => *elem,
                    _ => self.fresh_tv(),
                }
            }

            Expr::Slot => Type::View,

            Expr::MethodCall(obj, _, args) => {
                self.infer_expr(obj);
                for arg in args {
                    self.infer_expr(arg);
                }
                self.fresh_tv()
            }

            Expr::Await(inner) => {
                let inner_ty = self.infer_expr(inner);
                // Await unwraps a Promise/Future, for now just returns inner type
                inner_ty
            }

            Expr::Merge(streams) => {
                for s in streams {
                    self.infer_expr(s);
                }
                Type::Stream(Box::new(self.fresh_tv()))
            }

            Expr::RawString(_) => Type::String,
        }
    }

    /// Check if an effect is currently handled.
    fn is_effect_handled(&self, effect: &EffectType) -> bool {
        for handlers in &self.effect_handlers {
            if handlers.contains(effect) {
                return true;
            }
        }
        false
    }

    /// Get the accumulated errors.
    pub fn errors(&self) -> &[TypeError] {
        &self.errors
    }

    /// Get the resolved type environment.
    pub fn type_env(&self) -> &HashMap<String, Type> {
        &self.env
    }

    /// Check if there are any errors.
    pub fn has_errors(&self) -> bool {
        !self.errors.is_empty()
    }

    /// Convert accumulated type errors into `Diagnostic` structs.
    pub fn errors_as_diagnostics(&self) -> Vec<ds_diagnostic::Diagnostic> {
        self.errors.iter().map(|err| {
            let (line, col) = err.span.unwrap_or((0, 0));
            let span = ds_parser::Span {
                start: 0,
                end: 0,
                line,
                col,
            };

            let (code, message) = match &err.kind {
                crate::errors::TypeErrorKind::Mismatch { expected, found, context } => {
                    ("E0100", format!("expected {}, found {} — {}", expected.display(), found.display(), context))
                }
                crate::errors::TypeErrorKind::NotReactive { found, context } => {
                    ("E0101", format!("type {} is not reactive — {}", found.display(), context))
                }
                crate::errors::TypeErrorKind::UnhandledEffect { effect, function } => {
                    ("E0102", format!("unhandled effect `{}` in function `{}`", effect, function))
                }
                crate::errors::TypeErrorKind::ViewOutsideBlock { expr } => {
                    ("E0103", format!("view expression `{}` outside a `view` block", expr))
                }
                crate::errors::TypeErrorKind::UnboundVariable { name } => {
                    ("E0104", format!("unbound variable `{}`", name))
                }
                crate::errors::TypeErrorKind::InfiniteType { var, ty } => {
                    ("E0105", format!("infinite type: {} ~ {}", var, ty.display()))
                }
                crate::errors::TypeErrorKind::ArityMismatch { function, expected, found } => {
                    ("E0106", format!("`{}` expects {} argument(s), found {}", function, expected, found))
                }
                crate::errors::TypeErrorKind::MissingField { field, record_type } => {
                    ("E0107", format!("no field `{}` on {}", field, record_type.display()))
                }
                crate::errors::TypeErrorKind::RefinementViolation { type_name, predicate, value } => {
                    ("E0108", format!("value `{}` violates refinement `{}` ({})", value, type_name, predicate))
                }
                crate::errors::TypeErrorKind::TypeAliasCycle { name } => {
                    ("E0109", format!("type alias `{}` creates a cycle", name))
                }
                crate::errors::TypeErrorKind::NonExhaustiveMatch { enum_name, missing_variants } => {
                    ("E0110", format!("non-exhaustive match on `{}` — missing: {}", enum_name, missing_variants.join(", ")))
                }
            };

            ds_diagnostic::Diagnostic::error(message, span)
                .with_code(code)
        }).collect()
    }

    /// Format all errors for display.
    pub fn display_errors(&self) -> String {
        self.errors.iter()
            .map(|e| e.display())
            .collect::<Vec<_>>()
            .join("\n")
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use ds_parser::{Declaration, LetDecl, ViewDecl, Expr, Span, Container, ContainerKind, Element};

    fn span() -> Span {
        Span { start: 0, end: 0, line: 0, col: 0 }
    }

    fn make_program(decls: Vec<Declaration>) -> Program {
        Program { declarations: decls }
    }

    #[test]
    fn test_signal_types() {
        let mut checker = TypeChecker::new();
        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "count".to_string(),
                type_annotation: None,
                value: Expr::IntLit(0),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(!checker.has_errors(), "Errors: {}", checker.display_errors());

        let count_ty = checker.type_env().get("count").unwrap();
        assert!(count_ty.is_reactive());
        assert_eq!(count_ty.display(), "Signal<Int>");
    }

    #[test]
    fn test_derived_types() {
        let mut checker = TypeChecker::new();
        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "count".to_string(),
                type_annotation: None,
                value: Expr::IntLit(0),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
            Declaration::Let(LetDecl {
                name: "doubled".to_string(),
                type_annotation: None,
                value: Expr::BinOp(
                    Box::new(Expr::Ident("count".to_string())),
                    BinOp::Mul,
                    Box::new(Expr::IntLit(2)),
                ),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(!checker.has_errors(), "Errors: {}", checker.display_errors());

        let doubled_ty = checker.type_env().get("doubled").unwrap();
        assert_eq!(doubled_ty.display(), "Derived<Int>");
    }

    #[test]
    fn test_unbound_variable() {
        let mut checker = TypeChecker::new();
        let program = make_program(vec![
            Declaration::View(ViewDecl {
                name: "main".to_string(),
                params: vec![],
                body: Expr::Container(Container {
                    kind: ContainerKind::Column,
                    children: vec![
                        Expr::Ident("nonexistent".to_string()),
                    ],
                    props: vec![],
                }),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(checker.has_errors());

        let msg = checker.display_errors();
        assert!(msg.contains("UNBOUND VARIABLE"));
        assert!(msg.contains("nonexistent"));
    }

    #[test]
    fn test_type_display_in_env() {
        let mut checker = TypeChecker::new();
        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "name".to_string(),
                type_annotation: None,
                value: Expr::StringLit(ds_parser::StringLit {
                    segments: vec![ds_parser::StringSegment::Literal("hello".to_string())],
                }),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(!checker.has_errors());

        let name_ty = checker.type_env().get("name").unwrap();
        assert_eq!(name_ty.display(), "Signal<String>");
    }

    #[test]
    fn test_view_outside_block_error() {
        let mut checker = TypeChecker::new();
        // Manually add condition: an element expression outside a view
        let el_expr = Expr::Element(Element {
            tag: "button".to_string(),
            args: vec![],
            props: vec![],
            modifiers: vec![],
        });
        let ty = checker.infer_expr(&el_expr);
        assert_eq!(ty, Type::View);
        assert!(checker.has_errors());
        let msg = checker.display_errors();
        assert!(msg.contains("VIEW OUTSIDE BLOCK"));
    }

    #[test]
    fn test_type_annotation_basic() {
        let mut checker = TypeChecker::new();
        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "count".to_string(),
                type_annotation: Some(ds_parser::TypeExpr::Named("Int".to_string())),
                value: Expr::IntLit(42),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(!checker.has_errors(), "Errors: {}", checker.display_errors());
    }

    #[test]
    fn test_type_annotation_mismatch() {
        let mut checker = TypeChecker::new();
        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "count".to_string(),
                type_annotation: Some(ds_parser::TypeExpr::Named("Int".to_string())),
                value: Expr::StringLit(ds_parser::StringLit {
                    segments: vec![ds_parser::StringSegment::Literal("oops".to_string())],
                }),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(checker.has_errors());
        let msg = checker.display_errors();
        assert!(msg.contains("TYPE MISMATCH"));
    }

    #[test]
    fn test_refinement_passes_static() {
        let mut checker = TypeChecker::new();
        // let count: Int where value > 0 = 5
        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "count".to_string(),
                type_annotation: Some(ds_parser::TypeExpr::Refined {
                    base: Box::new(ds_parser::TypeExpr::Named("Int".to_string())),
                    predicate: Box::new(Expr::BinOp(
                        Box::new(Expr::Ident("value".to_string())),
                        BinOp::Gt,
                        Box::new(Expr::IntLit(0)),
                    )),
                }),
                value: Expr::IntLit(5),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(!checker.has_errors(), "Errors: {}", checker.display_errors());
    }

    #[test]
    fn test_refinement_violation_static() {
        let mut checker = TypeChecker::new();
        // let count: Int where value > 0 = -1
        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "count".to_string(),
                type_annotation: Some(ds_parser::TypeExpr::Refined {
                    base: Box::new(ds_parser::TypeExpr::Named("Int".to_string())),
                    predicate: Box::new(Expr::BinOp(
                        Box::new(Expr::Ident("value".to_string())),
                        BinOp::Gt,
                        Box::new(Expr::IntLit(0)),
                    )),
                }),
                value: Expr::IntLit(-1),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(checker.has_errors());
        let msg = checker.display_errors();
        assert!(msg.contains("REFINEMENT VIOLATED"), "Expected REFINEMENT VIOLATED, got: {}", msg);
    }

    #[test]
    fn test_type_alias_with_refinement() {
        let mut checker = TypeChecker::new();
        // type PositiveInt = Int where value > 0
        // let count: PositiveInt = 5
        let program = make_program(vec![
            Declaration::TypeAlias(ds_parser::TypeAliasDecl {
                name: "PositiveInt".to_string(),
                definition: ds_parser::TypeExpr::Refined {
                    base: Box::new(ds_parser::TypeExpr::Named("Int".to_string())),
                    predicate: Box::new(Expr::BinOp(
                        Box::new(Expr::Ident("value".to_string())),
                        BinOp::Gt,
                        Box::new(Expr::IntLit(0)),
                    )),
                },
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
            Declaration::Let(LetDecl {
                name: "count".to_string(),
                type_annotation: Some(ds_parser::TypeExpr::Named("PositiveInt".to_string())),
                value: Expr::IntLit(5),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(!checker.has_errors(), "Errors: {}", checker.display_errors());
    }

    #[test]
    fn test_type_alias_cycle_detection() {
        let mut checker = TypeChecker::new();
        // type Foo = Foo  (self-referential cycle)
        let program = make_program(vec![
            Declaration::TypeAlias(ds_parser::TypeAliasDecl {
                name: "Foo".to_string(),
                definition: ds_parser::TypeExpr::Named("Foo".to_string()),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(checker.has_errors());
        let msg = checker.display_errors();
        assert!(msg.contains("TYPE ALIAS CYCLE"), "Expected cycle error, got: {}", msg);
        assert!(msg.contains("Foo"));
    }

    #[test]
    fn test_refinement_violation_shows_readable_value() {
        let mut checker = TypeChecker::new();
        // let x: Int where value > 0 = -42
        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "x".to_string(),
                type_annotation: Some(ds_parser::TypeExpr::Refined {
                    base: Box::new(ds_parser::TypeExpr::Named("Int".to_string())),
                    predicate: Box::new(Expr::BinOp(
                        Box::new(Expr::Ident("value".to_string())),
                        BinOp::Gt,
                        Box::new(Expr::IntLit(0)),
                    )),
                }),
                value: Expr::IntLit(-42),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(checker.has_errors());
        let msg = checker.display_errors();
        assert!(msg.contains("-42"), "Error should show the literal value, got: {}", msg);
        assert!(msg.contains("REFINEMENT VIOLATED"));
    }

    // ── Hindley-Milner Unification tests ──────────────────────────

    #[test]
    fn test_unify_same_types() {
        let mut checker = TypeChecker::new();
        assert!(checker.unify(&Type::Int, &Type::Int).is_ok());
        assert!(checker.unify(&Type::String, &Type::String).is_ok());
        assert!(checker.unify(&Type::Bool, &Type::Bool).is_ok());
    }

    #[test]
    fn test_unify_type_variable() {
        let mut checker = TypeChecker::new();
        let tv = checker.fresh_tv();
        assert!(checker.unify(&tv, &Type::Int).is_ok());
        assert_eq!(checker.apply_subst(&tv), Type::Int);
    }

    #[test]
    fn test_unify_mismatch() {
        let mut checker = TypeChecker::new();
        assert!(checker.unify(&Type::Int, &Type::String).is_err());
    }

    #[test]
    fn test_unify_numeric_coercion() {
        let mut checker = TypeChecker::new();
        assert!(checker.unify(&Type::Int, &Type::Float).is_ok());
    }

    #[test]
    fn test_unify_signal_types() {
        let mut checker = TypeChecker::new();
        assert!(checker.unify(
            &Type::Signal(Box::new(Type::Int)),
            &Type::Signal(Box::new(Type::Int)),
        ).is_ok());
        // Signal<Int> vs Signal<String> should fail
        assert!(checker.unify(
            &Type::Signal(Box::new(Type::Int)),
            &Type::Signal(Box::new(Type::String)),
        ).is_err());
    }

    #[test]
    fn test_occurs_check_infinite_type() {
        let mut checker = TypeChecker::new();
        let tv = checker.fresh_tv();
        let array_of_tv = Type::Array(Box::new(tv.clone()));
        let result = checker.unify(&tv, &array_of_tv);
        assert!(result.is_err());
        match result.unwrap_err() {
            TypeErrorKind::InfiniteType { .. } => { /* expected */ }
            other => panic!("Expected InfiniteType, got {:?}", other),
        }
    }

    // ── Signal Graph Classification tests ─────────────────────────

    #[test]
    fn test_signal_graph_classification() {
        let mut checker = TypeChecker::new();
        let mut info = SignalInfo::default();
        info.signals.insert("count".to_string(), SignalClass::Source);
        info.signals.insert("doubled".to_string(), SignalClass::Derived);

        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "count".to_string(),
                type_annotation: None,
                value: Expr::IntLit(0),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
            Declaration::Let(LetDecl {
                name: "doubled".to_string(),
                type_annotation: None,
                value: Expr::BinOp(
                    Box::new(Expr::Ident("count".to_string())),
                    BinOp::Mul,
                    Box::new(Expr::IntLit(2)),
                ),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program_with_signals(&program, Some(&info));
        assert!(!checker.has_errors(), "Errors: {}", checker.display_errors());

        let count_ty = checker.type_env().get("count").unwrap();
        assert_eq!(count_ty.display(), "Signal<Int>");

        let doubled_ty = checker.type_env().get("doubled").unwrap();
        assert_eq!(doubled_ty.display(), "Derived<Int>");
    }

    // ── Effect Scope tests ────────────────────────────────────────

    #[test]
    fn test_effect_handler_in_view() {
        let mut checker = TypeChecker::new();
        // Dom effect should be auto-handled inside view
        checker.push_effect_scope(vec![EffectType::Dom]);
        assert!(checker.is_effect_handled(&EffectType::Dom));
        checker.pop_effect_scope();
        assert!(!checker.is_effect_handled(&EffectType::Dom));
    }

    #[test]
    fn test_list_type_unification() {
        let mut checker = TypeChecker::new();
        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "items".to_string(),
                type_annotation: None,
                value: Expr::List(vec![
                    Expr::IntLit(1),
                    Expr::IntLit(2),
                    Expr::IntLit(3),
                ]),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(!checker.has_errors(), "Errors: {}", checker.display_errors());

        let items_ty = checker.type_env().get("items").unwrap();
        assert_eq!(items_ty.display(), "Signal<[Int]>");
    }

    // ── Refinement Pass 2: Branch & Structural tests ──────────────

    #[test]
    fn test_dot_access_through_signal() {
        let mut checker = TypeChecker::new();
        // Simulate: let rec = { x: 10, y: 20 }
        // then access rec.x should work through Signal<Record> wrapper
        let mut fields = HashMap::new();
        fields.insert("x".to_string(), Type::Int);
        fields.insert("y".to_string(), Type::Int);
        let record_type = Type::Signal(Box::new(Type::Record(fields)));
        checker.env.insert("rec".to_string(), record_type);

        let expr = Expr::DotAccess(
            Box::new(Expr::Ident("rec".to_string())),
            "x".to_string(),
        );
        let ty = checker.infer_expr(&expr);
        assert_eq!(ty, Type::Int);
        assert!(!checker.has_errors(), "Errors: {}", checker.display_errors());
    }

    #[test]
    fn test_if_else_branch_unification() {
        let mut checker = TypeChecker::new();
        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "flag".to_string(),
                type_annotation: None,
                value: Expr::BoolLit(true),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
            Declaration::Let(LetDecl {
                name: "result".to_string(),
                type_annotation: None,
                value: Expr::If(
                    Box::new(Expr::Ident("flag".to_string())),
                    Box::new(Expr::IntLit(1)),
                    Box::new(Expr::IntLit(2)),
                ),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(!checker.has_errors(), "Errors: {}", checker.display_errors());
    }

    #[test]
    fn test_for_in_binds_item_type() {
        let mut checker = TypeChecker::new();
        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "nums".to_string(),
                type_annotation: None,
                value: Expr::List(vec![
                    Expr::IntLit(1),
                    Expr::IntLit(2),
                    Expr::IntLit(3),
                ]),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
            Declaration::Let(LetDecl {
                name: "total".to_string(),
                type_annotation: None,
                value: Expr::ForIn {
                    item: "n".to_string(),
                    index: None,
                    iter: Box::new(Expr::Ident("nums".to_string())),
                    body: Box::new(Expr::Ident("n".to_string())),
                },
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(!checker.has_errors(), "Errors: {}", checker.display_errors());

        // The iteration variable 'n' should have been bound as Int
        let n_ty = checker.type_env().get("n").unwrap();
        assert_eq!(*n_ty, Type::Int);
    }

    #[test]
    fn test_assign_type_compatibility() {
        let mut checker = TypeChecker::new();
        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "count".to_string(),
                type_annotation: None,
                value: Expr::IntLit(0),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        assert!(!checker.has_errors());

        // Assigning an Int to a Signal<Int> should be fine
        let assign_expr = Expr::Assign(
            Box::new(Expr::Ident("count".to_string())),
            ds_parser::AssignOp::Set,
            Box::new(Expr::IntLit(5)),
        );
        let ty = checker.infer_expr(&assign_expr);
        assert_eq!(ty, Type::Unit);
        assert!(!checker.has_errors(), "Errors: {}", checker.display_errors());
    }

    #[test]
    fn test_unify_fn_arg_types() {
        let mut checker = TypeChecker::new();
        // Register a function: add(Int, Int) -> Int
        let fn_type = Type::Fn {
            params: vec![Type::Int, Type::Int],
            ret: Box::new(Type::Int),
            effects: vec![EffectType::Pure],
        };
        checker.env.insert("add".to_string(), fn_type);

        // Call with correct types should work
        let call_expr = Expr::Call("add".to_string(), vec![
            Expr::IntLit(1),
            Expr::IntLit(2),
        ]);
        let ty = checker.infer_expr(&call_expr);
        assert_eq!(ty, Type::Int);
        assert!(!checker.has_errors(), "Errors: {}", checker.display_errors());
    }

    // ── v0.6 Integration Tests ──────────────────────────────

    /// Parse DreamStack source and type-check it.
    fn check_source(src: &str) -> TypeChecker {
        let mut lexer = ds_parser::Lexer::new(src);
        let tokens = lexer.tokenize();
        let mut parser = ds_parser::Parser::new(tokens);
        let program = parser.parse_program().expect("parse failed");
        let mut checker = TypeChecker::new();
        checker.check_program(&program);
        checker
    }

    #[test]
    fn test_counter_program_no_errors() {
        let checker = check_source(
            "let count = 0\nlet doubled = count * 2\nview main = column [ text \"hi\" ]"
        );
        assert!(!checker.has_errors(), "counter should have no type errors: {}", checker.display_errors());
    }

    #[test]
    fn test_errors_as_diagnostics_conversion() {
        let mut checker = TypeChecker::new();
        let program = make_program(vec![
            Declaration::Let(LetDecl {
                name: "x".to_string(),
                type_annotation: None,
                value: Expr::BinOp(
                    Box::new(Expr::IntLit(1)),
                    BinOp::Add,
                    Box::new(Expr::StringLit(ds_parser::StringLit {
                        segments: vec![ds_parser::StringSegment::Literal("hello".to_string())],
                    })),
                ),
                visibility: ds_parser::Visibility::Private,
                doc: None,
                span: span(),
            }),
        ]);
        checker.check_program(&program);
        let diags = checker.errors_as_diagnostics();
        // Whether or not there's a type error, the diagnostics should be well-formed
        for diag in &diags {
            assert!(!diag.message.is_empty());
            assert!(diag.code.is_some());
        }
    }

    #[test]
    fn test_derived_int_arithmetic() {
        let checker = check_source(
            "let a = 10\nlet b = 20\nlet sum = a + b"
        );
        assert!(!checker.has_errors(), "int arithmetic should not error: {}", checker.display_errors());
        assert_eq!(*checker.type_env().get("a").unwrap(), Type::Signal(Box::new(Type::Int)));
        assert_eq!(*checker.type_env().get("b").unwrap(), Type::Signal(Box::new(Type::Int)));
    }

    #[test]
    fn test_string_signal_type() {
        let checker = check_source("let name = \"world\"");
        assert!(!checker.has_errors(), "string let should not error: {}", checker.display_errors());
        assert_eq!(*checker.type_env().get("name").unwrap(), Type::Signal(Box::new(Type::String)));
    }

    #[test]
    fn test_bool_signal_type() {
        let checker = check_source("let active = true");
        assert!(!checker.has_errors(), "bool let should not error: {}", checker.display_errors());
        assert_eq!(*checker.type_env().get("active").unwrap(), Type::Signal(Box::new(Type::Bool)));
    }

    #[test]
    fn test_multiple_signals_type_env() {
        let checker = check_source(
            "let count = 0\nlet name = \"test\"\nlet active = false\nlet ratio = 3.14"
        );
        assert!(!checker.has_errors(), "multi-signal: {}", checker.display_errors());
        assert_eq!(*checker.type_env().get("count").unwrap(), Type::Signal(Box::new(Type::Int)));
        assert_eq!(*checker.type_env().get("name").unwrap(), Type::Signal(Box::new(Type::String)));
        assert_eq!(*checker.type_env().get("active").unwrap(), Type::Signal(Box::new(Type::Bool)));
        assert_eq!(*checker.type_env().get("ratio").unwrap(), Type::Signal(Box::new(Type::Float)));
    }

    // ── v0.9 Rust-Like Match Type Checker Tests ─────────────

    #[test]
    fn test_exhaustive_enum_match_no_error() {
        let checker = check_source(
            "enum Color { Red, Green, Blue }\nlet c = 0\nlet label = match c\n  Red -> \"r\"\n  Green -> \"g\"\n  Blue -> \"b\""
        );
        // Exhaustive match on Color — should not produce NonExhaustiveMatch
        let has_exhaustive_err = checker.display_errors().contains("non-exhaustive");
        assert!(!has_exhaustive_err, "exhaustive match should not error: {}", checker.display_errors());
    }

    #[test]
    fn test_non_exhaustive_enum_match_error() {
        let checker = check_source(
            "enum Color { Red, Green, Blue }\nlet c = 0\nlet label = match c\n  Red -> \"r\"\n  Green -> \"g\""
        );
        // Missing Blue — should produce non-exhaustive error
        let errors = checker.display_errors();
        assert!(errors.contains("non-exhaustive") || errors.contains("Blue") || checker.has_errors(),
            "non-exhaustive match should warn about Blue: {}", errors);
    }

    #[test]
    fn test_wildcard_always_exhaustive() {
        let checker = check_source(
            "enum Dir { Up, Down, Left, Right }\nlet d = 0\nlet label = match d\n  Up -> \"up\"\n  _ -> \"other\""
        );
        let has_exhaustive_err = checker.display_errors().contains("non-exhaustive");
        assert!(!has_exhaustive_err, "wildcard should make match exhaustive: {}", checker.display_errors());
    }

    #[test]
    fn test_match_return_type() {
        // All arms return String — match type should be String
        let checker = check_source(
            "let status = \"ok\"\nlet msg = match status\n  \"ok\" -> \"good\"\n  _ -> \"bad\""
        );
        assert!(!checker.has_errors(), "str match: {}", checker.display_errors());
    }

    #[test]
    fn test_match_expr_level_type() {
        // Expression-level match with int patterns
        let checker = check_source(
            "let n = 42\nlet label = match n\n  0 -> \"zero\"\n  1 -> \"one\"\n  _ -> \"many\""
        );
        assert!(!checker.has_errors(), "int-pattern match: {}", checker.display_errors());
    }
}


