//! Type system for the Lux language

#![allow(dead_code)]

use std::collections::{HashMap, HashSet};
use std::fmt;
use std::sync::atomic::{AtomicUsize, Ordering};

/// Unique ID for type variables
static NEXT_TYPE_VAR: AtomicUsize = AtomicUsize::new(0);

fn fresh_type_var() -> usize {
    NEXT_TYPE_VAR.fetch_add(1, Ordering::SeqCst)
}

/// Internal type representation
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Type {
    /// Type variable (for inference)
    Var(usize),
    /// Primitive types
    Int,
    Float,
    Bool,
    String,
    Char,
    Unit,
    /// Function type with effects and behavioral properties
    Function {
        params: Vec<Type>,
        return_type: Box<Type>,
        effects: EffectSet,
        properties: PropertySet,
    },
    /// Generic type application: List<Int>, Option<String>
    App {
        constructor: Box<Type>,
        args: Vec<Type>,
    },
    /// Named type (user-defined or built-in)
    Named(String),
    /// Tuple type
    Tuple(Vec<Type>),
    /// Record type
    Record(Vec<(String, Type)>),
    /// List type (sugar for App(List, [T]))
    List(Box<Type>),
    /// Option type (sugar for App(Option, [T]))
    Option(Box<Type>),
    /// Versioned type (e.g., User @v2)
    Versioned {
        base: Box<Type>,
        version: VersionInfo,
    },
    /// Error type (for type errors that shouldn't halt compilation)
    Error,
}

/// Version information for a versioned type
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum VersionInfo {
    /// Exactly this version
    Exact(u32),
    /// This version or later
    AtLeast(u32),
    /// Latest available version
    Latest,
}

impl Type {
    pub fn function(params: Vec<Type>, return_type: Type) -> Self {
        Type::Function {
            params,
            return_type: Box::new(return_type),
            effects: EffectSet::empty(),
            properties: PropertySet::empty(),
        }
    }

    pub fn function_with_effects(params: Vec<Type>, return_type: Type, effects: EffectSet) -> Self {
        Type::Function {
            params,
            return_type: Box::new(return_type),
            effects,
            properties: PropertySet::empty(),
        }
    }

    pub fn function_with_properties(
        params: Vec<Type>,
        return_type: Type,
        effects: EffectSet,
        properties: PropertySet,
    ) -> Self {
        Type::Function {
            params,
            return_type: Box::new(return_type),
            effects,
            properties,
        }
    }

    pub fn var() -> Self {
        Type::Var(fresh_type_var())
    }

    /// Check if this type contains a type variable
    pub fn contains_var(&self, var: usize) -> bool {
        match self {
            Type::Var(v) => *v == var,
            Type::Function {
                params,
                return_type,
                ..
            } => params.iter().any(|p| p.contains_var(var)) || return_type.contains_var(var),
            Type::App { constructor, args } => {
                constructor.contains_var(var) || args.iter().any(|a| a.contains_var(var))
            }
            Type::Tuple(elements) => elements.iter().any(|e| e.contains_var(var)),
            Type::Record(fields) => fields.iter().any(|(_, t)| t.contains_var(var)),
            Type::List(inner) | Type::Option(inner) => inner.contains_var(var),
            Type::Versioned { base, .. } => base.contains_var(var),
            _ => false,
        }
    }

    /// Apply a substitution to this type
    pub fn apply(&self, subst: &Substitution) -> Type {
        match self {
            Type::Var(v) => {
                if let Some(t) = subst.get(*v) {
                    t.apply(subst)
                } else {
                    self.clone()
                }
            }
            Type::Function {
                params,
                return_type,
                effects,
                properties,
            } => Type::Function {
                params: params.iter().map(|p| p.apply(subst)).collect(),
                return_type: Box::new(return_type.apply(subst)),
                effects: effects.clone(),
                properties: properties.clone(),
            },
            Type::App { constructor, args } => Type::App {
                constructor: Box::new(constructor.apply(subst)),
                args: args.iter().map(|a| a.apply(subst)).collect(),
            },
            Type::Tuple(elements) => Type::Tuple(elements.iter().map(|e| e.apply(subst)).collect()),
            Type::Record(fields) => Type::Record(
                fields
                    .iter()
                    .map(|(n, t)| (n.clone(), t.apply(subst)))
                    .collect(),
            ),
            Type::List(inner) => Type::List(Box::new(inner.apply(subst))),
            Type::Option(inner) => Type::Option(Box::new(inner.apply(subst))),
            Type::Versioned { base, version } => Type::Versioned {
                base: Box::new(base.apply(subst)),
                version: version.clone(),
            },
            _ => self.clone(),
        }
    }

