mq_check/

types.rs

//! Type representations for the mq type system.

use rustc_hash::FxHashMap;
use slotmap::SlotMap;
use std::collections::BTreeMap;
use std::fmt;

slotmap::new_key_type! {
    /// Unique identifier for type variables
    pub struct TypeVarId;
}

/// Represents a type in the mq type system
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Type {
    /// Integer type
    Int,
    /// Floating point type
    Float,
    /// Number type (unified numeric type)
    Number,
    /// String type
    String,
    /// Boolean type
    Bool,
    /// Symbol type
    Symbol,
    /// None/null type
    None,
    /// Markdown document type
    Markdown,
    /// Array type with element type
    Array(Box<Type>),
    /// Tuple type with known element types (e.g., `(number, string)`)
    ///
    /// Used when `--tuple` mode is enabled to track per-element types
    /// for small heterogeneous arrays like `[1, "hello"]`.
    Tuple(Vec<Type>),
    /// Dictionary type with key and value types
    Dict(Box<Type>, Box<Type>),
    /// Function type: arguments -> return type
    Function(Vec<Type>, Box<Type>),
    /// Union type: represents a value that could be one of multiple types
    /// Used for try/catch expressions with different branch types
    Union(Vec<Type>),
    /// Record type with known fields and optional row extension (row polymorphism).
    ///
    /// The first element is a map of field names to their types.
    /// The second element is the row tail: either `RowEmpty` (closed record)
    /// or `Var(id)` (open record that may have additional fields).
    ///
    /// Examples:
    /// - `{a: number, b: string}` → `Record({"a": Number, "b": String}, RowEmpty)`
    /// - `{a: number | r}` → `Record({"a": Number}, Var(r))`
    Record(BTreeMap<String, Type>, Box<Type>),
    /// Empty row — marks a closed record with no additional fields
    RowEmpty,
    /// Type variable for inference
    Var(TypeVarId),
    /// Never type (bottom type) — represents unreachable code paths.
    ///
    /// Produced when type narrowing eliminates all possible types from a union.
    /// For example, `Union(Array(String), Array(Number))` minus `is_array` predicate
    /// leaves no viable types, producing `Never` for the else-branch.
    Never,
}

impl Type {
    /// Creates a new function type
    pub fn function(params: Vec<Type>, ret: Type) -> Self {
        Type::Function(params, Box::new(ret))
    }

    /// Creates a new array type
    pub fn array(elem: Type) -> Self {
        Type::Array(Box::new(elem))
    }

    /// Creates a new tuple type with known element types
    pub fn tuple(elems: Vec<Type>) -> Self {
        Type::Tuple(elems)
    }

    /// Creates a new dict type
    pub fn dict(key: Type, value: Type) -> Self {
        Type::Dict(Box::new(key), Box::new(value))
    }

    /// Creates a new record type with known fields
    pub fn record(fields: BTreeMap<String, Type>, rest: Type) -> Self {
        Type::Record(fields, Box::new(rest))
    }

    /// Creates a new union type from two or more types
    /// Automatically normalizes the union (removes duplicates, flattens nested unions)
    pub fn union(types: Vec<Type>) -> Self {
        let mut normalized = Vec::with_capacity(types.len());
        for ty in types {
            match ty {
                // Flatten nested unions
                Type::Union(inner) => normalized.extend(inner),
                _ => normalized.push(ty),
            }
        }

        // Deduplicate and sort
        if normalized.len() <= 4 {
            // For small unions, avoid HashSet allocation with an O(n²) scan
            let mut i = 0;
            while i < normalized.len() {
                if normalized[..i].iter().any(|t| t == &normalized[i]) {
                    normalized.swap_remove(i);
                } else {
                    i += 1;
                }
            }
            normalized.sort_by_key(|t| t.discriminant());
        } else {
            // For larger unions, use a HashSet for O(n) deduplication
            let mut seen = rustc_hash::FxHashSet::default();
            normalized.retain(|t| seen.insert(t.clone()));
            normalized.sort_by_key(|t| t.discriminant());
        }

        // If only one type remains, return it directly
        if normalized.len() == 1 {
            normalized.into_iter().next().unwrap()
        } else {
            Type::Union(normalized)
        }
    }

