//! Unification algorithm for type inference.

use crate::TypeError;
use crate::constraint::{Constraint, ConstraintOrigin};
use crate::infer::InferenceContext;
use crate::types::{Type, TypeVarId};
use std::collections::{BTreeMap, HashSet};

/// Converts a Range to a simplified miette::SourceSpan for error reporting.
/// Note: This creates an approximate span based on line/column information.
/// For accurate byte-level spans, the original source text would be needed.
pub fn range_to_span(range: &mq_lang::Range) -> miette::SourceSpan {
    // Create an approximate offset based on line and column
    // This is a simple heuristic: assume average line length of 80 chars
    let line_offset = (range.start.line.saturating_sub(1) as usize) * 80;
    let offset = line_offset + range.start.column.saturating_sub(1);

    // Calculate approximate length based on end position
    let end_line_offset = (range.end.line.saturating_sub(1) as usize) * 80;
    let end_offset = end_line_offset + range.end.column.saturating_sub(1);
    let length = end_offset.saturating_sub(offset).max(1);

    miette::SourceSpan::new(offset.into(), length)
}

/// Solves type constraints through unification
pub fn solve_constraints(ctx: &mut InferenceContext) {
    let constraints = ctx.take_constraints();

    for constraint in constraints {
        match constraint {
            Constraint::Equal(t1, t2, range, origin) => {
                unify(ctx, &t1, &t2, range, &origin);
            }
        }
    }
}