    /// Get all free type variables in this type
    pub fn free_vars(&self) -> HashSet<usize> {
        match self {
            Type::Var(v) => {
                let mut set = HashSet::new();
                set.insert(*v);
                set
            }
            Type::Function {
                params,
                return_type,
                ..
            } => {
                let mut vars = HashSet::new();
                for p in params {
                    vars.extend(p.free_vars());
                }
                vars.extend(return_type.free_vars());
                vars
            }
            Type::App { constructor, args } => {
                let mut vars = constructor.free_vars();
                for a in args {
                    vars.extend(a.free_vars());
                }
                vars
            }
            Type::Tuple(elements) => {
                let mut vars = HashSet::new();
                for e in elements {
                    vars.extend(e.free_vars());
                }
                vars
            }
            Type::Record(fields) => {
                let mut vars = HashSet::new();
                for (_, t) in fields {
                    vars.extend(t.free_vars());
                }
                vars
            }
            Type::List(inner) | Type::Option(inner) => inner.free_vars(),
            Type::Versioned { base, .. } => base.free_vars(),
            _ => HashSet::new(),
        }
    }
}

impl fmt::Display for Type {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Type::Var(v) => write!(f, "?{}", v),
            Type::Int => write!(f, "Int"),
            Type::Float => write!(f, "Float"),
            Type::Bool => write!(f, "Bool"),
            Type::String => write!(f, "String"),
            Type::Char => write!(f, "Char"),
            Type::Unit => write!(f, "Unit"),
            Type::Function {
                params,
                return_type,
                effects,
                properties,
            } => {
                write!(f, "fn(")?;
                for (i, p) in params.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", p)?;
                }
                write!(f, "): {}", return_type)?;
                if !effects.is_empty() {
                    write!(f, " with {{{}}}", effects)?;
                }
                if !properties.is_empty() {
                    write!(f, " {}", properties)?;
                }
                Ok(())
            }
            Type::App { constructor, args } => {
                write!(f, "{}<", constructor)?;
                for (i, a) in args.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", a)?;
                }
                write!(f, ">")
            }
            Type::Named(name) => write!(f, "{}", name),
            Type::Tuple(elements) => {
                write!(f, "(")?;
                for (i, e) in elements.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}", e)?;
                }
                write!(f, ")")
            }
            Type::Record(fields) => {
                write!(f, "{{ ")?;
                for (i, (name, typ)) in fields.iter().enumerate() {
                    if i > 0 {
                        write!(f, ", ")?;
                    }
                    write!(f, "{}: {}", name, typ)?;
                }
                write!(f, " }}")
            }
            Type::List(inner) => write!(f, "List<{}>", inner),
            Type::Option(inner) => write!(f, "Option<{}>", inner),
            Type::Versioned { base, version } => {
                write!(f, "{} {}", base, version)
            }
            Type::Error => write!(f, "<error>"),
        }
    }
}

impl fmt::Display for VersionInfo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            VersionInfo::Exact(v) => write!(f, "@v{}", v),
            VersionInfo::AtLeast(v) => write!(f, "@v{}+", v),
            VersionInfo::Latest => write!(f, "@latest"),
        }
    }
}

/// A set of effects
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct EffectSet {
    pub effects: HashSet<String>,
}

impl EffectSet {
    pub fn empty() -> Self {
        Self {
            effects: HashSet::new(),
        }
    }

    pub fn single(effect: impl Into<String>) -> Self {
        let mut effects = HashSet::new();
        effects.insert(effect.into());
        Self { effects }
    }

    pub fn from_iter(iter: impl IntoIterator<Item = String>) -> Self {
        Self {
            effects: iter.into_iter().collect(),
        }
    }

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

    pub fn contains(&self, effect: &str) -> bool {
        self.effects.contains(effect)
    }

