// ============================================================================
// PFCL Type Checker — Type Representation
// ============================================================================
//
// The type language for PFCL. Mirrors the catalog's type system:
//   - Primitive types (Int, Float, Boolean, String, Unit)
//   - Type variables (for polymorphism)
//   - Type constructors applied to arguments (List<A>, Maybe<A>, etc.)
//   - Function types (A → B), represented as multi-arg for convenience
//   - Record types with row polymorphism ({ name: String, age: Int | r })
//   - Tuple types
//
// Every type variable is a unique integer assigned by a fresh-name supply.
// Substitution maps variables to types. Unification produces substitutions.

use crate::{Box, String, ToString, Vec, vec, BTreeMap};
use hashbrown::{HashMap, HashSet};
use core::fmt;

// ---------------------------------------------------------------------------
// Type Variables
// ---------------------------------------------------------------------------

/// A type variable — just a unique integer.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct TypeVar(pub u32);

impl fmt::Display for TypeVar {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Render as lowercase letters: 0→a, 1→b, ..., 25→z, 26→a1, etc.
        let idx = self.0 as usize;
        if idx < 26 {
            write!(f, "{}", (b'a' + idx as u8) as char)
        } else {
            write!(f, "{}{}", (b'a' + (idx % 26) as u8) as char, idx / 26)
        }
    }
}

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Type {
    /// A type variable (unification variable or quantified variable).
    Var(TypeVar),

    /// A primitive/named type constructor with zero or more arguments.
    /// Examples: `Int` → Con("Int", [])
    ///           `List<A>` → Con("List", [Var(a)])
    ///           `Maybe<Int>` → Con("Maybe", [Con("Int", [])])
    ///           `Either<A, B>` → Con("Either", [Var(a), Var(b)])
    ///           `Map<K, V>` → Con("Map", [Var(k), Var(v)])
    Con(String, Vec<Type>),

    /// Function type: inputs → output.
    /// `(Int, Int) → Bool` → Fun(vec![Int, Int], Bool)
    /// Single-arg: `A → B` → Fun(vec![A], B)
    Fun(Vec<Type>, Box<Type>),

    /// Record type with named fields and an optional row variable for
    /// extensibility (row polymorphism).
    /// `{ name: String, age: Int }` → Record(fields, None)
    /// `{ name: String | r }` → Record(fields, Some(TypeVar))
    Record(BTreeMap<String, Type>, Option<TypeVar>),

    /// Tuple type (positional product).
    /// `(Int, String)` → Tuple(vec![Int, String])
    Tuple(Vec<Type>),
}

// Convenience constructors
impl Type {
    pub fn int() -> Self {
        Type::Con("Int".into(), vec![])
    }
    pub fn float() -> Self {
        Type::Con("Float".into(), vec![])
    }
    pub fn boolean() -> Self {
        Type::Con("Boolean".into(), vec![])
    }
    pub fn string() -> Self {
        Type::Con("String".into(), vec![])
    }
    pub fn unit() -> Self {
        Type::Con("Unit".into(), vec![])
    }
    pub fn nothing() -> Self {
        // Nothing is a value of type Maybe<A> for any A — handled specially
        // during inference. Here we represent the *type* Maybe<A>.
        // The literal Nothing is typed as Maybe<fresh_var> at each use site.
        Type::Con("Maybe".into(), vec![Type::Var(TypeVar(u32::MAX))])
    }
    pub fn list(elem: Type) -> Self {
        Type::Con("List".into(), vec![elem])
    }
    pub fn maybe(inner: Type) -> Self {
        Type::Con("Maybe".into(), vec![inner])
    }
    pub fn either(a: Type, b: Type) -> Self {
        Type::Con("Either".into(), vec![a, b])
    }
    pub fn var(v: u32) -> Self {
        Type::Var(TypeVar(v))
    }
    pub fn fun1(input: Type, output: Type) -> Self {
        Type::Fun(vec![input], Box::new(output))
    }
    pub fn fun(inputs: Vec<Type>, output: Type) -> Self {
        Type::Fun(inputs, Box::new(output))
    }

    /// App type constructor: `App<R, S>` represents an application
    /// with resource type R and state type S.
    ///
    /// Semantically: `App<R, S> = (R, S, Event) -> (S, Commands)`
    pub fn app(res: Type, state: Type) -> Self {
        Type::Con("App".into(), vec![res, state])
    }

    /// The Event type used in app handler signatures.
    pub fn event() -> Self {
        Type::Con("Event".into(), vec![])
    }

    /// The Commands type returned by app handlers.
    pub fn commands() -> Self {
        Type::Con("Commands".into(), vec![])
    }

    /// The AppConfig type for runtime configuration.
    pub fn app_config() -> Self {
        Type::Con("AppConfig".into(), vec![])
    }