/// Unifies two types
pub fn unify(
    ctx: &mut InferenceContext,
    t1: &Type,
    t2: &Type,
    range: Option<mq_lang::Range>,
    origin: &ConstraintOrigin,
) {
    match (t1, t2) {
        // Never (bottom type) unifies with anything without error
        (Type::Never, _) | (_, Type::Never) => {}

        // Same concrete types unify trivially
        (Type::Int, Type::Int)
        | (Type::Float, Type::Float)
        | (Type::Number, Type::Number)
        | (Type::String, Type::String)
        | (Type::Bool, Type::Bool)
        | (Type::Symbol, Type::Symbol)
        | (Type::None, Type::None)
        | (Type::Markdown, Type::Markdown)
        | (Type::Bytes, Type::Bytes) => {}

        // Type variables
        (Type::Var(v1), Type::Var(v2)) if v1 == v2 => {}

        (Type::Var(_), _) | (_, Type::Var(_)) => {
            // Resolve type variable chains iteratively before unifying
            let t1 = resolve_var_chain(ctx, t1);
            let t2 = resolve_var_chain(ctx, t2);

            // After resolution, if neither is a Var, unify the resolved types
            let (var, ty) = match (&t1, &t2) {
                (Type::Var(v1), Type::Var(v2)) if v1 == v2 => return,
                (Type::Var(var), ty) => (*var, ty),
                (ty, Type::Var(var)) => (*var, ty),
                _ => {
                    unify(ctx, &t1, &t2, range, origin);
                    return;
                }
            };

            // Transitive occurs check: ensure var doesn't occur in ty even through
            // substitution chains, preventing cyclic bindings like
            // Var(a)→Union(Var(b),T) + Var(b)→Union(Var(a),T).
            let mut oc_visited = HashSet::new();
            if occurs_check_transitive(ctx, var, ty, &mut oc_visited) {
                // Binding would create an infinite type (e.g. Var(a) = Tuple(Var(a))).
                // This arises from valid dynamic patterns like `if (is_array(x)): x else: [x]`
                // where type-narrowing hasn't yet constrained the then-branch to Array.
                // Skip the binding silently — the variable remains free/polymorphic, which is
                // acceptable for mq's dynamic type system.
                return;
            }

            // Bind the type variable
            ctx.bind_type_var(var, ty.clone());
        }

        // Arrays
        (Type::Array(elem1), Type::Array(elem2)) => unify(ctx, elem1, elem2, range, origin),

        // Tuples: same-length tuples unify element-wise
        (Type::Tuple(elems1), Type::Tuple(elems2)) => {
            if elems1.len() != elems2.len() {
                ctx.report_mismatch(t1, t2, range, origin);
                return;
            }
            for (e1, e2) in elems1.iter().zip(elems2.iter()) {
                unify(ctx, e1, e2, range, origin);
            }
        }

        // Tuple ↔ Array: unify each tuple element with the array element type
        (Type::Tuple(elems), Type::Array(elem)) | (Type::Array(elem), Type::Tuple(elems)) => {
            for e in elems {
                unify(ctx, e, elem, range, origin);
            }
        }

        // Dictionaries
        (Type::Dict(k1, v1), Type::Dict(k2, v2)) => {
            unify(ctx, k1, k2, range, origin);
            unify(ctx, v1, v2, range, origin);
        }

        // RowEmpty ↔ RowEmpty
        (Type::RowEmpty, Type::RowEmpty) => {}

        // RowEmpty ↔ Dict: a closed row is compatible with any Dict
        // (all known fields were already matched in Record ↔ Dict)
        (Type::RowEmpty, Type::Dict(_, _)) | (Type::Dict(_, _), Type::RowEmpty) => {}

        // RowEmpty ↔ Record: closed row can absorb an empty record
        (Type::RowEmpty, Type::Record(fields, rest)) | (Type::Record(fields, rest), Type::RowEmpty)
            if fields.is_empty() =>
        {
            unify(ctx, rest, &Type::RowEmpty, range, origin);
        }
        (Type::RowEmpty, Type::Record(..)) | (Type::Record(..), Type::RowEmpty) => {
            ctx.report_mismatch(t1, t2, range, origin);
        }

        // Record ↔ Record (row polymorphism)
        (Type::Record(f1, r1), Type::Record(f2, r2)) => {
            unify_records(ctx, f1, r1, f2, r2, range, origin);
        }

        // Record ↔ Dict compatibility
        (Type::Record(fields, rest), Type::Dict(k, v)) | (Type::Dict(k, v), Type::Record(fields, rest)) => {
            // All record keys are strings → unify k with String
            unify(ctx, k, &Type::String, range, origin);
            // All record field values must unify with the dict value type
            for field_ty in fields.values() {
                unify(ctx, field_ty, v, range, origin);
            }
            // The rest of the row must also be compatible with the dict
            unify(ctx, rest, &Type::Dict(k.clone(), v.clone()), range, origin);
        }

        // Functions
        (Type::Function(params1, ret1), Type::Function(params2, ret2)) => {
            if params1.len() != params2.len() {
                ctx.add_error(TypeError::WrongArity {
                    expected: params1.len(),
                    found: params2.len(),
                    span: range.as_ref().map(range_to_span),
                    location: range,
                    context: None,
                });
                return;
            }

            // Unify parameter types
            for (p1, p2) in params1.iter().zip(params2.iter()) {
                unify(ctx, p1, p2, range, origin);
            }

            // Unify return types
            unify(ctx, ret1, ret2, range, origin);
        }

        // Union vs Union: compatible when every member of each union has a corresponding
        // member (by discriminant) in the other union, or when a type variable is present.
        // This must come before the general Union arm to avoid incorrectly comparing each
        // member of the first union against the whole-Union discriminant of the second.
        (Type::Union(types1), Type::Union(types2)) => {
            let member_matches = |m1: &Type, members: &[Type]| {
                members.iter().any(|m2| {
                    std::mem::discriminant(m1) == std::mem::discriminant(m2)
                        || matches!(m1, Type::Var(_))
                        || matches!(m2, Type::Var(_))
                })
            };
            let all_matched = types1.iter().all(|m1| member_matches(m1, types2))
                && types2.iter().all(|m2| member_matches(m2, types1));
            if !all_matched {
                ctx.report_mismatch(t1, t2, range, origin);
            }
        }

        // Union types: a union can unify with a non-union type if any of its members can.
        (Type::Union(types), other) | (other, Type::Union(types)) => {
            // Check if the other type matches any member of the union
            let matches_any = types.iter().any(|t| {
                // Try to check if types can match without adding errors
                match (t, other) {
                    (t1, t2) if std::mem::discriminant(t1) == std::mem::discriminant(t2) => true,
                    (Type::Var(_), _) | (_, Type::Var(_)) => true,
                    _ => false,
                }
            });

            if !matches_any {
                // No member of the union can unify with the other type - report error
                ctx.report_mismatch(t1, t2, range, origin);
            }
            // If at least one member matches, allow it (union type semantics)
        }

        // Mismatch
        _ => {
            ctx.report_mismatch(t1, t2, range, origin);
        }
    }
}