    pub fn union(&self, other: &EffectSet) -> EffectSet {
        EffectSet {
            effects: self.effects.union(&other.effects).cloned().collect(),
        }
    }

    pub fn is_subset(&self, other: &EffectSet) -> bool {
        self.effects.is_subset(&other.effects)
    }

    pub fn insert(&mut self, effect: impl Into<String>) {
        self.effects.insert(effect.into());
    }
}

impl fmt::Display for EffectSet {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let effects: Vec<_> = self.effects.iter().collect();
        for (i, e) in effects.iter().enumerate() {
            if i > 0 {
                write!(f, ", ")?;
            }
            write!(f, "{}", e)?;
        }
        Ok(())
    }
}

/// Set of behavioral properties for a function
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct PropertySet {
    pub properties: HashSet<Property>,
}

/// A behavioral property (mirrors BehavioralProperty from AST but for type system)
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Property {
    Pure,
    Total,
    Idempotent,
    Deterministic,
    Commutative,
}

impl From<crate::ast::BehavioralProperty> for Property {
    fn from(p: crate::ast::BehavioralProperty) -> Self {
        match p {
            crate::ast::BehavioralProperty::Pure => Property::Pure,
            crate::ast::BehavioralProperty::Total => Property::Total,
            crate::ast::BehavioralProperty::Idempotent => Property::Idempotent,
            crate::ast::BehavioralProperty::Deterministic => Property::Deterministic,
            crate::ast::BehavioralProperty::Commutative => Property::Commutative,
        }
    }
}

impl fmt::Display for Property {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Property::Pure => write!(f, "pure"),
            Property::Total => write!(f, "total"),
            Property::Idempotent => write!(f, "idempotent"),
            Property::Deterministic => write!(f, "deterministic"),
            Property::Commutative => write!(f, "commutative"),
        }
    }
}

impl PropertySet {
    pub fn empty() -> Self {
        Self {
            properties: HashSet::new(),
        }
    }

    pub fn from_ast(props: &[crate::ast::BehavioralProperty]) -> Self {
        Self {
            properties: props.iter().copied().map(Property::from).collect(),
        }
    }

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

    pub fn contains(&self, prop: Property) -> bool {
        self.properties.contains(&prop)
    }

    pub fn is_pure(&self) -> bool {
        self.contains(Property::Pure)
    }

    pub fn is_total(&self) -> bool {
        self.contains(Property::Total)
    }

    pub fn insert(&mut self, prop: Property) {
        self.properties.insert(prop);
    }

    /// Check if this property set is a superset of another (satisfies constraints)
    pub fn satisfies(&self, required: &PropertySet) -> bool {
        required.properties.is_subset(&self.properties)
    }

    /// Intersection of two property sets (for composition)
    pub fn intersection(&self, other: &PropertySet) -> PropertySet {
        PropertySet {
            properties: self
                .properties
                .intersection(&other.properties)
                .copied()
                .collect(),
        }
    }
}

impl fmt::Display for PropertySet {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let props: Vec<_> = self.properties.iter().collect();
        for (i, p) in props.iter().enumerate() {
            if i > 0 {
                write!(f, ", ")?;
            }
            write!(f, "is {}", p)?;
        }
        Ok(())
    }
}

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

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

    pub fn get(&self, var: usize) -> Option<&Type> {
        self.mapping.get(&var)
    }

    pub fn insert(&mut self, var: usize, typ: Type) {
        self.mapping.insert(var, typ);
    }

    pub fn compose(&self, other: &Substitution) -> Substitution {
        let mut result = Substitution::new();

        // Apply self to all mappings in other
        for (var, typ) in &other.mapping {
            result.insert(*var, typ.apply(self));
        }

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

        result
    }
}

/// Type scheme (polymorphic type)
#[derive(Debug, Clone)]
pub struct TypeScheme {
    pub vars: Vec<usize>,
    pub typ: Type,
}

impl TypeScheme {
    pub fn mono(typ: Type) -> Self {
        Self {
            vars: Vec::new(),
            typ,
        }
    }

    pub fn poly(vars: Vec<usize>, typ: Type) -> Self {
        Self { vars, typ }
    }