    /// Check if this type is an App<R, S> constructor.
    pub fn is_app(&self) -> bool {
        matches!(self, Type::Con(name, args) if name == "App" && args.len() == 2)
    }

    /// Extract (R, S) from an App<R, S> type, if it is one.
    pub fn as_app(&self) -> Option<(&Type, &Type)> {
        match self {
            Type::Con(name, args) if name == "App" && args.len() == 2 => {
                Some((&args[0], &args[1]))
            }
            _ => None,
        }
    }

    /// Expand `App<R, S>` into its handler function type:
    /// `(R, S, Event) -> (S, Commands)`
    pub fn app_handler_type(res: Type, state: Type) -> Self {
        Type::Fun(
            vec![res, state.clone(), Type::event()],
            Box::new(Type::Tuple(vec![state, Type::commands()])),
        )
    }

    /// Collect all free type variables in this type.
    pub fn free_vars(&self) -> HashSet<TypeVar> {
        let mut vars = HashSet::new();
        self.collect_free_vars(&mut vars);
        vars
    }

    fn collect_free_vars(&self, vars: &mut HashSet<TypeVar>) {
        match self {
            Type::Var(v) => {
                vars.insert(*v);
            }
            Type::Con(_, args) => {
                for a in args {
                    a.collect_free_vars(vars);
                }
            }
            Type::Fun(inputs, output) => {
                for i in inputs {
                    i.collect_free_vars(vars);
                }
                output.collect_free_vars(vars);
            }
            Type::Record(fields, row) => {
                for t in fields.values() {
                    t.collect_free_vars(vars);
                }
                if let Some(r) = row {
                    vars.insert(*r);
                }
            }
            Type::Tuple(elems) => {
                for e in elems {
                    e.collect_free_vars(vars);
                }
            }
        }
    }

    /// Apply a substitution to this type, producing a new type.
    pub fn apply_subst(&self, subst: &Substitution) -> Type {
        match self {
            Type::Var(v) => {
                if let Some(t) = subst.lookup(*v) {
                    // Apply recursively in case the substitution chains
                    t.apply_subst(subst)
                } else {
                    self.clone()
                }
            }
            Type::Con(name, args) => {
                Type::Con(name.clone(), args.iter().map(|a| a.apply_subst(subst)).collect())
            }
            Type::Fun(inputs, output) => Type::Fun(
                inputs.iter().map(|i| i.apply_subst(subst)).collect(),
                Box::new(output.apply_subst(subst)),
            ),
            Type::Record(fields, row) => {
                let new_fields: BTreeMap<String, Type> = fields
                    .iter()
                    .map(|(k, v)| (k.clone(), v.apply_subst(subst)))
                    .collect();
                let new_row = row.and_then(|r| {
                    if let Some(t) = subst.lookup(r) {
                        // If the row variable maps to a Record, merge fields
                        // Otherwise keep it as-is (will be caught by unification)
                        match t.apply_subst(subst) {
                            Type::Var(v) => Some(v),
                            _ => Some(r), // propagation handled in unification
                        }
                    } else {
                        Some(r)
                    }
                });
                Type::Record(new_fields, new_row)
            }
            Type::Tuple(elems) => {
                Type::Tuple(elems.iter().map(|e| e.apply_subst(subst)).collect())
            }
        }
    }

    /// Check if a type variable occurs in this type (for occurs check).
    pub fn occurs(&self, v: TypeVar) -> bool {
        match self {
            Type::Var(w) => *w == v,
            Type::Con(_, args) => args.iter().any(|a| a.occurs(v)),
            Type::Fun(inputs, output) => {
                inputs.iter().any(|i| i.occurs(v)) || output.occurs(v)
            }
            Type::Record(fields, row) => {
                fields.values().any(|t| t.occurs(v))
                    || row.map_or(false, |r| r == v)
            }
            Type::Tuple(elems) => elems.iter().any(|e| e.occurs(v)),
        }
    }
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Type::Var(v) => write!(f, "{}", v),
            Type::Con(name, args) if args.is_empty() => write!(f, "{}", name),
            Type::Con(name, args) => {
                write!(f, "{}<", name)?;
                for (i, a) in args.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", a)?;
                }
                write!(f, ">")
            }
            Type::Fun(inputs, output) => {
                if inputs.len() == 1 {
                    let inp = &inputs[0];
                    // Parenthesize function inputs
                    match inp {
                        Type::Fun(_, _) => write!(f, "({})", inp)?,
                        _ => write!(f, "{}", inp)?,
                    }
                } else {
                    write!(f, "(")?;
                    for (i, inp) in inputs.iter().enumerate() {
                        if i > 0 {
                            write!(f, ", ")?;
                        }
                        write!(f, "{}", inp)?;
                    }
                    write!(f, ")")?;
                }
                write!(f, " → {}", output)
            }
            Type::Record(fields, row) => {
                write!(f, "{{ ")?;
                for (i, (k, v)) in fields.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}: {}", k, v)?;
                }
                if let Some(r) = row {
                    if !fields.is_empty() {
                        write!(f, " ")?;
                    }
                    write!(f, "| {}", r)?;
                }
                write!(f, " }}")
            }
            Type::Tuple(elems) => {
                write!(f, "(")?;
                for (i, e) in elems.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", e)?;
                }
                write!(f, ")")
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Type Schemes (polymorphic types with ∀-quantified variables)
// ---------------------------------------------------------------------------