/// Unifies two record types using row polymorphism.
///
/// Given `Record(f1, r1)` and `Record(f2, r2)`:
/// 1. Common fields have their types unified
/// 2. Fields only in f1 are pushed into r2 via a fresh row variable
/// 3. Fields only in f2 are pushed into r1 via a fresh row variable
/// 4. If no unique fields exist on either side, r1 and r2 are unified directly
fn unify_records(
    ctx: &mut InferenceContext,
    f1: &BTreeMap<String, Type>,
    r1: &Type,
    f2: &BTreeMap<String, Type>,
    r2: &Type,
    range: Option<mq_lang::Range>,
    origin: &ConstraintOrigin,
) {
    // Unify common fields
    for (k, v1) in f1 {
        if let Some(v2) = f2.get(k) {
            unify(ctx, v1, v2, range, origin);
        }
    }

    let only1: BTreeMap<String, Type> = f1
        .iter()
        .filter(|(k, _)| !f2.contains_key(k.as_str()))
        .map(|(k, v)| (k.clone(), v.clone()))
        .collect();
    let only2: BTreeMap<String, Type> = f2
        .iter()
        .filter(|(k, _)| !f1.contains_key(k.as_str()))
        .map(|(k, v)| (k.clone(), v.clone()))
        .collect();

    match (only1.is_empty(), only2.is_empty()) {
        (true, true) => {
            // No unique fields — just unify the row tails
            unify(ctx, r1, r2, range, origin);
        }
        (false, true) => {
            // f1 has extra fields; r2 must accommodate them
            let fresh = ctx.fresh_var();
            unify(ctx, r2, &Type::record(only1, Type::Var(fresh)), range, origin);
            unify(ctx, r1, &Type::Var(fresh), range, origin);
        }
        (true, false) => {
            // f2 has extra fields; r1 must accommodate them
            let fresh = ctx.fresh_var();
            unify(ctx, r1, &Type::record(only2, Type::Var(fresh)), range, origin);
            unify(ctx, r2, &Type::Var(fresh), range, origin);
        }
        (false, false) => {
            // Both sides have unique fields
            let fresh = ctx.fresh_var();
            unify(ctx, r1, &Type::record(only2, Type::Var(fresh)), range, origin);
            unify(ctx, r2, &Type::record(only1, Type::Var(fresh)), range, origin);
        }
    }
}

/// Occurs check: ensures a type variable doesn't occur directly in a type.
///
/// This prevents infinite types like T = [T].
#[cfg_attr(not(test), allow(dead_code))]
fn occurs_check(var: TypeVarId, ty: &Type) -> bool {
    match ty {
        Type::Var(v) => var == *v,
        Type::Array(elem) => occurs_check(var, elem),
        Type::Tuple(elems) => elems.iter().any(|e| occurs_check(var, e)),
        Type::Dict(key, value) => occurs_check(var, key) || occurs_check(var, value),
        Type::Function(params, ret) => params.iter().any(|p| occurs_check(var, p)) || occurs_check(var, ret),
        Type::Union(types) => types.iter().any(|t| occurs_check(var, t)),
        Type::Record(fields, rest) => fields.values().any(|v| occurs_check(var, v)) || occurs_check(var, rest),
        _ => false,
    }
}