    /// Instantiate this scheme with fresh type variables
    pub fn instantiate(&self) -> Type {
        let mut subst = Substitution::new();
        for &var in &self.vars {
            subst.insert(var, Type::var());
        }
        self.typ.apply(&subst)
    }

    pub fn free_vars(&self) -> HashSet<usize> {
        let mut vars = self.typ.free_vars();
        for v in &self.vars {
            vars.remove(v);
        }
        vars
    }
}

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

/// Effect definition
#[derive(Debug, Clone)]
pub struct EffectDef {
    pub name: String,
    pub type_params: Vec<String>,
    pub operations: Vec<EffectOpDef>,
}

/// Effect operation definition
#[derive(Debug, Clone)]
pub struct EffectOpDef {
    pub name: String,
    pub params: Vec<(String, Type)>,
    pub return_type: Type,
}

/// Type environment
#[derive(Debug, Clone, Default)]
pub struct TypeEnv {
    /// Variable bindings
    pub bindings: HashMap<String, TypeScheme>,
    /// Type definitions
    pub types: HashMap<String, TypeDef>,
    /// Effect definitions
    pub effects: HashMap<String, EffectDef>,
    /// Handler types
    pub handlers: HashMap<String, HandlerDef>,
    /// Trait definitions
    pub traits: HashMap<String, TraitDef>,
    /// Trait implementations: (trait_name, type) -> impl
    pub trait_impls: Vec<TraitImpl>,
}

/// Type definition
#[derive(Debug, Clone)]
pub enum TypeDef {
    Alias(Type),
    Record(Vec<(String, Type)>),
    Enum(Vec<VariantDef>),
}

/// Variant definition
#[derive(Debug, Clone)]
pub struct VariantDef {
    pub name: String,
    pub fields: VariantFieldsDef,
}

/// Variant fields
#[derive(Debug, Clone)]
pub enum VariantFieldsDef {
    Unit,
    Tuple(Vec<Type>),
    Record(Vec<(String, Type)>),
}

/// Handler definition
#[derive(Debug, Clone)]
pub struct HandlerDef {
    pub name: String,
    pub effect: String,
    pub params: Vec<(String, Type)>,
}

/// Trait definition
#[derive(Debug, Clone)]
pub struct TraitDef {
    pub name: String,
    /// Type parameters for the trait (e.g., Functor<F>)
    pub type_params: Vec<String>,
    /// Super traits that must be implemented
    pub super_traits: Vec<TraitBoundDef>,
    /// Method signatures
    pub methods: Vec<TraitMethodDef>,
}

/// A trait bound in type definitions
#[derive(Debug, Clone)]
pub struct TraitBoundDef {
    pub trait_name: String,
    pub type_args: Vec<Type>,
}

/// A method signature in a trait
#[derive(Debug, Clone)]
pub struct TraitMethodDef {
    pub name: String,
    pub type_params: Vec<String>,
    pub params: Vec<(String, Type)>,
    pub return_type: Type,
    /// Whether this method has a default implementation
    pub has_default: bool,
}

/// A trait implementation
#[derive(Debug, Clone)]
pub struct TraitImpl {
    pub trait_name: String,
    pub trait_args: Vec<Type>,
    /// The type this impl is for
    pub target_type: Type,
    /// Type parameters on the impl (e.g., impl<T> Show for List<T>)
    pub type_params: Vec<String>,
    /// Constraints on type parameters (e.g., where T: Show)
    pub constraints: Vec<(String, Vec<TraitBoundDef>)>,
    /// Method implementations
    pub methods: HashMap<String, Type>,
}

/// A trait constraint on a type variable
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TraitConstraintDef {
    pub type_var: String,
    pub bounds: Vec<TraitBoundDef>,
}

impl PartialEq for TraitBoundDef {
    fn eq(&self, other: &Self) -> bool {
        self.trait_name == other.trait_name && self.type_args == other.type_args
    }
}

impl Eq for TraitBoundDef {}

impl TypeEnv {
    pub fn new() -> Self {
        Self::default()
    }