    /// Removes a type from a union by discriminant, returning the remaining type.
    ///
    /// If this is not a union, returns self unchanged.
    /// If only one type remains after subtraction, returns it directly (unwrapped).
    /// If all members are removed (e.g., `Union(Array(String), Array(Number))` minus
    /// `Array`), returns `Type::Never` to signal an unreachable code path.
    pub fn subtract(&self, exclude: &Type) -> Type {
        match self {
            Type::Union(members) => {
                let remaining: Vec<Type> = members
                    .iter()
                    .filter(|t| std::mem::discriminant(*t) != std::mem::discriminant(exclude))
                    .cloned()
                    .collect();
                if remaining.is_empty() {
                    Type::Never
                } else {
                    Type::union(remaining)
                }
            }
            _ => self.clone(),
        }
    }

    /// Checks if this type is nullable (is `None` or a `Union` containing `None`).
    pub fn is_nullable(&self) -> bool {
        match self {
            Type::None => true,
            Type::Union(members) => members.iter().any(|m| matches!(m, Type::None)),
            _ => false,
        }
    }

    /// Checks if this is the bottom type (never/unreachable).
    pub fn is_never(&self) -> bool {
        matches!(self, Type::Never)
    }

    /// Returns a numeric discriminant for ordering purposes.
    /// Used by `Type::union` to sort union members without allocating.
    fn discriminant(&self) -> u8 {
        match self {
            Type::Int => 0,
            Type::Float => 1,
            Type::Number => 2,
            Type::String => 3,
            Type::Bool => 4,
            Type::Symbol => 5,
            Type::None => 6,
            Type::Markdown => 7,
            Type::Array(_) => 8,
            Type::Tuple(_) => 9,
            Type::Dict(_, _) => 10,
            Type::Function(_, _) => 11,
            Type::Union(_) => 12,
            Type::Record(_, _) => 13,
            Type::RowEmpty => 14,
            Type::Var(_) => 15,
            Type::Never => 16,
        }
    }

    /// Checks if this is a type variable
    pub fn is_var(&self) -> bool {
        matches!(self, Type::Var(_))
    }

    /// Checks if this type contains no free type variables (is fully concrete)
    pub fn is_concrete(&self) -> bool {
        self.free_vars().is_empty()
    }

    /// Checks if this is a union type
    pub fn is_union(&self) -> bool {
        matches!(self, Type::Union(_))
    }

    /// Gets the type variable ID if this is a type variable
    pub fn as_var(&self) -> Option<TypeVarId> {
        match self {
            Type::Var(id) => Some(*id),
            _ => None,
        }
    }

    /// Substitutes type variables according to the given substitution
    pub fn apply_subst(&self, subst: &Substitution) -> Type {
        if subst.is_empty() {
            return self.clone();
        }
        match self {
            Type::Var(id) => subst.lookup(*id).map_or_else(|| self.clone(), |t| t.apply_subst(subst)),
            Type::Array(elem) => Type::Array(Box::new(elem.apply_subst(subst))),
            Type::Tuple(elems) => Type::Tuple(elems.iter().map(|e| e.apply_subst(subst)).collect()),
            Type::Dict(key, value) => Type::Dict(Box::new(key.apply_subst(subst)), Box::new(value.apply_subst(subst))),
            Type::Function(params, ret) => {
                let new_params = params.iter().map(|p| p.apply_subst(subst)).collect();
                Type::Function(new_params, Box::new(ret.apply_subst(subst)))
            }
            Type::Union(types) => {
                let new_types = types.iter().map(|t| t.apply_subst(subst)).collect();
                Type::union(new_types)
            }
            Type::Record(fields, rest) => {
                let new_fields = fields.iter().map(|(k, v)| (k.clone(), v.apply_subst(subst))).collect();
                Type::Record(new_fields, Box::new(rest.apply_subst(subst)))
            }
            // Never and other ground types have no type variables to substitute
            _ => self.clone(),
        }
    }

    /// Gets all free type variables in this type
    pub fn free_vars(&self) -> Vec<TypeVarId> {
        match self {
            Type::Var(id) => vec![*id],
            Type::Array(elem) => elem.free_vars(),
            Type::Tuple(elems) => elems.iter().flat_map(|e| e.free_vars()).collect(),
            Type::Dict(key, value) => {
                let mut vars = key.free_vars();
                vars.extend(value.free_vars());
                vars
            }
            Type::Function(params, ret) => {
                let mut vars: Vec<TypeVarId> = params.iter().flat_map(|p| p.free_vars()).collect();
                vars.extend(ret.free_vars());
                vars
            }
            Type::Union(types) => types.iter().flat_map(|t| t.free_vars()).collect(),
            Type::Record(fields, rest) => {
                let mut vars: Vec<TypeVarId> = fields.values().flat_map(|v| v.free_vars()).collect();
                vars.extend(rest.free_vars());
                vars
            }
            _ => Vec::new(),
        }
    }

