// ============================================================================
// PFCL Type Checker — Unification
// ============================================================================
//
// Robinson's unification algorithm extended for:
//   - Row polymorphism (record types with row variables)
//   - Multi-argument function types
//   - Tuple types
//
// Produces either a Substitution or a TypeError with context about what
// failed and why — this is how gap reports get generated.

use crate::types::*;
use crate::{String, ToString, Vec, vec, BTreeMap};
use core::fmt;

// ---------------------------------------------------------------------------
// Type Errors
// ---------------------------------------------------------------------------

/// A type error with enough context to generate a useful gap report.
#[derive(Debug, Clone)]
pub struct TypeError {
    pub kind: TypeErrorKind,
    /// Breadcrumb trail of context for where the error occurred.
    pub context: Vec<ErrorContext>,
}

#[derive(Debug, Clone)]
pub enum TypeErrorKind {
    /// Two types could not be unified.
    Mismatch {
        expected: Type,
        actual: Type,
    },
    /// A type variable occurs in the type it's being unified with (infinite type).
    OccursCheck {
        var: TypeVar,
        ty: Type,
    },
    /// Function called with wrong number of arguments.
    ArityMismatch {
        expected: usize,
        actual: usize,
        fn_type: Type,
    },
    /// Tried to apply a non-function type.
    NotAFunction {
        actual: Type,
    },
    /// A record field was expected but missing.
    MissingField {
        field: String,
        record_type: Type,
    },
    /// A catalog function's type signature could not be found.
    UnknownCatalogEntry {
        hash: String,
        name: String,
    },
    /// Tuple length mismatch.
    TupleLengthMismatch {
        expected: usize,
        actual: usize,
    },
    /// Match arms have incompatible types.
    IncompatibleMatchArms {
        arm_index: usize,
        expected: Type,
        actual: Type,
    },
    /// Constructor pattern doesn't match the scrutinee type.
    PatternTypeMismatch {
        pattern_desc: String,
        expected: Type,
    },
    /// A gap: the composition needs a function of this type but none exists.
    Gap {
        needed: Type,
        description: String,
    },
}

/// Contextual breadcrumb for where in the expression tree the error occurred.
#[derive(Debug, Clone)]
pub enum ErrorContext {
    InApplication { fn_name: String },
    InArgument { position: usize, fn_name: String },
    InLetBinding { name: String },
    InLetBody { name: String },
    InLambdaBody,
    InLambdaParam { name: String, position: usize },
    InMatchScrutinee,
    InMatchArm { index: usize },
    InFieldAccess { field: String },
    InPipeStage { position: usize },
    InListElement { index: usize },
    InRecordField { name: String },
    InTupleElement { index: usize },
    InConstructorArg { name: String, position: usize },
    InIfCondition,
    InIfBranch { is_then: bool },
}

impl TypeError {
    pub fn mismatch(expected: Type, actual: Type) -> Self {
        TypeError {
            kind: TypeErrorKind::Mismatch { expected, actual },
            context: vec![],
        }
    }

    pub fn occurs(var: TypeVar, ty: Type) -> Self {
        TypeError {
            kind: TypeErrorKind::OccursCheck { var, ty },
            context: vec![],
        }
    }

    pub fn arity(expected: usize, actual: usize, fn_type: Type) -> Self {
        TypeError {
            kind: TypeErrorKind::ArityMismatch {
                expected,
                actual,
                fn_type,
            },
            context: vec![],
        }
    }

    pub fn not_a_function(actual: Type) -> Self {
        TypeError {
            kind: TypeErrorKind::NotAFunction { actual },
            context: vec![],
        }
    }

    pub fn missing_field(field: String, record_type: Type) -> Self {
        TypeError {
            kind: TypeErrorKind::MissingField {
                field,
                record_type,
            },
            context: vec![],
        }
    }

    pub fn unknown_catalog(hash: String, name: String) -> Self {
        TypeError {
            kind: TypeErrorKind::UnknownCatalogEntry { hash, name },
            context: vec![],
        }
    }

    /// Add a context breadcrumb (pushed from inner to outer).
    pub fn with_context(mut self, ctx: ErrorContext) -> Self {
        self.context.push(ctx);
        self
    }
}