    /// Create environment with built-in types
    pub fn with_builtins() -> Self {
        let mut env = Self::new();

        // Add built-in types
        env.types
            .insert("Int".to_string(), TypeDef::Alias(Type::Int));
        env.types
            .insert("Float".to_string(), TypeDef::Alias(Type::Float));
        env.types
            .insert("Bool".to_string(), TypeDef::Alias(Type::Bool));
        env.types
            .insert("String".to_string(), TypeDef::Alias(Type::String));
        env.types
            .insert("Char".to_string(), TypeDef::Alias(Type::Char));
        env.types
            .insert("Unit".to_string(), TypeDef::Alias(Type::Unit));

        // Add Option type
        env.types.insert(
            "Option".to_string(),
            TypeDef::Enum(vec![
                VariantDef {
                    name: "None".to_string(),
                    fields: VariantFieldsDef::Unit,
                },
                VariantDef {
                    name: "Some".to_string(),
                    fields: VariantFieldsDef::Tuple(vec![Type::Var(0)]),
                },
            ]),
        );

        // Add Result type
        env.types.insert(
            "Result".to_string(),
            TypeDef::Enum(vec![
                VariantDef {
                    name: "Ok".to_string(),
                    fields: VariantFieldsDef::Tuple(vec![Type::Var(0)]),
                },
                VariantDef {
                    name: "Err".to_string(),
                    fields: VariantFieldsDef::Tuple(vec![Type::Var(1)]),
                },
            ]),
        );

        // Add Console effect
        env.effects.insert(
            "Console".to_string(),
            EffectDef {
                name: "Console".to_string(),
                type_params: Vec::new(),
                operations: vec![
                    EffectOpDef {
                        name: "print".to_string(),
                        params: vec![("msg".to_string(), Type::String)],
                        return_type: Type::Unit,
                    },
                    EffectOpDef {
                        name: "read".to_string(),
                        params: Vec::new(),
                        return_type: Type::String,
                    },
                ],
            },
        );

        // Add Fail effect
        env.effects.insert(
            "Fail".to_string(),
            EffectDef {
                name: "Fail".to_string(),
                type_params: Vec::new(),
                operations: vec![EffectOpDef {
                    name: "fail".to_string(),
                    params: vec![("msg".to_string(), Type::String)],
                    return_type: Type::Var(0), // Returns any type (never returns)
                }],
            },
        );

        // Add State effect
        env.effects.insert(
            "State".to_string(),
            EffectDef {
                name: "State".to_string(),
                type_params: vec!["S".to_string()],
                operations: vec![
                    EffectOpDef {
                        name: "get".to_string(),
                        params: Vec::new(),
                        return_type: Type::Var(0), // S
                    },
                    EffectOpDef {
                        name: "put".to_string(),
                        params: vec![("value".to_string(), Type::Var(0))], // S
                        return_type: Type::Unit,
                    },
                ],
            },
        );

        // Add Reader effect
        env.effects.insert(
            "Reader".to_string(),
            EffectDef {
                name: "Reader".to_string(),
                type_params: vec!["R".to_string()],
                operations: vec![EffectOpDef {
                    name: "ask".to_string(),
                    params: Vec::new(),
                    return_type: Type::Var(0), // R
                }],
            },
        );

        // Add Some and Ok, Err constructors
        // Some : fn(a) -> Option<a>
        let a = Type::var();
        env.bind(
            "Some",
            TypeScheme::mono(Type::function(vec![a.clone()], Type::Option(Box::new(a)))),
        );

        // None : Option<a>
        env.bind(
            "None",
            TypeScheme::mono(Type::Option(Box::new(Type::var()))),
        );

        // Ok : fn(a) -> Result<a, e>
        let a = Type::var();
        let e = Type::var();
        env.bind(
            "Ok",
            TypeScheme::mono(Type::function(
                vec![a.clone()],
                Type::App {
                    constructor: Box::new(Type::Named("Result".to_string())),
                    args: vec![a, e],
                },
            )),
        );

        // Err : fn(e) -> Result<a, e>
        let a = Type::var();
        let e = Type::var();
        env.bind(
            "Err",
            TypeScheme::mono(Type::function(
                vec![e.clone()],
                Type::App {
                    constructor: Box::new(Type::Named("Result".to_string())),
                    args: vec![a, e],
                },
            )),
        );

        // Add stdlib modules as record types
        // List module
        let list_module_type = Type::Record(vec![
            (
                "map".to_string(),
                Type::function(
                    vec![
                        Type::List(Box::new(Type::var())),
                        Type::function(vec![Type::var()], Type::var()),
                    ],
                    Type::List(Box::new(Type::var())),
                ),
            ),
            (
                "filter".to_string(),
                Type::function(
                    vec![
                        Type::List(Box::new(Type::var())),
                        Type::function(vec![Type::var()], Type::Bool),
                    ],
                    Type::List(Box::new(Type::var())),
                ),
            ),
            (
                "fold".to_string(),
                Type::function(
                    vec![
                        Type::List(Box::new(Type::var())),
                        Type::var(),
                        Type::function(vec![Type::var(), Type::var()], Type::var()),
                    ],
                    Type::var(),
                ),
            ),
            (
                "head".to_string(),
                Type::function(
                    vec![Type::List(Box::new(Type::var()))],
                    Type::Option(Box::new(Type::var())),
                ),
            ),
            (
                "tail".to_string(),
                Type::function(
                    vec![Type::List(Box::new(Type::var()))],
                    Type::Option(Box::new(Type::List(Box::new(Type::var())))),
                ),
            ),
            (
                "concat".to_string(),
                Type::function(
                    vec![
                        Type::List(Box::new(Type::var())),
                        Type::List(Box::new(Type::var())),
                    ],
                    Type::List(Box::new(Type::var())),
                ),
            ),
            (
                "reverse".to_string(),
                Type::function(
                    vec![Type::List(Box::new(Type::var()))],
                    Type::List(Box::new(Type::var())),
                ),
            ),
            (
                "length".to_string(),
                Type::function(vec![Type::List(Box::new(Type::var()))], Type::Int),
            ),
            (
                "get".to_string(),
                Type::function(
                    vec![Type::List(Box::new(Type::var())), Type::Int],
                    Type::Option(Box::new(Type::var())),
                ),
            ),
            (
                "range".to_string(),
                Type::function(vec![Type::Int, Type::Int], Type::List(Box::new(Type::Int))),
            ),
        ]);
        env.bind("List", TypeScheme::mono(list_module_type));

        // String module
        let string_module_type = Type::Record(vec![
            (
                "split".to_string(),
                Type::function(
                    vec![Type::String, Type::String],
                    Type::List(Box::new(Type::String)),
                ),
            ),
            (
                "join".to_string(),
                Type::function(
                    vec![Type::List(Box::new(Type::String)), Type::String],
                    Type::String,
                ),
            ),
            (
                "trim".to_string(),
                Type::function(vec![Type::String], Type::String),
            ),
            (
                "contains".to_string(),
                Type::function(vec![Type::String, Type::String], Type::Bool),
            ),
            (
                "replace".to_string(),
                Type::function(vec![Type::String, Type::String, Type::String], Type::String),
            ),
            (
                "length".to_string(),
                Type::function(vec![Type::String], Type::Int),
            ),
            (
                "chars".to_string(),
                Type::function(vec![Type::String], Type::List(Box::new(Type::Char))),
            ),
            (
                "lines".to_string(),
                Type::function(vec![Type::String], Type::List(Box::new(Type::String))),
            ),
        ]);
        env.bind("String", TypeScheme::mono(string_module_type));

        // Option module
        let option_module_type = Type::Record(vec![
            (
                "map".to_string(),
                Type::function(
                    vec![
                        Type::Option(Box::new(Type::var())),
                        Type::function(vec![Type::var()], Type::var()),
                    ],
                    Type::Option(Box::new(Type::var())),
                ),
            ),
            (
                "flatMap".to_string(),
                Type::function(
                    vec![
                        Type::Option(Box::new(Type::var())),
                        Type::function(vec![Type::var()], Type::Option(Box::new(Type::var()))),
                    ],
                    Type::Option(Box::new(Type::var())),
                ),
            ),
            (
                "getOrElse".to_string(),
                Type::function(
                    vec![Type::Option(Box::new(Type::var())), Type::var()],
                    Type::var(),
                ),
            ),
            (
                "isSome".to_string(),
                Type::function(vec![Type::Option(Box::new(Type::var()))], Type::Bool),
            ),
            (
                "isNone".to_string(),
                Type::function(vec![Type::Option(Box::new(Type::var()))], Type::Bool),
            ),
        ]);
        env.bind("Option", TypeScheme::mono(option_module_type));

        // Result module
        let result_type = Type::App {
            constructor: Box::new(Type::Named("Result".to_string())),
            args: vec![Type::var(), Type::var()],
        };
        let result_module_type = Type::Record(vec![
            (
                "map".to_string(),
                Type::function(
                    vec![
                        result_type.clone(),
                        Type::function(vec![Type::var()], Type::var()),
                    ],
                    result_type.clone(),
                ),
            ),
            (
                "flatMap".to_string(),
                Type::function(
                    vec![
                        result_type.clone(),
                        Type::function(vec![Type::var()], result_type.clone()),
                    ],
                    result_type.clone(),
                ),
            ),
            (
                "getOrElse".to_string(),
                Type::function(vec![result_type.clone(), Type::var()], Type::var()),
            ),
            (
                "isOk".to_string(),
                Type::function(vec![result_type.clone()], Type::Bool),
            ),
            (
                "isErr".to_string(),
                Type::function(vec![result_type], Type::Bool),
            ),
        ]);
        env.bind("Result", TypeScheme::mono(result_module_type));

        // Utility functions
        env.bind(
            "print",
            TypeScheme::mono(Type::function(vec![Type::var()], Type::Unit)),
        );
        env.bind(
            "toString",
            TypeScheme::mono(Type::function(vec![Type::var()], Type::String)),
        );
        env.bind(
            "typeOf",
            TypeScheme::mono(Type::function(vec![Type::var()], Type::String)),
        );

        env
    }