    /// Checks if this type can match with another type (for overload resolution).
    ///
    /// This is a weaker check than unification - it returns true if the types
    /// could potentially be unified, but doesn't require them to be identical.
    /// Type variables always match.
    pub fn can_match(&self, other: &Type) -> bool {
        match (self, other) {
            // Never (bottom type) can match anything
            (Type::Never, _) | (_, Type::Never) => true,

            // Type variables always match
            (Type::Var(_), _) | (_, Type::Var(_)) => true,

            // Union types match if any of their constituent types can match
            (Type::Union(types), other) => types.iter().any(|t| t.can_match(other)),
            (other, Type::Union(types)) => types.iter().any(|t| other.can_match(t)),

            // Concrete types must match exactly
            (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) => true,

            // Arrays match if their element types can match
            (Type::Array(elem1), Type::Array(elem2)) => elem1.can_match(elem2),

            // Tuples match if they have the same length and all elements can match
            (Type::Tuple(elems1), Type::Tuple(elems2)) => {
                elems1.len() == elems2.len() && elems1.iter().zip(elems2.iter()).all(|(e1, e2)| e1.can_match(e2))
            }

            // Tuple can match Array (compatibility)
            (Type::Tuple(_), Type::Array(_)) | (Type::Array(_), Type::Tuple(_)) => true,

            // Dicts match if both key and value types can match
            (Type::Dict(k1, v1), Type::Dict(k2, v2)) => k1.can_match(k2) && v1.can_match(v2),

            // Functions match if they have the same arity and all parameter/return types can match
            (Type::Function(params1, ret1), Type::Function(params2, ret2)) => {
                params1.len() == params2.len()
                    && params1.iter().zip(params2.iter()).all(|(p1, p2)| p1.can_match(p2))
                    && ret1.can_match(ret2)
            }

            // Records match if common fields can match and rests can match
            (Type::Record(f1, r1), Type::Record(f2, r2)) => {
                // Common fields must match
                for (k, v1) in f1 {
                    if let Some(v2) = f2.get(k)
                        && !v1.can_match(v2)
                    {
                        return false;
                    }
                }
                r1.can_match(r2)
            }

            // Record can match Dict (compatibility)
            (Type::Record(_, _), Type::Dict(_, _)) | (Type::Dict(_, _), Type::Record(_, _)) => true,

            // RowEmpty matches RowEmpty
            (Type::RowEmpty, Type::RowEmpty) => true,

            // Everything else doesn't match
            _ => false,
        }
    }

    /// Strict branch compatibility check for if-expression unification decisions.
    ///
    /// Unlike `can_match`, this treats an unresolved type variable (`Var`) as incompatible
    /// with any known concrete type. This is used when deciding whether if-expression
    /// branches should be unified or placed in a Union type: if one branch has `None` in
    /// a tuple element position and another has `Var` (which might resolve to `Number`),
    /// we conservatively choose Union to avoid a spurious unification error later.
    ///
    /// - `Var(a)` vs `Var(b)` → `true` (both unknown; let unification decide)
    /// - `Var(a)` vs `concrete` → `false` (Var might not match the concrete type)
    /// - `concrete` vs `Var(a)` → `false` (symmetric)
    /// - All other cases mirror `can_match`
    pub fn can_branch_unify_with(&self, other: &Type) -> bool {
        match (self, other) {
            // Never (bottom type) can unify with anything
            (Type::Never, _) | (_, Type::Never) => true,

            // Two type variables are compatible (unification will sort them out)
            (Type::Var(_), Type::Var(_)) => true,

            // Var vs concrete — unknown; do NOT assume compatible
            (Type::Var(_), _) | (_, Type::Var(_)) => false,

            // Union types: strict — any member must strictly match
            (Type::Union(types), other) => types.iter().any(|t| t.can_branch_unify_with(other)),
            (other, Type::Union(types)) => types.iter().any(|t| other.can_branch_unify_with(t)),

            // Concrete types must match exactly
            (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) => true,

            // Arrays: recurse strictly
            (Type::Array(elem1), Type::Array(elem2)) => elem1.can_branch_unify_with(elem2),

            // Tuples: same length and all elements strictly match
            (Type::Tuple(elems1), Type::Tuple(elems2)) => {
                elems1.len() == elems2.len()
                    && elems1
                        .iter()
                        .zip(elems2.iter())
                        .all(|(e1, e2)| e1.can_branch_unify_with(e2))
            }

            // Dicts: recurse strictly
            (Type::Dict(k1, v1), Type::Dict(k2, v2)) => k1.can_branch_unify_with(k2) && v1.can_branch_unify_with(v2),

            // Functions: same arity and all param/ret types strictly match
            (Type::Function(params1, ret1), Type::Function(params2, ret2)) => {
                params1.len() == params2.len()
                    && params1
                        .iter()
                        .zip(params2.iter())
                        .all(|(p1, p2)| p1.can_branch_unify_with(p2))
                    && ret1.can_branch_unify_with(ret2)
            }

            // Everything else is incompatible
            _ => false,
        }
    }