impl fmt::Display for TypeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Print context path (outer to inner)
        if !self.context.is_empty() {
            write!(f, "In ")?;
            for (i, ctx) in self.context.iter().rev().enumerate() {
                if i > 0 {
                    write!(f, " → ")?;
                }
                match ctx {
                    ErrorContext::InApplication { fn_name } => write!(f, "application of `{}`", fn_name)?,
                    ErrorContext::InArgument { position, fn_name } => {
                        write!(f, "argument {} of `{}`", position + 1, fn_name)?
                    }
                    ErrorContext::InLetBinding { name } => write!(f, "let binding `{}`", name)?,
                    ErrorContext::InLetBody { name } => write!(f, "body of let `{}`", name)?,
                    ErrorContext::InLambdaBody => write!(f, "lambda body")?,
                    ErrorContext::InLambdaParam { name, position } => {
                        write!(f, "lambda param `{}` (position {})", name, position)?
                    }
                    ErrorContext::InMatchScrutinee => write!(f, "match scrutinee")?,
                    ErrorContext::InMatchArm { index } => write!(f, "match arm {}", index + 1)?,
                    ErrorContext::InFieldAccess { field } => write!(f, "field `.{}`", field)?,
                    ErrorContext::InPipeStage { position } => {
                        write!(f, "pipe stage {}", position + 1)?
                    }
                    ErrorContext::InListElement { index } => {
                        write!(f, "list element {}", index)?
                    }
                    ErrorContext::InRecordField { name } => {
                        write!(f, "record field `{}`", name)?
                    }
                    ErrorContext::InTupleElement { index } => {
                        write!(f, "tuple element {}", index)?
                    }
                    ErrorContext::InConstructorArg { name, position } => {
                        write!(f, "argument {} of constructor `{}`", position + 1, name)?
                    }
                    ErrorContext::InIfCondition => write!(f, "if condition")?,
                    ErrorContext::InIfBranch { is_then } => {
                        write!(f, "{} branch", if *is_then { "then" } else { "else" })?
                    }
                }
            }
            writeln!(f, ":")?;
        }

        // Print the error itself
        match &self.kind {
            TypeErrorKind::Mismatch { expected, actual } => {
                write!(f, "Type mismatch: expected `{}`, found `{}`", expected, actual)
            }
            TypeErrorKind::OccursCheck { var, ty } => {
                write!(f, "Infinite type: `{}` occurs in `{}`", var, ty)
            }
            TypeErrorKind::ArityMismatch {
                expected,
                actual,
                fn_type,
            } => write!(
                f,
                "Arity mismatch: function of type `{}` expects {} argument(s), got {}",
                fn_type, expected, actual
            ),
            TypeErrorKind::NotAFunction { actual } => {
                write!(f, "Not a function: tried to apply `{}`", actual)
            }
            TypeErrorKind::MissingField { field, record_type } => {
                write!(
                    f,
                    "Missing field `{}` in record type `{}`",
                    field, record_type
                )
            }
            TypeErrorKind::UnknownCatalogEntry { hash, name } => {
                write!(
                    f,
                    "Unknown catalog entry: `{}` (hash: {})",
                    name, hash
                )
            }
            TypeErrorKind::TupleLengthMismatch { expected, actual } => {
                write!(
                    f,
                    "Tuple length mismatch: expected {}, found {}",
                    expected, actual
                )
            }
            TypeErrorKind::IncompatibleMatchArms {
                arm_index,
                expected,
                actual,
            } => write!(
                f,
                "Match arm {} has type `{}` but expected `{}` (from first arm)",
                arm_index + 1,
                actual,
                expected
            ),
            TypeErrorKind::PatternTypeMismatch {
                pattern_desc,
                expected,
            } => write!(
                f,
                "Pattern `{}` incompatible with scrutinee type `{}`",
                pattern_desc, expected
            ),
            TypeErrorKind::Gap { needed, description } => {
                write!(
                    f,
                    "GAP: Need a function of type `{}` — {}",
                    needed, description
                )
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Unification
// ---------------------------------------------------------------------------

/// Unify two types, producing a substitution or a type error.
pub fn unify(t1: &Type, t2: &Type) -> Result<Substitution, TypeError> {
    match (t1, t2) {
        // Two identical type variables
        (Type::Var(a), Type::Var(b)) if a == b => Ok(Substitution::empty()),

        // Variable on the left — bind it
        (Type::Var(v), t) => bind(*v, t),

        // Variable on the right — bind it
        (t, Type::Var(v)) => bind(*v, t),

        // Two type constructors — same name, unify args pointwise
        (Type::Con(n1, args1), Type::Con(n2, args2)) => {
            if n1 != n2 {
                return Err(TypeError::mismatch(t1.clone(), t2.clone()));
            }
            if args1.len() != args2.len() {
                return Err(TypeError::mismatch(t1.clone(), t2.clone()));
            }
            let mut subst = Substitution::empty();
            for (a1, a2) in args1.iter().zip(args2.iter()) {
                let a1_s = a1.apply_subst(&subst);
                let a2_s = a2.apply_subst(&subst);
                let s = unify(&a1_s, &a2_s)?;
                subst = s.compose(&subst);
            }
            Ok(subst)
        }

        // Two function types — unify inputs pointwise, then output
        (Type::Fun(ins1, out1), Type::Fun(ins2, out2)) => {
            if ins1.len() != ins2.len() {
                return Err(TypeError::arity(ins1.len(), ins2.len(), t1.clone()));
            }
            let mut subst = Substitution::empty();
            for (i1, i2) in ins1.iter().zip(ins2.iter()) {
                let i1_s = i1.apply_subst(&subst);
                let i2_s = i2.apply_subst(&subst);
                let s = unify(&i1_s, &i2_s)?;
                subst = s.compose(&subst);
            }
            let o1_s = out1.apply_subst(&subst);
            let o2_s = out2.apply_subst(&subst);
            let s = unify(&o1_s, &o2_s)?;
            Ok(s.compose(&subst))
        }

        // Two record types — unify shared fields, handle row variables
        (Type::Record(fields1, row1), Type::Record(fields2, row2)) => {
            unify_records(fields1, *row1, fields2, *row2)
        }

        // Two tuple types — unify elementwise
        (Type::Tuple(elems1), Type::Tuple(elems2)) => {
            if elems1.len() != elems2.len() {
                return Err(TypeError {
                    kind: TypeErrorKind::TupleLengthMismatch {
                        expected: elems1.len(),
                        actual: elems2.len(),
                    },
                    context: vec![],
                });
            }
            let mut subst = Substitution::empty();
            for (e1, e2) in elems1.iter().zip(elems2.iter()) {
                let e1_s = e1.apply_subst(&subst);
                let e2_s = e2.apply_subst(&subst);
                let s = unify(&e1_s, &e2_s)?;
                subst = s.compose(&subst);
            }
            Ok(subst)
        }

        // Everything else is a mismatch
        _ => Err(TypeError::mismatch(t1.clone(), t2.clone())),
    }
}

/// Bind a type variable to a type, with occurs check.
fn bind(v: TypeVar, t: &Type) -> Result<Substitution, TypeError> {
    // If t is already this variable, no binding needed
    if let Type::Var(w) = t {
        if *w == v {
            return Ok(Substitution::empty());
        }
    }
    // Occurs check: v must not appear in t
    if t.occurs(v) {
        return Err(TypeError::occurs(v, t.clone()));
    }
    Ok(Substitution::single(v, t.clone()))
}

/// Unify two record types with row polymorphism.
///
/// Strategy:
/// 1. Unify all shared fields pointwise.
/// 2. Fields only in one record but not the other:
///    - If the other has a row variable, the extra fields get absorbed into it.
///    - If no row variable, it's a type error (missing fields).
/// 3. Row variables: if both have row vars, unify them (they represent the
///    "rest" of each record).
fn unify_records(
    fields1: &BTreeMap<String, Type>,
    row1: Option<TypeVar>,
    fields2: &BTreeMap<String, Type>,
    row2: Option<TypeVar>,
) -> Result<Substitution, TypeError> {
    let keys1: HashSet<&String> = fields1.keys().collect();
    let keys2: HashSet<&String> = fields2.keys().collect();
    let shared: Vec<&&String> = keys1.intersection(&keys2).collect();
    let only1: Vec<&&String> = keys1.difference(&keys2).collect();
    let only2: Vec<&&String> = keys2.difference(&keys1).collect();

    // Unify shared fields
    let mut subst = Substitution::empty();
    for key in &shared {
        let t1 = fields1[**key].apply_subst(&subst);
        let t2 = fields2[**key].apply_subst(&subst);
        let s = unify(&t1, &t2)?;
        subst = s.compose(&subst);
    }

    // Handle fields only in record 1
    if !only1.is_empty() {
        match row2 {
            Some(r2) => {
                // The extra fields from record 1 must be part of record 2's row
                let extra: BTreeMap<String, Type> = only1
                    .iter()
                    .map(|k| ((**k).clone(), fields1[**k].apply_subst(&subst)))
                    .collect();
                // Bind r2 to a record containing the extra fields
                // (possibly with its own row variable if row1 exists)
                let tail_row = if row1.is_some() && row2.is_some() {
                    // Both have row vars — we'll handle below
                    row1
                } else {
                    None
                };
                let extra_record = Type::Record(extra, tail_row);
                let s = bind(r2, &extra_record.apply_subst(&subst))?;
                subst = s.compose(&subst);
            }
            None => {
                // No row variable to absorb extras — error
                let missing = only1.first().unwrap();
                return Err(TypeError::missing_field(
                    (***missing).clone(),
                    Type::Record(fields2.clone(), row2),
                ));
            }
        }
    }

    // Handle fields only in record 2
    if !only2.is_empty() {
        match row1 {
            Some(r1) => {
                let extra: BTreeMap<String, Type> = only2
                    .iter()
                    .map(|k| ((**k).clone(), fields2[**k].apply_subst(&subst)))
                    .collect();
                // Don't set tail_row again if we already handled it above
                let tail_row = if only1.is_empty() { row2 } else { None };
                let extra_record = Type::Record(extra, tail_row);
                let r1_applied = if let Some(t) = subst.lookup(r1) {
                    // r1 already bound — unify with extra
                    let t_clone = t.clone();
                    let s = unify(&t_clone, &extra_record.apply_subst(&subst))?;
                    subst = s.compose(&subst);
                    return Ok(subst);
                } else {
                    r1
                };
                let s = bind(r1_applied, &extra_record.apply_subst(&subst))?;
                subst = s.compose(&subst);
            }
            None => {
                let missing = only2.first().unwrap();
                return Err(TypeError::missing_field(
                    (***missing).clone(),
                    Type::Record(fields1.clone(), row1),
                ));
            }
        }
    }

    // If both had row variables and no extra fields on either side,
    // unify the row variables with each other
    if only1.is_empty() && only2.is_empty() {
        match (row1, row2) {
            (Some(r1), Some(r2)) if r1 != r2 => {
                let r1_t = Type::Var(r1).apply_subst(&subst);
                let r2_t = Type::Var(r2).apply_subst(&subst);
                let s = unify(&r1_t, &r2_t)?;
                subst = s.compose(&subst);
            }
            _ => {} // same variable, or one/both missing — nothing to do
        }
    }

    Ok(subst)
}

use hashbrown::HashSet;

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn var(n: u32) -> Type {
        Type::Var(TypeVar(n))
    }

    #[test]
    fn unify_same_primitive() {
        let s = unify(&Type::int(), &Type::int()).unwrap();
        assert!(s.is_empty());
    }

    #[test]
    fn unify_different_primitives() {
        let r = unify(&Type::int(), &Type::string());
        assert!(r.is_err());
    }

    #[test]
    fn unify_var_with_concrete() {
        let s = unify(&var(0), &Type::int()).unwrap();
        assert_eq!(s.lookup(TypeVar(0)), Some(&Type::int()));
    }

    #[test]
    fn unify_var_with_var() {
        let s = unify(&var(0), &var(1)).unwrap();
        // One should be bound to the other
        assert!(s.lookup(TypeVar(0)).is_some() || s.lookup(TypeVar(1)).is_some());
    }

    #[test]
    fn unify_list_types() {
        let t1 = Type::list(var(0));
        let t2 = Type::list(Type::int());
        let s = unify(&t1, &t2).unwrap();
        assert_eq!(var(0).apply_subst(&s), Type::int());
    }

    #[test]
    fn unify_function_types() {
        let t1 = Type::fun(vec![var(0)], Type::boolean());
        let t2 = Type::fun(vec![Type::int()], Type::boolean());
        let s = unify(&t1, &t2).unwrap();
        assert_eq!(var(0).apply_subst(&s), Type::int());
    }

    #[test]
    fn unify_function_arity_mismatch() {
        let t1 = Type::fun(vec![Type::int()], Type::boolean());
        let t2 = Type::fun(vec![Type::int(), Type::int()], Type::boolean());
        assert!(unify(&t1, &t2).is_err());
    }

    #[test]
    fn unify_nested_generic() {
        // List<Maybe<A>> ~ List<Maybe<Int>>
        let t1 = Type::list(Type::maybe(var(0)));
        let t2 = Type::list(Type::maybe(Type::int()));
        let s = unify(&t1, &t2).unwrap();
        assert_eq!(var(0).apply_subst(&s), Type::int());
    }

    #[test]
    fn occurs_check() {
        // A ~ List<A> should fail
        let t1 = var(0);
        let t2 = Type::list(var(0));
        let r = unify(&t1, &t2);
        assert!(r.is_err());
        match r.unwrap_err().kind {
            TypeErrorKind::OccursCheck { .. } => {}
            other => panic!("Expected OccursCheck, got {:?}", other),
        }
    }

    #[test]
    fn unify_tuples() {
        let t1 = Type::Tuple(vec![var(0), Type::string()]);
        let t2 = Type::Tuple(vec![Type::int(), var(1)]);
        let s = unify(&t1, &t2).unwrap();
        assert_eq!(var(0).apply_subst(&s), Type::int());
        assert_eq!(var(1).apply_subst(&s), Type::string());
    }

    #[test]
    fn unify_records_exact() {
        let mut f1 = BTreeMap::new();
        f1.insert("x".into(), Type::int());
        f1.insert("y".into(), Type::string());
        let mut f2 = BTreeMap::new();
        f2.insert("x".into(), Type::int());
        f2.insert("y".into(), Type::string());

        let t1 = Type::Record(f1, None);
        let t2 = Type::Record(f2, None);
        let s = unify(&t1, &t2).unwrap();
        assert!(s.is_empty());
    }

    #[test]
    fn unify_records_row_polymorphism() {
        // { x: Int | r } ~ { x: Int, y: String }
        let mut f1 = BTreeMap::new();
        f1.insert("x".into(), Type::int());
        let row = TypeVar(99);
        let t1 = Type::Record(f1, Some(row));

        let mut f2 = BTreeMap::new();
        f2.insert("x".into(), Type::int());
        f2.insert("y".into(), Type::string());
        let t2 = Type::Record(f2, None);

        let s = unify(&t1, &t2).unwrap();
        // Row variable should be bound to { y: String }
        let row_val = s.lookup(row).expect("row var should be bound");
        match row_val {
            Type::Record(fields, None) => {
                assert_eq!(fields.len(), 1);
                assert_eq!(fields["y"], Type::string());
            }
            other => panic!("Expected record, got {:?}", other),
        }
    }

    #[test]
    fn unify_records_missing_field_no_row() {
        // { x: Int } ~ { x: Int, y: String } with no row var → error
        let mut f1 = BTreeMap::new();
        f1.insert("x".into(), Type::int());
        let t1 = Type::Record(f1, None);

        let mut f2 = BTreeMap::new();
        f2.insert("x".into(), Type::int());
        f2.insert("y".into(), Type::string());
        let t2 = Type::Record(f2, None);

        assert!(unify(&t1, &t2).is_err());
    }

    #[test]
    fn substitution_compose() {
        // s1: a ↦ Int, s2: b ↦ a
        // compose(s1, s2) should give: b ↦ Int, a ↦ Int
        let s1 = Substitution::single(TypeVar(0), Type::int());
        let s2 = Substitution::single(TypeVar(1), var(0));
        let composed = s1.compose(&s2);
        assert_eq!(composed.lookup(TypeVar(1)), Some(&Type::int()));
        assert_eq!(composed.lookup(TypeVar(0)), Some(&Type::int()));
    }

    #[test]
    fn type_scheme_instantiation_free_vars() {
        // ∀ a. a → List<a>
        let scheme = TypeScheme {
            vars: vec![TypeVar(0)],
            ty: Type::fun1(var(0), Type::list(var(0))),
        };
        // Free vars should be empty (a is quantified)
        assert!(scheme.free_vars().is_empty());
    }
}