/// A type scheme: ∀ a b c. Type
/// The `vars` are the universally quantified type variables.
/// When instantiated, each gets a fresh unification variable.
#[derive(Debug, Clone)]
pub struct TypeScheme {
    pub vars: Vec<TypeVar>,
    pub ty: Type,
}

impl TypeScheme {
    /// A monomorphic type (no quantified variables).
    pub fn mono(ty: Type) -> Self {
        TypeScheme {
            vars: vec![],
            ty,
        }
    }

    /// Free variables = free vars of the type minus the quantified vars.
    pub fn free_vars(&self) -> HashSet<TypeVar> {
        let mut fv = self.ty.free_vars();
        for v in &self.vars {
            fv.remove(v);
        }
        fv
    }

    /// Apply a substitution (only affects free variables, not quantified ones).
    pub fn apply_subst(&self, subst: &Substitution) -> TypeScheme {
        // Filter out bindings for quantified variables
        let filtered = subst.without(&self.vars);
        TypeScheme {
            vars: self.vars.clone(),
            ty: self.ty.apply_subst(&filtered),
        }
    }
}

impl fmt::Display for TypeScheme {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.vars.is_empty() {
            write!(f, "{}", self.ty)
        } else {
            write!(f, "∀")?;
            for (i, v) in self.vars.iter().enumerate() {
                if i > 0 {
                    write!(f, " ")?;
                }
                write!(f, " {}", v)?;
            }
            write!(f, ". {}", self.ty)
        }
    }
}

// ---------------------------------------------------------------------------
// Substitution
// ---------------------------------------------------------------------------

/// A mapping from type variables to types.
/// This is the core output of unification.
#[derive(Debug, Clone)]
pub struct Substitution {
    map: HashMap<TypeVar, Type>,
}

impl Substitution {
    pub fn empty() -> Self {
        Substitution {
            map: HashMap::new(),
        }
    }

    pub fn single(v: TypeVar, t: Type) -> Self {
        let mut map = HashMap::new();
        map.insert(v, t);
        Substitution { map }
    }

    pub fn lookup(&self, v: TypeVar) -> Option<&Type> {
        self.map.get(&v)
    }

    pub fn insert(&mut self, v: TypeVar, t: Type) {
        self.map.insert(v, t);
    }

    /// Compose two substitutions: (self ∘ other).
    /// Apply self to all types in other, then add self's bindings.
    pub fn compose(&self, other: &Substitution) -> Substitution {
        let mut result: HashMap<TypeVar, Type> = other
            .map
            .iter()
            .map(|(v, t)| (*v, t.apply_subst(self)))
            .collect();
        // Add bindings from self that aren't in other
        for (v, t) in &self.map {
            result.entry(*v).or_insert_with(|| t.clone());
        }
        Substitution { map: result }
    }

    /// Return a substitution with certain variables removed.
    pub fn without(&self, vars: &[TypeVar]) -> Substitution {
        let exclude: HashSet<TypeVar> = vars.iter().cloned().collect();
        Substitution {
            map: self
                .map
                .iter()
                .filter(|(v, _)| !exclude.contains(*v))
                .map(|(v, t)| (*v, t.clone()))
                .collect(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    pub fn iter(&self) -> impl Iterator<Item = (&TypeVar, &Type)> {
        self.map.iter()
    }
}

impl fmt::Display for Substitution {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{{")?;
        for (i, (v, t)) in self.map.iter().enumerate() {
            if i > 0 {
                write!(f, ", ")?;
            }
            write!(f, "{} ↦ {}", v, t)?;
        }
        write!(f, "}}")
    }
}

// ---------------------------------------------------------------------------
// Fresh variable supply
// ---------------------------------------------------------------------------

/// Generates fresh type variables with a monotonically increasing counter.
#[derive(Debug)]
pub struct FreshSupply {
    next: u32,
}

impl FreshSupply {
    pub fn new() -> Self {
        FreshSupply { next: 0 }
    }

    pub fn fresh(&mut self) -> TypeVar {
        let v = TypeVar(self.next);
        self.next += 1;
        v
    }

    pub fn fresh_type(&mut self) -> Type {
        Type::Var(self.fresh())
    }
}