    /// Computes a match score for overload resolution.
    /// Higher scores indicate better matches. Returns None if types cannot match.
    ///
    /// Scoring:
    /// - Exact match: 100
    /// - Union type: best match among variants (slightly penalized)
    /// - Type variable: 10
    /// - Structural match (array/dict/function): sum of component scores
    pub fn match_score(&self, other: &Type) -> Option<u32> {
        if !self.can_match(other) {
            return None;
        }

        match (self, other) {
            // Never (bottom type) can match anything but with lowest score
            (Type::Never, _) | (_, Type::Never) => Some(1),

            // Exact matches get highest score
            (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) => Some(100),

            // Type variables get low score (prefer concrete types).
            // This arm must come BEFORE the union arms so that a Var parameter
            // (e.g. `to_number: (Var) -> Number`) scores 10 against a union arg
            // rather than 0 (the union arm penalises by -15, giving 10-15=0).
            (Type::Var(_), _) | (_, Type::Var(_)) => Some(10),

            // Union types: take the best match among all variants, but penalize
            (Type::Union(types), other) => types
                .iter()
                .filter_map(|t| t.match_score(other))
                .max()
                .map(|s| s.saturating_sub(15)),
            (other, Type::Union(types)) => types
                .iter()
                .filter_map(|t| other.match_score(t))
                .max()
                .map(|s| s.saturating_sub(15)),

            // Arrays: structural match scores higher than bare type variable
            (Type::Array(elem1), Type::Array(elem2)) => elem1.match_score(elem2).map(|s| s + 20),

            // Tuples: structural match on all elements
            (Type::Tuple(elems1), Type::Tuple(elems2)) if elems1.len() == elems2.len() => {
                let total: u32 = elems1
                    .iter()
                    .zip(elems2.iter())
                    .map(|(e1, e2)| e1.match_score(e2).unwrap_or(0))
                    .sum();
                Some(total / elems1.len() as u32 + 20)
            }

            // Tuple ↔ Array compatibility (lower score than direct Tuple match)
            (Type::Tuple(_), Type::Array(_)) | (Type::Array(_), Type::Tuple(_)) => Some(15),

            // Dicts: structural match scores higher than bare type variable
            (Type::Dict(k1, v1), Type::Dict(k2, v2)) => {
                let key_score = k1.match_score(k2)?;
                let val_score = v1.match_score(v2)?;
                Some((key_score + val_score) / 2 + 20)
            }

            // Records: structural match on fields
            (Type::Record(f1, r1), Type::Record(f2, r2)) => {
                let mut total = 0u32;
                let mut count = 0u32;
                for (k, v1) in f1 {
                    if let Some(v2) = f2.get(k) {
                        total += v1.match_score(v2)?;
                        count += 1;
                    }
                }
                let field_score = if count > 0 { total / count } else { 10 };
                let rest_score = r1.match_score(r2).unwrap_or(10);
                Some(field_score + rest_score + 20)
            }

            // Record ↔ Dict compatibility (lower score than direct Record match)
            (Type::Record(_, _), Type::Dict(_, _)) | (Type::Dict(_, _), Type::Record(_, _)) => Some(15),

            (Type::RowEmpty, Type::RowEmpty) => Some(100),

            // Functions: sum all parameter scores and return score
            (Type::Function(params1, ret1), Type::Function(params2, ret2)) => {
                let param_score: u32 = params1
                    .iter()
                    .zip(params2.iter())
                    .map(|(p1, p2)| p1.match_score(p2).unwrap_or(0))
                    .sum();
                let ret_score = ret1.match_score(ret2)?;
                Some(param_score + ret_score)
            }

            _ => None,
        }
    }
}