    pub fn lookup(&self, name: &str) -> Option<&TypeScheme> {
        self.bindings.get(name)
    }

    pub fn bind(&mut self, name: impl Into<String>, scheme: TypeScheme) {
        self.bindings.insert(name.into(), scheme);
    }

    pub fn extend(&self, name: impl Into<String>, scheme: TypeScheme) -> TypeEnv {
        let mut env = self.clone();
        env.bind(name, scheme);
        env
    }

    pub fn lookup_effect(&self, name: &str) -> Option<&EffectDef> {
        self.effects.get(name)
    }

    pub fn lookup_effect_op(&self, effect: &str, op: &str) -> Option<&EffectOpDef> {
        self.effects
            .get(effect)?
            .operations
            .iter()
            .find(|o| o.name == op)
    }

    /// Generalize a type to a type scheme
    pub fn generalize(&self, typ: &Type) -> TypeScheme {
        let env_vars: HashSet<usize> = self.bindings.values().flat_map(|s| s.free_vars()).collect();

        let type_vars: Vec<usize> = typ
            .free_vars()
            .into_iter()
            .filter(|v| !env_vars.contains(v))
            .collect();

        TypeScheme::poly(type_vars, typ.clone())
    }
}

/// Unification of types
pub fn unify(t1: &Type, t2: &Type) -> Result<Substitution, String> {
    match (t1, t2) {
        // Same type
        (Type::Int, Type::Int)
        | (Type::Float, Type::Float)
        | (Type::Bool, Type::Bool)
        | (Type::String, Type::String)
        | (Type::Char, Type::Char)
        | (Type::Unit, Type::Unit)
        | (Type::Error, _)
        | (_, Type::Error) => Ok(Substitution::new()),

        // Named types
        (Type::Named(a), Type::Named(b)) if a == b => Ok(Substitution::new()),

        // Type variables
        (Type::Var(v), t) | (t, Type::Var(v)) => {
            if let Type::Var(v2) = t {
                if v == v2 {
                    return Ok(Substitution::new());
                }
            }
            if t.contains_var(*v) {
                return Err(format!("Occurs check failed: ?{} occurs in {}", v, t));
            }
            let mut subst = Substitution::new();
            subst.insert(*v, t.clone());
            Ok(subst)
        }

        // Function types
        (
            Type::Function {
                params: p1,
                return_type: r1,
                effects: e1,
                properties: props1,
            },
            Type::Function {
                params: p2,
                return_type: r2,
                effects: e2,
                properties: props2,
            },
        ) => {
            if p1.len() != p2.len() {
                return Err(format!(
                    "Function arity mismatch: expected {} params, got {}",
                    p1.len(),
                    p2.len()
                ));
            }

            // Function's required effects (e1) must be a subset of available effects (e2)
            // A pure function (empty effects) can be called anywhere
            // A function requiring {Logger} can be called in context with {Logger} or {Logger, Console}
            if !e1.is_subset(&e2) {
                return Err(format!(
                    "Effect mismatch: expected {{{}}}, got {{{}}}",
                    e1, e2
                ));
            }

            // Properties are not checked during unification - they are guarantees, not constraints
            // A pure function can be used anywhere a function is expected
            // Property requirements are enforced via where clauses, not type unification
            let _ = (props1, props2); // Acknowledge but don't enforce

            let mut subst = Substitution::new();
            for (a, b) in p1.iter().zip(p2.iter()) {
                let s = unify(&a.apply(&subst), &b.apply(&subst))?;
                subst = subst.compose(&s);
            }
            let s = unify(&r1.apply(&subst), &r2.apply(&subst))?;
            Ok(subst.compose(&s))
        }

        // Type applications
        (
            Type::App {
                constructor: c1,
                args: a1,
            },
            Type::App {
                constructor: c2,
                args: a2,
            },
        ) => {
            if a1.len() != a2.len() {
                return Err(format!(
                    "Type argument count mismatch: {} vs {}",
                    a1.len(),
                    a2.len()
                ));
            }
            let mut subst = unify(c1, c2)?;
            for (a, b) in a1.iter().zip(a2.iter()) {
                let s = unify(&a.apply(&subst), &b.apply(&subst))?;
                subst = subst.compose(&s);
            }
            Ok(subst)
        }

        // Tuples
        (Type::Tuple(a), Type::Tuple(b)) => {
            if a.len() != b.len() {
                return Err(format!("Tuple size mismatch: {} vs {}", a.len(), b.len()));
            }
            let mut subst = Substitution::new();
            for (x, y) in a.iter().zip(b.iter()) {
                let s = unify(&x.apply(&subst), &y.apply(&subst))?;
                subst = subst.compose(&s);
            }
            Ok(subst)
        }

        // Records (structural)
        (Type::Record(a), Type::Record(b)) => {
            if a.len() != b.len() {
                return Err("Record field count mismatch".to_string());
            }
            let mut subst = Substitution::new();
            for ((n1, t1), (n2, t2)) in a.iter().zip(b.iter()) {
                if n1 != n2 {
                    return Err(format!("Record field name mismatch: {} vs {}", n1, n2));
                }
                let s = unify(&t1.apply(&subst), &t2.apply(&subst))?;
                subst = subst.compose(&s);
            }
            Ok(subst)
        }

        // List
        (Type::List(a), Type::List(b)) => unify(a, b),

        // Option
        (Type::Option(a), Type::Option(b)) => unify(a, b),

        // Versioned types
        (
            Type::Versioned {
                base: b1,
                version: v1,
            },
            Type::Versioned {
                base: b2,
                version: v2,
            },
        ) => {
            // Check version compatibility
            if !versions_compatible(v1, v2) {
                return Err(format!(
                    "Version mismatch: {} is not compatible with {}",
                    v1, v2
                ));
            }
            unify(b1, b2)
        }

        // Versioned type with non-versioned: treat as the base type
        (Type::Versioned { base, .. }, other) | (other, Type::Versioned { base, .. }) => {
            unify(base, other)
        }

        _ => Err(format!("Cannot unify {} with {}", t1, t2)),
    }
}

/// Check if two version constraints are compatible
fn versions_compatible(v1: &VersionInfo, v2: &VersionInfo) -> bool {
    match (v1, v2) {
        // Same exact version
        (VersionInfo::Exact(a), VersionInfo::Exact(b)) => a == b,

        // AtLeast is compatible if the other version is >= the minimum
        (VersionInfo::AtLeast(min), VersionInfo::Exact(v))
        | (VersionInfo::Exact(v), VersionInfo::AtLeast(min)) => v >= min,

        // Two AtLeast constraints: use the higher minimum
        (VersionInfo::AtLeast(_), VersionInfo::AtLeast(_)) => true, // compatible, use max

        // Latest is compatible with anything
        (VersionInfo::Latest, _) | (_, VersionInfo::Latest) => true,
    }
}