/// Transitive occurs check that follows the substitution map.
///
/// Prevents cycles like `Var(a) → Union(Var(b), String)` + `Var(b) →
/// Union(Var(a), String)` from being created in the substitution map.
/// Uses a visited set to handle mutual references without looping.
fn occurs_check_transitive(
    ctx: &InferenceContext,
    var: TypeVarId,
    ty: &Type,
    visited: &mut HashSet<TypeVarId>,
) -> bool {
    match ty {
        Type::Var(v) => {
            if *v == var {
                return true;
            }
            if !visited.insert(*v) {
                return false; // Already explored this path
            }
            if let Some(bound) = ctx.get_type_var(*v) {
                occurs_check_transitive(ctx, var, &bound, visited)
            } else {
                false
            }
        }
        Type::Array(elem) => occurs_check_transitive(ctx, var, elem, visited),
        Type::Tuple(elems) => elems.iter().any(|e| occurs_check_transitive(ctx, var, e, visited)),
        Type::Dict(key, value) => {
            occurs_check_transitive(ctx, var, key, visited) || occurs_check_transitive(ctx, var, value, visited)
        }
        Type::Function(params, ret) => {
            params.iter().any(|p| occurs_check_transitive(ctx, var, p, visited))
                || occurs_check_transitive(ctx, var, ret, visited)
        }
        Type::Union(types) => types.iter().any(|t| occurs_check_transitive(ctx, var, t, visited)),
        Type::Record(fields, rest) => {
            fields.values().any(|v| occurs_check_transitive(ctx, var, v, visited))
                || occurs_check_transitive(ctx, var, rest, visited)
        }
        _ => false,
    }
}

/// Applies substitutions to resolve all type variables.
///
/// Uses a visited set to detect and break cycles in the substitution map,
/// preventing stack overflow from mutually recursive type variable bindings.
pub fn apply_substitution(ctx: &InferenceContext, ty: &Type) -> Type {
    let mut visited = HashSet::new();
    apply_substitution_inner(ctx, ty, &mut visited)
}

fn apply_substitution_inner(ctx: &InferenceContext, ty: &Type, visited: &mut HashSet<TypeVarId>) -> Type {
    match ty {
        Type::Var(var) => {
            if !visited.insert(*var) {
                return ty.clone(); // Cycle detected — return var unresolved
            }
            let result = if let Some(bound) = ctx.get_type_var(*var) {
                apply_substitution_inner(ctx, &bound, visited)
            } else {
                ty.clone()
            };
            visited.remove(var);
            result
        }
        Type::Array(elem) => Type::Array(Box::new(apply_substitution_inner(ctx, elem, visited))),
        Type::Tuple(elems) => Type::Tuple(
            elems
                .iter()
                .map(|e| apply_substitution_inner(ctx, e, visited))
                .collect(),
        ),
        Type::Dict(key, value) => Type::Dict(
            Box::new(apply_substitution_inner(ctx, key, visited)),
            Box::new(apply_substitution_inner(ctx, value, visited)),
        ),
        Type::Function(params, ret) => {
            let new_params = params
                .iter()
                .map(|p| apply_substitution_inner(ctx, p, visited))
                .collect();
            Type::Function(new_params, Box::new(apply_substitution_inner(ctx, ret, visited)))
        }
        Type::Union(types) => {
            let new_types = types
                .iter()
                .map(|t| apply_substitution_inner(ctx, t, visited))
                .collect();
            Type::union(new_types)
        }
        Type::Record(fields, rest) => {
            let new_fields = fields
                .iter()
                .map(|(k, v)| (k.clone(), apply_substitution_inner(ctx, v, visited)))
                .collect();
            Type::Record(new_fields, Box::new(apply_substitution_inner(ctx, rest, visited)))
        }
        _ => ty.clone(),
    }
}

/// Follows a type variable substitution chain iteratively until reaching
/// a non-variable type or an unbound variable.
fn resolve_var_chain(ctx: &InferenceContext, ty: &Type) -> Type {
    let mut current = ty.clone();
    loop {
        match &current {
            Type::Var(var) => {
                if let Some(bound) = ctx.get_type_var(*var) {
                    current = bound;
                } else {
                    return current;
                }
            }
            _ => return current,
        }
    }
}

/// Gets all free type variables in a type after applying current substitutions
pub fn free_vars(ctx: &InferenceContext, ty: &Type) -> HashSet<TypeVarId> {
    let resolved = apply_substitution(ctx, ty);
    let mut vars = HashSet::new();
    collect_free_vars(&resolved, &mut vars);
    vars
}