impl Type {
    /// Formats the type as a string, resolving type variables to their readable names.
    /// This is used for better error messages.
    pub fn display_resolved(&self) -> String {
        match self {
            Type::Int => "int".to_string(),
            Type::Float => "float".to_string(),
            Type::Number => "number".to_string(),
            Type::String => "string".to_string(),
            Type::Bool => "bool".to_string(),
            Type::Symbol => "symbol".to_string(),
            Type::None => "none".to_string(),
            Type::Markdown => "markdown".to_string(),
            Type::Array(elem) => format!("[{}]", elem.display_resolved()),
            Type::Tuple(elems) => {
                let elems_str = elems
                    .iter()
                    .map(|e| e.display_resolved())
                    .collect::<Vec<_>>()
                    .join(", ");
                format!("({})", elems_str)
            }
            Type::Dict(key, value) => format!("{{{}: {}}}", key.display_resolved(), value.display_resolved()),
            Type::Record(fields, rest) => {
                let fields_str = fields
                    .iter()
                    .map(|(k, v)| format!("{}: {}", k, v.display_resolved()))
                    .collect::<Vec<_>>()
                    .join(", ");
                match rest.as_ref() {
                    Type::RowEmpty => format!("{{{}}}", fields_str),
                    _ => {
                        if fields_str.is_empty() {
                            format!("{{| {}}}", rest.display_resolved())
                        } else {
                            format!("{{{} | {}}}", fields_str, rest.display_resolved())
                        }
                    }
                }
            }
            Type::RowEmpty => "{}".to_string(),
            Type::Function(params, ret) => {
                let params_str = params
                    .iter()
                    .map(|p| p.display_resolved())
                    .collect::<Vec<_>>()
                    .join(", ");
                format!("({}) -> {}", params_str, ret.display_resolved())
            }
            Type::Union(types) => {
                let types_str = types
                    .iter()
                    .map(|t| t.display_resolved())
                    .collect::<Vec<_>>()
                    .join(" | ");
                format!("({})", types_str)
            }
            Type::Var(id) => {
                // Convert TypeVarId to a readable name like 'a, 'b, 'c, etc.
                type_var_name(*id)
            }
            Type::Never => "never".to_string(),
        }
    }

    /// Formats the type with renumbered type variables starting from `'a`.
    ///
    /// This produces clean, sequential type variable names regardless of internal
    /// slotmap indices, which can be very large due to builtin registrations.
    /// For example, instead of `'y32` or `'x3`, this produces `'a`, `'b`, etc.
    pub fn display_renumbered(&self) -> String {
        let mut var_map = FxHashMap::default();
        let mut counter = 0usize;
        self.fmt_renumbered(&mut var_map, &mut counter)
    }

    /// Internal helper for renumbered formatting.
    pub(crate) fn fmt_renumbered(&self, var_map: &mut FxHashMap<TypeVarId, usize>, counter: &mut usize) -> String {
        match self {
            Type::Int => "int".to_string(),
            Type::Float => "float".to_string(),
            Type::Number => "number".to_string(),
            Type::String => "string".to_string(),
            Type::Bool => "bool".to_string(),
            Type::Symbol => "symbol".to_string(),
            Type::None => "none".to_string(),
            Type::Markdown => "markdown".to_string(),
            Type::Array(elem) => format!("[{}]", elem.fmt_renumbered(var_map, counter)),
            Type::Tuple(elems) => {
                let elems_str = elems
                    .iter()
                    .map(|e| e.fmt_renumbered(var_map, counter))
                    .collect::<Vec<_>>()
                    .join(", ");
                format!("({})", elems_str)
            }
            Type::Dict(key, value) => {
                format!(
                    "{{{}: {}}}",
                    key.fmt_renumbered(var_map, counter),
                    value.fmt_renumbered(var_map, counter)
                )
            }
            Type::Function(params, ret) => {
                let params_str = params
                    .iter()
                    .map(|p| p.fmt_renumbered(var_map, counter))
                    .collect::<Vec<_>>()
                    .join(", ");
                format!("({}) -> {}", params_str, ret.fmt_renumbered(var_map, counter))
            }
            Type::Union(types) => {
                let types_str = types
                    .iter()
                    .map(|t| t.fmt_renumbered(var_map, counter))
                    .collect::<Vec<_>>()
                    .join(" | ");
                format!("({})", types_str)
            }
            Type::Record(fields, rest) => {
                let fields_str = fields
                    .iter()
                    .map(|(k, v)| format!("{}: {}", k, v.fmt_renumbered(var_map, counter)))
                    .collect::<Vec<_>>()
                    .join(", ");
                match rest.as_ref() {
                    Type::RowEmpty => format!("{{{}}}", fields_str),
                    _ => {
                        let rest_str = rest.fmt_renumbered(var_map, counter);
                        if fields_str.is_empty() {
                            format!("{{| {}}}", rest_str)
                        } else {
                            format!("{{{} | {}}}", fields_str, rest_str)
                        }
                    }
                }
            }
            Type::RowEmpty => "{}".to_string(),
            Type::Never => "never".to_string(),
            Type::Var(id) => {
                let index = *var_map.entry(*id).or_insert_with(|| {
                    let i = *counter;
                    *counter += 1;
                    i
                });
                format_var_name(index)
            }
        }
    }
}

