//! Unification algorithm for type inference

use crate::middleend::types::{MonoType, Substitution, TyVar};
use anyhow::{bail, Result};
use std::collections::HashMap;

/// Unification engine for type inference
pub struct Unifier {
    subst: Substitution,
}

impl Unifier {
    #[must_use]
    pub fn new() -> Self {
        Unifier {
            subst: HashMap::new(),
        }
    }

    /// Get the current substitution
    #[must_use]
    pub fn substitution(&self) -> &Substitution {
        &self.subst
    }

    /// Apply current substitution to a type
    #[must_use]
    pub fn apply(&self, ty: &MonoType) -> MonoType {
        ty.substitute(&self.subst)
    }

    /// Unify two types, updating the substitution
    ///
    /// # Errors
    ///
    /// Returns an error if the types cannot be unified (type mismatch or occurs check failure)
    /// # Errors
    ///
    /// Returns an error if the operation fails
    pub fn unify(&mut self, t1: &MonoType, t2: &MonoType) -> Result<()> {
        let t1 = self.apply(t1);
        let t2 = self.apply(t2);

        match (t1, t2) {
            (MonoType::Var(v1), MonoType::Var(v2)) if v1 == v2 => Ok(()),
            (MonoType::Var(v), t) | (t, MonoType::Var(v)) => self.bind(&v, &t),
            (MonoType::Int, MonoType::Int)
            | (MonoType::Float, MonoType::Float)
            | (MonoType::Bool, MonoType::Bool)
            | (MonoType::String, MonoType::String)
            | (MonoType::Unit, MonoType::Unit) => Ok(()),
            (MonoType::Named(n1), MonoType::Named(n2)) if n1 == n2 => Ok(()),
            (MonoType::Function(a1, r1), MonoType::Function(a2, r2)) => {
                self.unify(&a1, &a2)?;
                self.unify(&r1, &r2)
            }
            (MonoType::List(e1), MonoType::List(e2)) => self.unify(&e1, &e2),
            (MonoType::Optional(i1), MonoType::Optional(i2)) => self.unify(&i1, &i2),
            (MonoType::Result(ok1, err1), MonoType::Result(ok2, err2)) => {
                self.unify(&ok1, &ok2)?;
                self.unify(&err1, &err2)
            }
            (MonoType::DataFrame(cols1), MonoType::DataFrame(cols2)) => {
                // DataFrames unify if they have the same columns with the same types
                if cols1.len() != cols2.len() {
                    bail!("Cannot unify DataFrames with different column counts");
                }
                for ((name1, ty1), (name2, ty2)) in cols1.iter().zip(cols2.iter()) {
                    if name1 != name2 {
                        bail!(
                            "Cannot unify DataFrames with different column names: {} vs {}",
                            name1,
                            name2
                        );
                    }
                    self.unify(ty1, ty2)?;
                }
                Ok(())
            }
            (MonoType::Series(ty1), MonoType::Series(ty2)) => self.unify(&ty1, &ty2),
            (t1, t2) => bail!("Cannot unify {} with {}", t1, t2),
        }
    }

    /// Bind a type variable to a type
    fn bind(&mut self, var: &TyVar, ty: &MonoType) -> Result<()> {
        // Occurs check: ensure var doesn't occur in ty
        if Self::occurs(var, ty) {
            bail!("Infinite type: {} occurs in {}", var, ty);
        }

        // Apply the binding
        self.subst.insert(var.clone(), ty.clone());

        // Update existing substitutions
        let updated: Substitution = self
            .subst
            .iter()
            .map(|(k, v)| {
                if k == var {
                    (k.clone(), ty.clone())
                } else {
                    (k.clone(), v.substitute(&[(var.clone(), ty.clone())].into()))
                }
            })
            .collect();
        self.subst = updated;

        Ok(())
    }

    /// Check if a type variable occurs in a type (occurs check)
    fn occurs(var: &TyVar, ty: &MonoType) -> bool {
        match ty {
            MonoType::Var(v) => v == var,
            MonoType::Function(arg, ret) => Self::occurs(var, arg) || Self::occurs(var, ret),
            MonoType::List(elem) => Self::occurs(var, elem),
            MonoType::Optional(inner) | MonoType::Series(inner) | MonoType::Reference(inner) => {
                Self::occurs(var, inner)
            }
            MonoType::Result(ok, err) => Self::occurs(var, ok) || Self::occurs(var, err),
            MonoType::DataFrame(columns) => {
                columns.iter().any(|(_, col_ty)| Self::occurs(var, col_ty))
            }
            MonoType::Tuple(types) => types.iter().any(|ty| Self::occurs(var, ty)),
            _ => false,
        }
    }

    /// Solve a type variable to its final type
    #[must_use]
    pub fn solve(&self, var: &TyVar) -> MonoType {
        self.subst
            .get(var)
            .map_or_else(|| MonoType::Var(var.clone()), |ty| self.apply(ty))
    }
}

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

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::panic)]
mod tests {
    use super::*;

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

    #[test]
    fn test_unify_different_types() {
        let mut unifier = Unifier::new();
        assert!(unifier.unify(&MonoType::Int, &MonoType::Bool).is_err());
        assert!(unifier.unify(&MonoType::String, &MonoType::Int).is_err());
    }

    #[test]
    fn test_unify_with_var() {
        let mut unifier = Unifier::new();
        let var = TyVar(0);

        assert!(unifier
            .unify(&MonoType::Var(var.clone()), &MonoType::Int)
            .is_ok());
        assert_eq!(unifier.solve(&var), MonoType::Int);
    }

    #[test]
    fn test_unify_functions() {
        let mut unifier = Unifier::new();
        let var = TyVar(0);

        let f1 = MonoType::Function(
            Box::new(MonoType::Int),
            Box::new(MonoType::Var(var.clone())),
        );
        let f2 = MonoType::Function(Box::new(MonoType::Int), Box::new(MonoType::Bool));

        assert!(unifier.unify(&f1, &f2).is_ok());
        assert_eq!(unifier.solve(&var), MonoType::Bool);
    }

    #[test]
    fn test_unify_lists() {
        let mut unifier = Unifier::new();
        let var = TyVar(0);

        let l1 = MonoType::List(Box::new(MonoType::Var(var.clone())));
        let l2 = MonoType::List(Box::new(MonoType::String));

        assert!(unifier.unify(&l1, &l2).is_ok());
        assert_eq!(unifier.solve(&var), MonoType::String);
    }

    #[test]
    fn test_occurs_check() {
        let mut unifier = Unifier::new();
        let var = TyVar(0);

        // Try to unify τ0 with [τ0] - should fail (infinite type)
        let infinite = MonoType::List(Box::new(MonoType::Var(var.clone())));
        assert!(unifier.unify(&MonoType::Var(var), &infinite).is_err());
    }

    #[test]
    fn test_transitive_unification() {
        let mut unifier = Unifier::new();
        let var1 = TyVar(0);
        let var2 = TyVar(1);

        // τ0 = τ1
        assert!(unifier
            .unify(&MonoType::Var(var1.clone()), &MonoType::Var(var2.clone()))
            .is_ok());

        // τ1 = Int
        assert!(unifier
            .unify(&MonoType::Var(var2.clone()), &MonoType::Int)
            .is_ok());

        // Now τ0 should also be Int
        assert_eq!(unifier.solve(&var1), MonoType::Int);
        assert_eq!(unifier.solve(&var2), MonoType::Int);
    }
}