fn collect_free_vars(ty: &Type, vars: &mut HashSet<TypeVarId>) {
    match ty {
        Type::Var(var) => {
            vars.insert(*var);
        }
        Type::Array(elem) => collect_free_vars(elem, vars),
        Type::Tuple(elems) => {
            for e in elems {
                collect_free_vars(e, vars);
            }
        }
        Type::Dict(key, value) => {
            collect_free_vars(key, vars);
            collect_free_vars(value, vars);
        }
        Type::Function(params, ret) => {
            for param in params {
                collect_free_vars(param, vars);
            }
            collect_free_vars(ret, vars);
        }
        Type::Record(fields, rest) => {
            for v in fields.values() {
                collect_free_vars(v, vars);
            }
            collect_free_vars(rest, vars);
        }
        _ => {}
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::TypeVarContext;

    #[test]
    fn test_unify_concrete_types() {
        let mut ctx = InferenceContext::new();
        unify(&mut ctx, &Type::Number, &Type::Number, None, &ConstraintOrigin::General);
        assert!(ctx.take_errors().is_empty());

        unify(&mut ctx, &Type::String, &Type::Number, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());
    }

    #[test]
    fn test_unify_type_vars() {
        let mut var_ctx = TypeVarContext::new();
        let mut ctx = InferenceContext::new();

        let var1 = var_ctx.fresh();
        let var2 = var_ctx.fresh();

        // Unify var1 with Number
        unify(
            &mut ctx,
            &Type::Var(var1),
            &Type::Number,
            None,
            &ConstraintOrigin::General,
        );
        assert!(ctx.take_errors().is_empty());

        // Unify var2 with var1 (should transitively become Number)
        unify(
            &mut ctx,
            &Type::Var(var2),
            &Type::Var(var1),
            None,
            &ConstraintOrigin::General,
        );
        assert!(ctx.take_errors().is_empty());
    }

    #[test]
    fn test_unify_arrays() {
        let mut ctx = InferenceContext::new();
        let arr1 = Type::array(Type::Number);
        let arr2 = Type::array(Type::Number);
        unify(&mut ctx, &arr1, &arr2, None, &ConstraintOrigin::General);
        assert!(ctx.take_errors().is_empty());

        let arr3 = Type::array(Type::String);
        unify(&mut ctx, &arr1, &arr3, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());
    }

    #[test]
    fn test_occurs_check() {
        let mut var_ctx = TypeVarContext::new();
        let var = var_ctx.fresh();

        // T = [T] should fail occurs check
        let recursive = Type::array(Type::Var(var));
        assert!(occurs_check(var, &recursive));

        // T = number should pass occurs check
        assert!(!occurs_check(var, &Type::Number));
    }

    #[test]
    fn test_range_to_span() {
        let range = mq_lang::Range {
            start: mq_lang::Position { line: 1, column: 1 },
            end: mq_lang::Position { line: 1, column: 11 },
        };
        let span = range_to_span(&range);
        assert_eq!(span.offset(), 0);
        assert_eq!(span.len(), 10);

        let range2 = mq_lang::Range {
            start: mq_lang::Position { line: 2, column: 5 },
            end: mq_lang::Position { line: 2, column: 10 },
        };
        let span2 = range_to_span(&range2);
        assert_eq!(span2.offset(), 80 + 4);
        assert_eq!(span2.len(), 5);
    }

    #[test]
    fn test_unify_tuples() {
        let mut ctx = InferenceContext::new();
        let t1 = Type::tuple(vec![Type::Number, Type::String]);
        let t2 = Type::tuple(vec![Type::Number, Type::String]);
        unify(&mut ctx, &t1, &t2, None, &ConstraintOrigin::General);
        assert!(ctx.take_errors().is_empty());

        let t3 = Type::tuple(vec![Type::Number]);
        unify(&mut ctx, &t1, &t3, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());

        let t4 = Type::tuple(vec![Type::String, Type::Number]);
        unify(&mut ctx, &t1, &t4, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());
    }

    #[test]
    fn test_unify_tuple_array() {
        let mut ctx = InferenceContext::new();
        let tuple = Type::tuple(vec![Type::Number, Type::Number]);
        let array = Type::array(Type::Number);

        unify(&mut ctx, &tuple, &array, None, &ConstraintOrigin::General);
        assert!(ctx.take_errors().is_empty());

        let tuple2 = Type::tuple(vec![Type::Number, Type::String]);
        unify(&mut ctx, &tuple2, &array, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());
    }

    #[test]
    fn test_unify_dicts() {
        let mut ctx = InferenceContext::new();
        let d1 = Type::dict(Type::String, Type::Number);
        let d2 = Type::dict(Type::String, Type::Number);
        unify(&mut ctx, &d1, &d2, None, &ConstraintOrigin::General);
        assert!(ctx.take_errors().is_empty());

        let d3 = Type::dict(Type::Number, Type::Number);
        unify(&mut ctx, &d1, &d3, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());
    }

    #[test]
    fn test_unify_records() {
        let mut ctx = InferenceContext::new();

        // Identical records
        let r1 = Type::record([("a".to_string(), Type::Number)].into_iter().collect(), Type::RowEmpty);
        let r2 = Type::record([("a".to_string(), Type::Number)].into_iter().collect(), Type::RowEmpty);
        unify(&mut ctx, &r1, &r2, None, &ConstraintOrigin::General);
        assert!(ctx.take_errors().is_empty());

        // Field type mismatch
        let r3 = Type::record([("a".to_string(), Type::String)].into_iter().collect(), Type::RowEmpty);
        unify(&mut ctx, &r1, &r3, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());

        // Row polymorphism: open record
        let var = ctx.fresh_var();
        let r_open = Type::record([("a".to_string(), Type::Number)].into_iter().collect(), Type::Var(var));
        let r_closed = Type::record(
            [("a".to_string(), Type::Number), ("b".to_string(), Type::String)]
                .into_iter()
                .collect(),
            Type::RowEmpty,
        );
        unify(&mut ctx, &r_open, &r_closed, None, &ConstraintOrigin::General);
        assert!(ctx.take_errors().is_empty());

        // Check if var was bound to {b: String}
        let resolved_var = ctx.resolve_type(&Type::Var(var));
        match resolved_var {
            Type::Record(fields, rest) => {
                assert_eq!(fields.get("b"), Some(&Type::String));
                assert_eq!(*rest, Type::RowEmpty);
            }
            _ => panic!("Expected record type, found {:?}", resolved_var),
        }
    }

    #[test]
    fn test_unify_record_dict() {
        let mut ctx = InferenceContext::new();
        let record = Type::record([("a".to_string(), Type::Number)].into_iter().collect(), Type::RowEmpty);
        let dict = Type::dict(Type::String, Type::Number);

        unify(&mut ctx, &record, &dict, None, &ConstraintOrigin::General);
        assert!(ctx.take_errors().is_empty());

        let dict2 = Type::dict(Type::String, Type::String);
        unify(&mut ctx, &record, &dict2, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());
    }

    #[test]
    fn test_unify_unions() {
        let mut ctx = InferenceContext::new();
        let union = Type::union(vec![Type::Number, Type::String]);

        // Unify union with one of its members
        unify(&mut ctx, &union, &Type::Number, None, &ConstraintOrigin::General);
        assert!(ctx.take_errors().is_empty());

        unify(&mut ctx, &Type::String, &union, None, &ConstraintOrigin::General);
        assert!(ctx.take_errors().is_empty());

        // Unify union with incompatible type
        unify(&mut ctx, &union, &Type::Bool, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());

        // Unify union with another union that has different members — should fail.
        // Union(Number, String) vs Union(String, Bool): Number has no match in [String, Bool].
        let union2 = Type::union(vec![Type::String, Type::Bool]);
        unify(&mut ctx, &union, &union2, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());

        // Unify identical unions — should succeed.
        let union3 = Type::union(vec![Type::Number, Type::String]);
        unify(&mut ctx, &union, &union3, None, &ConstraintOrigin::General);
        assert!(ctx.take_errors().is_empty());
    }

    #[test]
    fn test_unify_functions() {
        let mut ctx = InferenceContext::new();
        let f1 = Type::function(vec![Type::Number], Type::String);
        let f2 = Type::function(vec![Type::Number], Type::String);
        unify(&mut ctx, &f1, &f2, None, &ConstraintOrigin::General);
        assert!(ctx.take_errors().is_empty());

        // Arity mismatch
        let f3 = Type::function(vec![Type::Number, Type::Number], Type::String);
        unify(&mut ctx, &f1, &f3, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());

        // Param mismatch
        let f4 = Type::function(vec![Type::String], Type::String);
        unify(&mut ctx, &f1, &f4, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());

        // Return mismatch
        let f5 = Type::function(vec![Type::Number], Type::Number);
        unify(&mut ctx, &f1, &f5, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());
    }

    #[test]
    fn test_occurs_check_transitive() {
        let mut ctx = InferenceContext::new();
        let v1 = ctx.fresh_var();
        let v2 = ctx.fresh_var();

        // v1 -> v2
        ctx.bind_type_var(v1, Type::Var(v2));

        let mut visited = HashSet::new();
        // v2 occurs in v1? Yes, because v1 -> v2
        assert!(occurs_check_transitive(&ctx, v2, &Type::Var(v1), &mut visited));

        // Cycle: v1 -> v2, v2 -> v1
        ctx.bind_type_var(v2, Type::Var(v1));
        let mut visited = HashSet::new();
        // Should not stack overflow
        assert!(occurs_check_transitive(&ctx, v1, &Type::Var(v1), &mut visited));
    }

    #[test]
    fn test_apply_substitution() {
        let mut ctx = InferenceContext::new();
        let v1 = ctx.fresh_var();
        let v2 = ctx.fresh_var();

        ctx.bind_type_var(v1, Type::array(Type::Var(v2)));
        ctx.bind_type_var(v2, Type::Number);

        let t = Type::Var(v1);
        let applied = apply_substitution(&ctx, &t);
        assert_eq!(applied, Type::array(Type::Number));

        // Cycle
        let v3 = ctx.fresh_var();
        let v4 = ctx.fresh_var();
        ctx.bind_type_var(v3, Type::Var(v4));
        ctx.bind_type_var(v4, Type::Var(v3));
        let applied_cycle = apply_substitution(&ctx, &Type::Var(v3));
        assert!(matches!(applied_cycle, Type::Var(_)));
    }

    #[test]
    fn test_free_vars() {
        let mut ctx = InferenceContext::new();
        let v1 = ctx.fresh_var();
        let v2 = ctx.fresh_var();

        let t = Type::tuple(vec![Type::Var(v1), Type::array(Type::Var(v2))]);
        let vars = free_vars(&ctx, &t);
        assert_eq!(vars.len(), 2);
        assert!(vars.contains(&v1));
        assert!(vars.contains(&v2));

        ctx.bind_type_var(v1, Type::Number);
        let vars2 = free_vars(&ctx, &t);
        assert_eq!(vars2.len(), 1);
        assert!(vars2.contains(&v2));
    }

    #[test]
    fn test_solve_constraints() {
        let mut ctx = InferenceContext::new();
        let v1 = ctx.fresh_var();
        ctx.add_constraint(Constraint::Equal(
            Type::Var(v1),
            Type::Number,
            None,
            ConstraintOrigin::General,
        ));

        solve_constraints(&mut ctx);
        assert_eq!(ctx.resolve_type(&Type::Var(v1)), Type::Number);
    }

    #[test]
    fn test_unify_row_empty() {
        let mut ctx = InferenceContext::new();

        // RowEmpty <-> RowEmpty
        unify(
            &mut ctx,
            &Type::RowEmpty,
            &Type::RowEmpty,
            None,
            &ConstraintOrigin::General,
        );
        assert!(ctx.take_errors().is_empty());

        // RowEmpty <-> Dict
        let dict = Type::dict(Type::String, Type::Number);
        unify(&mut ctx, &Type::RowEmpty, &dict, None, &ConstraintOrigin::General);
        assert!(ctx.take_errors().is_empty());

        // RowEmpty <-> Empty Record
        let empty_record = Type::record(BTreeMap::new(), Type::RowEmpty);
        unify(
            &mut ctx,
            &Type::RowEmpty,
            &empty_record,
            None,
            &ConstraintOrigin::General,
        );
        assert!(ctx.take_errors().is_empty());

        // RowEmpty <-> Non-empty Record (should fail)
        let record = Type::record([("a".to_string(), Type::Number)].into_iter().collect(), Type::RowEmpty);
        unify(&mut ctx, &Type::RowEmpty, &record, None, &ConstraintOrigin::General);
        assert!(!ctx.take_errors().is_empty());
    }
}