/// Converts a TypeVarId to a readable name like `'a`, `'b`, ..., `'z`, `'a1`, `'b1`, etc.
///
/// Uses the slotmap key's index (via `KeyData`) for a reliable numeric index,
/// then maps it to a human-readable alphabetic name.
fn type_var_name(id: TypeVarId) -> String {
    use slotmap::Key;
    let index = id.data().as_ffi() as u32 as usize;
    format_var_name(index)
}

/// Formats a type variable name from a sequential index.
///
/// Maps index 0 → `'a`, 1 → `'b`, ..., 25 → `'z`, 26 → `'a1`, etc.
pub fn format_var_name(index: usize) -> String {
    let letter = (b'a' + (index % 26) as u8) as char;
    let suffix = index / 26;
    if suffix == 0 {
        format!("'{}", letter)
    } else {
        format!("'{}{}", letter, suffix)
    }
}

/// Formats a list of types as a comma-separated string using renumbered display.
///
/// Convenience helper used in error messages and overload reporting.
pub(crate) fn format_type_list(types: &[Type]) -> String {
    types
        .iter()
        .map(|t| t.display_renumbered())
        .collect::<Vec<_>>()
        .join(", ")
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Use display_renumbered() to produce clean, sequential type variable names
        // regardless of internal slotmap indices.
        write!(f, "{}", self.display_renumbered())
    }
}

/// Type scheme for polymorphic types (generalized types)
///
/// A type scheme represents a polymorphic type by quantifying over type variables.
/// For example: forall a b. (a -> b) -> [a] -> [b]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TypeScheme {
    /// Quantified type variables
    pub quantified: Vec<TypeVarId>,
    /// The actual type
    pub ty: Type,
}

impl TypeScheme {
    /// Creates a monomorphic type scheme (no quantified variables)
    pub fn mono(ty: Type) -> Self {
        Self {
            quantified: Vec::new(),
            ty,
        }
    }

    /// Creates a polymorphic type scheme
    pub fn poly(quantified: Vec<TypeVarId>, ty: Type) -> Self {
        Self { quantified, ty }
    }

    /// Instantiates this type scheme with fresh type variables
    pub fn instantiate(&self, ctx: &mut TypeVarContext) -> Type {
        if self.quantified.is_empty() {
            return self.ty.clone();
        }

        // Create fresh type variables for each quantified variable
        let mut subst = Substitution::empty();
        for var_id in &self.quantified {
            let fresh = ctx.fresh();
            subst.insert(*var_id, Type::Var(fresh));
        }

        self.ty.apply_subst(&subst)
    }

    /// Generalizes a type into a type scheme
    pub fn generalize(ty: Type, env_vars: &[TypeVarId]) -> Self {
        let ty_vars = ty.free_vars();
        let quantified: Vec<TypeVarId> = ty_vars.into_iter().filter(|v| !env_vars.contains(v)).collect();
        Self::poly(quantified, ty)
    }
}

impl fmt::Display for TypeScheme {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.quantified.is_empty() {
            // Monomorphic type — renumber for clean display
            write!(f, "{}", self.ty.display_renumbered())
        } else {
            // Polymorphic type — renumber quantified vars to 'a, 'b, ...
            let mut var_map: FxHashMap<TypeVarId, usize> = FxHashMap::default();
            for (i, var) in self.quantified.iter().enumerate() {
                var_map.insert(*var, i);
            }
            let mut counter = self.quantified.len();

            write!(f, "forall ")?;
            for (i, var) in self.quantified.iter().enumerate() {
                if i > 0 {
                    write!(f, " ")?;
                }
                write!(f, "{}", format_var_name(var_map[var]))?;
            }
            write!(f, ". {}", self.ty.fmt_renumbered(&mut var_map, &mut counter))
        }
    }
}

/// Type variable context for generating fresh type variables
pub struct TypeVarContext {
    vars: SlotMap<TypeVarId, Option<Type>>,
}

impl TypeVarContext {
    /// Creates a new type variable context
    pub fn new() -> Self {
        Self {
            vars: SlotMap::with_key(),
        }
    }

    /// Generates a fresh type variable
    pub fn fresh(&mut self) -> TypeVarId {
        self.vars.insert(None)
    }

    /// Gets the resolved type for a type variable
    pub fn get(&self, var: TypeVarId) -> Option<&Type> {
        self.vars.get(var).and_then(|opt| opt.as_ref())
    }

    /// Sets the resolved type for a type variable
    pub fn set(&mut self, var: TypeVarId, ty: Type) {
        if let Some(slot) = self.vars.get_mut(var) {
            *slot = Some(ty);
        }
    }

    /// Checks if a type variable is resolved
    pub fn is_resolved(&self, var: TypeVarId) -> bool {
        self.vars.get(var).and_then(|opt| opt.as_ref()).is_some()
    }
}

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

/// Type substitution mapping type variables to types
#[derive(Debug, Clone, Default)]
pub struct Substitution {
    map: FxHashMap<TypeVarId, Type>,
}

impl Substitution {
    /// Creates an empty substitution
    pub fn empty() -> Self {
        Self {
            map: FxHashMap::default(),
        }
    }

    /// Inserts a substitution
    pub fn insert(&mut self, var: TypeVarId, ty: Type) {
        self.map.insert(var, ty);
    }

    /// Looks up a type variable in the substitution
    pub fn lookup(&self, var: TypeVarId) -> Option<&Type> {
        self.map.get(&var)
    }

    /// Returns true if the substitution has no bindings
    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    /// Composes two substitutions
    pub fn compose(&self, other: &Substitution) -> Substitution {
        let mut result = Substitution::empty();

        // Apply other to all types in self
        for (var, ty) in &self.map {
            result.insert(*var, ty.apply_subst(other));
        }

        // Add mappings from other that aren't in self
        for (var, ty) in &other.map {
            if !self.map.contains_key(var) {
                result.insert(*var, ty.clone());
            }
        }

        result
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rstest::rstest;

    #[test]
    fn test_type_display() {
        assert_eq!(Type::Number.to_string(), "number");
        assert_eq!(Type::String.to_string(), "string");
        assert_eq!(Type::array(Type::Number).to_string(), "[number]");
        assert_eq!(
            Type::function(vec![Type::Number, Type::String], Type::Bool).to_string(),
            "(number, string) -> bool"
        );
    }

    #[test]
    fn test_type_var_context() {
        let mut ctx = TypeVarContext::new();
        let var1 = ctx.fresh();
        let var2 = ctx.fresh();
        assert_ne!(var1, var2);
    }

    #[test]
    fn test_substitution() {
        let mut ctx = TypeVarContext::new();
        let var = ctx.fresh();
        let ty = Type::Var(var);

        let mut subst = Substitution::empty();
        subst.insert(var, Type::Number);

        let result = ty.apply_subst(&subst);
        assert_eq!(result, Type::Number);
    }

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

        let scheme = TypeScheme::poly(vec![var], Type::Var(var));
        let inst1 = scheme.instantiate(&mut ctx);
        let inst2 = scheme.instantiate(&mut ctx);

        // Each instantiation should create fresh variables
        assert_ne!(inst1, inst2);
    }

    #[test]
    fn test_can_match_concrete_types() {
        assert!(Type::Number.can_match(&Type::Number));
        assert!(Type::String.can_match(&Type::String));
        assert!(!Type::Number.can_match(&Type::String));
    }

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

        // Type variables can match anything
        assert!(Type::Var(var).can_match(&Type::Number));
        assert!(Type::Number.can_match(&Type::Var(var)));
        assert!(Type::Var(var).can_match(&Type::String));
    }

    #[test]
    fn test_can_match_arrays() {
        let arr_num = Type::array(Type::Number);
        let arr_str = Type::array(Type::String);

        assert!(arr_num.can_match(&arr_num));
        assert!(!arr_num.can_match(&arr_str));
    }

    #[test]
    fn test_can_match_functions() {
        let func1 = Type::function(vec![Type::Number], Type::String);
        let func2 = Type::function(vec![Type::Number], Type::String);
        let func3 = Type::function(vec![Type::String], Type::String);

        assert!(func1.can_match(&func2));
        assert!(!func1.can_match(&func3));
    }

    #[test]
    fn test_match_score() {
        // Exact matches get highest score
        assert_eq!(Type::Number.match_score(&Type::Number), Some(100));
        assert_eq!(Type::String.match_score(&Type::String), Some(100));

        // Type variables get lower score
        let mut ctx = TypeVarContext::new();
        let var = ctx.fresh();
        assert_eq!(Type::Var(var).match_score(&Type::Number), Some(10));

        // Incompatible types return None
        assert_eq!(Type::Number.match_score(&Type::String), None);
    }

    #[rstest]
    #[case(vec![Type::Number, Type::Number], Type::Number)]
    #[case(vec![Type::Number, Type::String], Type::union(vec![Type::Number, Type::String]))]
    #[case(vec![Type::union(vec![Type::Number, Type::String]), Type::Bool], Type::union(vec![Type::Number, Type::String, Type::Bool]))]
    #[case(vec![Type::Number, Type::String, Type::Number], Type::union(vec![Type::Number, Type::String]))]
    fn test_type_union(#[case] types: Vec<Type>, #[case] expected: Type) {
        assert_eq!(Type::union(types), expected);
    }

    #[rstest]
    #[case(Type::union(vec![Type::Number, Type::String]), &Type::Number, Type::String)]
    #[case(Type::union(vec![Type::Number, Type::String, Type::Bool]), &Type::String, Type::union(vec![Type::Number, Type::Bool]))]
    #[case(Type::Number, &Type::Number, Type::Number)]
    fn test_type_subtract(#[case] ty: Type, #[case] exclude: &Type, #[case] expected: Type) {
        assert_eq!(ty.subtract(exclude), expected);
    }

    #[rstest]
    #[case(Type::Number, Type::Number, true)]
    #[case(Type::Number, Type::String, false)]
    #[case(Type::array(Type::Number), Type::array(Type::Number), true)]
    #[case(Type::array(Type::Number), Type::array(Type::String), false)]
    #[case(Type::tuple(vec![Type::Number]), Type::array(Type::Number), true)]
    #[case(Type::record(BTreeMap::from([("a".to_string(), Type::Number)]), Type::RowEmpty), Type::dict(Type::String, Type::Number), true)]
    fn test_can_match_complex(#[case] t1: Type, #[case] t2: Type, #[case] expected: bool) {
        assert_eq!(t1.can_match(&t2), expected);
    }

    #[rstest]
    #[case(Type::Number, Type::Number, true)]
    #[case(Type::Number, Type::String, false)]
    fn test_can_branch_unify_with(#[case] t1: Type, #[case] t2: Type, #[case] expected: bool) {
        assert_eq!(t1.can_branch_unify_with(&t2), expected);
    }

    #[test]
    fn test_can_branch_unify_with_vars() {
        let mut ctx = TypeVarContext::new();
        let v1 = ctx.fresh();
        let v2 = ctx.fresh();
        assert!(Type::Var(v1).can_branch_unify_with(&Type::Var(v2)));
        assert!(!Type::Var(v1).can_branch_unify_with(&Type::Number));
    }

    #[test]
    fn test_substitution_compose() {
        let mut ctx = TypeVarContext::new();
        let v1 = ctx.fresh();
        let v2 = ctx.fresh();

        let mut s1 = Substitution::empty();
        s1.insert(v1, Type::Var(v2));

        let mut s2 = Substitution::empty();
        s2.insert(v2, Type::Number);

        let s3 = s1.compose(&s2);
        assert_eq!(s3.lookup(v1), Some(&Type::Number));
        assert_eq!(s3.lookup(v2), Some(&Type::Number));
    }

    #[test]
    fn test_type_scheme_generalize() {
        let mut ctx = TypeVarContext::new();
        let v1 = ctx.fresh();
        let v2 = ctx.fresh();

        let ty = Type::function(vec![Type::Var(v1)], Type::Var(v2));
        let env_vars = vec![v1];

        let scheme = TypeScheme::generalize(ty, &env_vars);
        assert_eq!(scheme.quantified, vec![v2]);
    }

    #[test]
    fn test_display_renumbered() {
        let mut ctx = TypeVarContext::new();
        let v1 = ctx.fresh();
        let v2 = ctx.fresh();
        let ty = Type::function(vec![Type::Var(v1)], Type::Var(v2));

        // Renumbered should always start from 'a
        assert_eq!(ty.display_renumbered(), "('a) -> 'b");
    }
}