//! Type environment for tracking variable bindings during type inference.

use std::collections::HashSet;
use std::rc::Rc;

use indexmap::IndexMap;
use maat_ast::{NamedType, TypeExpr};

use crate::unify::Substitution;
use crate::{
    EnumDef, FnType, ImplDef, MethodSig, StructDef, TraitDef, Type, TypeScheme, TypeVarId,
    VariantDef, VariantKind, resolve_type_expr,
};

pub struct TypeEnv {
    bindings: Vec<(String, TypeScheme)>,
    scope_starts: Vec<usize>,
    next_var: TypeVarId,
    structs: IndexMap<String, StructDef>,
    enums: IndexMap<String, EnumDef>,
    traits: IndexMap<String, TraitDef>,
    impls: Vec<ImplDef>,
    builtin_method_schemes: IndexMap<String, BuiltinMethodScheme>,
}

#[derive(Debug, Clone)]
struct BuiltinMethodScheme {
    forall: Vec<TypeVarId>,
    self_type: Type,
    fn_type: Type,
}

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

impl TypeEnv {
    pub fn new() -> Self {
        Self {
            bindings: Vec::new(),
            scope_starts: vec![0],
            next_var: 0,
            structs: IndexMap::new(),
            enums: IndexMap::new(),
            traits: IndexMap::new(),
            impls: Vec::new(),
            builtin_method_schemes: IndexMap::new(),
        }
    }

    pub fn register_builtins(&mut self) {
        self.register_builtin_methods();
        self.register_builtin_enums();
        self.register_option_result_methods();
        self.register_builtin_ctors();
        self.register_numeric_conversions();
        self.register_default_fns();
        self.register_cmp_builtins();
    }

    fn register_builtin_ctors(&mut self) {
        // Vector::new() -> Vector<T> (i.e., [T])
        let vec_t_id = self.next_var;
        self.next_var += 1;
        let vec_ty = Type::Vector(Box::new(Type::Var(vec_t_id)));
        self.define_scheme(
            "Vector::new",
            TypeScheme {
                forall: vec![vec_t_id],
                ty: Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(vec_ty),
                }),
            },
        );

        // Set::new() -> Set<T>
        let set_t_id = self.next_var;
        self.next_var += 1;
        let set_ty = Type::Set(Box::new(Type::Var(set_t_id)));
        self.define_scheme(
            "Set::new",
            TypeScheme {
                forall: vec![set_t_id],
                ty: Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(set_ty),
                }),
            },
        );

        // Map::new() -> Map<K, V>
        let map_k_id = self.next_var;
        self.next_var += 1;
        let map_v_id = self.next_var;
        self.next_var += 1;
        let map_ty = Type::Map(Box::new(Type::Var(map_k_id)), Box::new(Type::Var(map_v_id)));
        self.define_scheme(
            "Map::new",
            TypeScheme {
                forall: vec![map_k_id, map_v_id],
                ty: Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(map_ty),
                }),
            },
        );
    }

    fn register_numeric_conversions(&mut self) {
        let targets: &[(&str, Type)] = &[
            ("i16::from", Type::I16),
            ("i32::from", Type::I32),
            ("i64::from", Type::I64),
            ("i128::from", Type::I128),
            ("u16::from", Type::U16),
            ("u32::from", Type::U32),
            ("u64::from", Type::U64),
            ("u128::from", Type::U128),
        ];
        for &(name, ref ret) in targets {
            let param_id = self.next_var;
            self.next_var += 1;
            self.define_scheme(
                name,
                TypeScheme {
                    forall: vec![param_id],
                    ty: Type::Function(FnType {
                        params: vec![Type::Var(param_id)],
                        ret: Box::new(ret.clone()),
                    }),
                },
            );
        }
    }

    fn register_default_fns(&mut self) {
        let defaults: &[(&str, Type)] = &[
            ("i8::default", Type::I8),
            ("i16::default", Type::I16),
            ("i32::default", Type::I32),
            ("i64::default", Type::I64),
            ("i128::default", Type::I128),
            ("u8::default", Type::U8),
            ("u16::default", Type::U16),
            ("u32::default", Type::U32),
            ("u64::default", Type::U64),
            ("u128::default", Type::U128),
            ("usize::default", Type::Usize),
            ("isize::default", Type::Isize),
            ("bool::default", Type::Bool),
            ("str::default", Type::Str),
        ];
        for &(name, ref ret) in defaults {
            self.define_scheme(
                name,
                TypeScheme {
                    forall: vec![],
                    ty: Type::Function(FnType {
                        params: vec![],
                        ret: Box::new(ret.clone()),
                    }),
                },
            );
        }
    }

    fn register_cmp_builtins(&mut self) {
        let t_id = self.next_var;
        self.next_var += 1;
        let t = Type::Var(t_id);
        let forall = vec![t_id];

        self.define_scheme(
            "cmp::min",
            TypeScheme {
                forall: forall.clone(),
                ty: Type::Function(FnType {
                    params: vec![t.clone(), t.clone()],
                    ret: Box::new(t.clone()),
                }),
            },
        );
        self.define_scheme(
            "cmp::max",
            TypeScheme {
                forall: forall.clone(),
                ty: Type::Function(FnType {
                    params: vec![t.clone(), t.clone()],
                    ret: Box::new(t.clone()),
                }),
            },
        );
        self.define_scheme(
            "cmp::clamp",
            TypeScheme {
                forall,
                ty: Type::Function(FnType {
                    params: vec![t.clone(), t.clone(), t.clone()],
                    ret: Box::new(t),
                }),
            },
        );
    }

    fn register_builtin_methods(&mut self) {
        let elem_id = self.next_var;
        self.next_var += 1;
        let elem = Type::Var(elem_id);
        let vector_ty = Type::Vector(Box::new(elem.clone()));
        let forall = vec![elem_id];
        // impl Vector<T>
        let vec_methods = [
            (
                "Vector::len",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Usize),
                }),
            ),
            (
                "Vector::first",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Enum(Rc::from("Option"), vec![elem.clone()])),
                }),
            ),
            (
                "Vector::last",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Enum(Rc::from("Option"), vec![elem.clone()])),
                }),
            ),
            (
                "Vector::split_first",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(vector_ty.clone()),
                }),
            ),
            (
                "Vector::push",
                Type::Function(FnType {
                    params: vec![elem.clone()],
                    ret: Box::new(vector_ty.clone()),
                }),
            ),
            (
                "Vector::join",
                Type::Function(FnType {
                    params: vec![Type::Str],
                    ret: Box::new(Type::Str),
                }),
            ),
            (
                "Vector::rev",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(vector_ty.clone()),
                }),
            ),
            (
                "Vector::count",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Usize),
                }),
            ),
            (
                "Vector::take",
                Type::Function(FnType {
                    params: vec![Type::Usize],
                    ret: Box::new(vector_ty.clone()),
                }),
            ),
            (
                "Vector::skip",
                Type::Function(FnType {
                    params: vec![Type::Usize],
                    ret: Box::new(vector_ty.clone()),
                }),
            ),
            (
                "Vector::dedup",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(vector_ty.clone()),
                }),
            ),
            (
                "Vector::chain",
                Type::Function(FnType {
                    params: vec![vector_ty.clone()],
                    ret: Box::new(vector_ty.clone()),
                }),
            ),
            (
                "Vector::contains",
                Type::Function(FnType {
                    params: vec![elem.clone()],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "Vector::sum",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(elem.clone()),
                }),
            ),
            (
                "Vector::product",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(elem.clone()),
                }),
            ),
            (
                "Vector::min",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Enum(Rc::from("Option"), vec![elem.clone()])),
                }),
            ),
            (
                "Vector::max",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Enum(Rc::from("Option"), vec![elem.clone()])),
                }),
            ),
        ];
        for (name, fn_type) in vec_methods {
            self.builtin_method_schemes.insert(
                name.to_string(),
                BuiltinMethodScheme {
                    forall: forall.clone(),
                    self_type: vector_ty.clone(),
                    fn_type,
                },
            );
        }
        // Higher-order Vector<T> methods requiring an additional type variable.
        let out_id = self.next_var;
        self.next_var += 1;
        let out = Type::Var(out_id);
        let vector_out = Type::Vector(Box::new(out.clone()));
        let vec_ho_methods: Vec<(&str, Vec<TypeVarId>, Type)> = vec![
            (
                "Vector::map",
                vec![elem_id, out_id],
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![Type::Var(elem_id)],
                        ret: Box::new(out.clone()),
                    })],
                    ret: Box::new(vector_out),
                }),
            ),
            (
                "Vector::filter",
                forall.clone(),
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![Type::Var(elem_id)],
                        ret: Box::new(Type::Bool),
                    })],
                    ret: Box::new(vector_ty.clone()),
                }),
            ),
            (
                "Vector::fold",
                vec![elem_id, out_id],
                Type::Function(FnType {
                    params: vec![
                        out.clone(),
                        Type::Function(FnType {
                            params: vec![out.clone(), Type::Var(elem_id)],
                            ret: Box::new(out.clone()),
                        }),
                    ],
                    ret: Box::new(Type::Var(out_id)),
                }),
            ),
            (
                "Vector::any",
                forall.clone(),
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![Type::Var(elem_id)],
                        ret: Box::new(Type::Bool),
                    })],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "Vector::all",
                forall.clone(),
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![Type::Var(elem_id)],
                        ret: Box::new(Type::Bool),
                    })],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "Vector::find",
                forall.clone(),
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![Type::Var(elem_id)],
                        ret: Box::new(Type::Bool),
                    })],
                    ret: Box::new(Type::Enum(Rc::from("Option"), vec![Type::Var(elem_id)])),
                }),
            ),
            (
                "Vector::position",
                forall.clone(),
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![Type::Var(elem_id)],
                        ret: Box::new(Type::Bool),
                    })],
                    ret: Box::new(Type::Enum(Rc::from("Option"), vec![Type::Usize])),
                }),
            ),
            (
                "Vector::for_each",
                forall.clone(),
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![Type::Var(elem_id)],
                        ret: Box::new(Type::Unit),
                    })],
                    ret: Box::new(Type::Unit),
                }),
            ),
            (
                "Vector::enumerate",
                forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Vector(Box::new(Type::Tuple(vec![
                        Type::Usize,
                        Type::Var(elem_id),
                    ])))),
                }),
            ),
            (
                "Vector::zip",
                vec![elem_id, out_id],
                Type::Function(FnType {
                    params: vec![Type::Vector(Box::new(out.clone()))],
                    ret: Box::new(Type::Vector(Box::new(Type::Tuple(vec![
                        Type::Var(elem_id),
                        out.clone(),
                    ])))),
                }),
            ),
            (
                "Vector::flat_map",
                vec![elem_id, out_id],
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![Type::Var(elem_id)],
                        ret: Box::new(Type::Vector(Box::new(out.clone()))),
                    })],
                    ret: Box::new(Type::Vector(Box::new(out.clone()))),
                }),
            ),
            (
                "Vector::windows",
                forall.clone(),
                Type::Function(FnType {
                    params: vec![Type::Usize],
                    ret: Box::new(Type::Vector(Box::new(vector_ty.clone()))),
                }),
            ),
            (
                "Vector::chunks",
                forall.clone(),
                Type::Function(FnType {
                    params: vec![Type::Usize],
                    ret: Box::new(Type::Vector(Box::new(vector_ty.clone()))),
                }),
            ),
        ];
        for (name, ho_forall, fn_type) in vec_ho_methods {
            self.builtin_method_schemes.insert(
                name.to_string(),
                BuiltinMethodScheme {
                    forall: ho_forall,
                    self_type: vector_ty.clone(),
                    fn_type,
                },
            );
        }
        // impl str
        let str_methods = [
            (
                "str::len",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Usize),
                }),
            ),
            (
                "str::trim",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Str),
                }),
            ),
            (
                "str::contains",
                Type::Function(FnType {
                    params: vec![Type::Str],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "str::starts_with",
                Type::Function(FnType {
                    params: vec![Type::Str],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "str::ends_with",
                Type::Function(FnType {
                    params: vec![Type::Str],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "str::split",
                Type::Function(FnType {
                    params: vec![Type::Str],
                    ret: Box::new(Type::Vector(Box::new(Type::Str))),
                }),
            ),
        ];
        let parse_error_ty = Type::Enum(Rc::from("ParseIntError"), vec![]);
        let parse_result = |int_ty: Type| -> Type {
            Type::Enum(Rc::from("Result"), vec![int_ty, parse_error_ty.clone()])
        };
        let parse_methods: [(&str, Type); 12] = [
            ("str::parse_int", parse_result(Type::I64)),
            ("str::parse_i8", parse_result(Type::I8)),
            ("str::parse_i16", parse_result(Type::I16)),
            ("str::parse_i32", parse_result(Type::I32)),
            ("str::parse_i64", parse_result(Type::I64)),
            ("str::parse_i128", parse_result(Type::I128)),
            ("str::parse_u8", parse_result(Type::U8)),
            ("str::parse_u16", parse_result(Type::U16)),
            ("str::parse_u32", parse_result(Type::U32)),
            ("str::parse_u64", parse_result(Type::U64)),
            ("str::parse_u128", parse_result(Type::U128)),
            ("str::parse_usize", parse_result(Type::Usize)),
        ];
        let all_str_methods = str_methods
            .into_iter()
            .chain(parse_methods.map(|(name, ret)| {
                (
                    name,
                    Type::Function(FnType {
                        params: vec![],
                        ret: Box::new(ret),
                    }),
                )
            }));
        for (name, fn_type) in all_str_methods {
            self.builtin_method_schemes.insert(
                name.to_string(),
                BuiltinMethodScheme {
                    forall: vec![],
                    self_type: Type::Str,
                    fn_type,
                },
            );
        }
        // impl char
        let char_methods: [(&str, Type); 3] = [
            (
                "char::is_alphabetic",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "char::is_numeric",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "char::to_string",
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Str),
                }),
            ),
        ];
        for (name, fn_type) in char_methods {
            self.builtin_method_schemes.insert(
                name.to_string(),
                BuiltinMethodScheme {
                    forall: vec![],
                    self_type: Type::Char,
                    fn_type,
                },
            );
        }
        // impl Set<T>
        let set_elem_id = self.next_var;
        self.next_var += 1;
        let set_elem = Type::Var(set_elem_id);
        let set_ty = Type::Set(Box::new(set_elem.clone()));
        let set_forall = vec![set_elem_id];
        let set_methods = [
            (
                "Set::insert",
                set_forall.clone(),
                Type::Function(FnType {
                    params: vec![set_elem.clone()],
                    ret: Box::new(set_ty.clone()),
                }),
            ),
            (
                "Set::contains",
                set_forall.clone(),
                Type::Function(FnType {
                    params: vec![set_elem.clone()],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "Set::remove",
                set_forall.clone(),
                Type::Function(FnType {
                    params: vec![set_elem.clone()],
                    ret: Box::new(set_ty.clone()),
                }),
            ),
            (
                "Set::len",
                set_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Usize),
                }),
            ),
            (
                "Set::to_vector",
                set_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Vector(Box::new(set_elem))),
                }),
            ),
        ];
        for (name, forall, fn_type) in set_methods {
            self.builtin_method_schemes.insert(
                name.to_string(),
                BuiltinMethodScheme {
                    forall,
                    self_type: set_ty.clone(),
                    fn_type,
                },
            );
        }
        // impl Map<K, V>
        let map_k_id = self.next_var;
        self.next_var += 1;
        let map_v_id = self.next_var;
        self.next_var += 1;
        let map_k = Type::Var(map_k_id);
        let map_v = Type::Var(map_v_id);
        let map_ty = Type::Map(Box::new(map_k.clone()), Box::new(map_v.clone()));
        let map_forall = vec![map_k_id, map_v_id];
        let map_methods = [
            (
                "Map::insert",
                map_forall.clone(),
                Type::Function(FnType {
                    params: vec![map_k.clone(), map_v.clone()],
                    ret: Box::new(map_ty.clone()),
                }),
            ),
            (
                "Map::get",
                map_forall.clone(),
                Type::Function(FnType {
                    params: vec![map_k.clone()],
                    ret: Box::new(Type::Enum(Rc::from("Option"), vec![map_v])),
                }),
            ),
            (
                "Map::contains_key",
                map_forall.clone(),
                Type::Function(FnType {
                    params: vec![map_k.clone()],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "Map::remove",
                map_forall.clone(),
                Type::Function(FnType {
                    params: vec![map_k.clone()],
                    ret: Box::new(map_ty.clone()),
                }),
            ),
            (
                "Map::len",
                vec![],
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Usize),
                }),
            ),
            (
                "Map::keys",
                map_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Vector(Box::new(map_k))),
                }),
            ),
            (
                "Map::values",
                map_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Vector(Box::new(Type::Var(map_v_id)))),
                }),
            ),
        ];
        for (name, forall, fn_type) in map_methods {
            self.builtin_method_schemes.insert(
                name.to_string(),
                BuiltinMethodScheme {
                    forall,
                    self_type: map_ty.clone(),
                    fn_type,
                },
            );
        }
    }

    fn register_option_result_methods(&mut self) {
        // impl Option<T>
        let t_id = self.next_var;
        self.next_var += 1;
        let t = Type::Var(t_id);
        let option_ty = Type::Enum(Rc::from("Option"), vec![t.clone()]);
        let option_forall = vec![t_id];

        let u_id = self.next_var;
        self.next_var += 1;
        let u = Type::Var(u_id);
        let option_u = Type::Enum(Rc::from("Option"), vec![u.clone()]);

        let option_methods: Vec<(&str, Vec<TypeVarId>, Type)> = vec![
            (
                "Option::unwrap",
                option_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(t.clone()),
                }),
            ),
            (
                "Option::unwrap_or",
                option_forall.clone(),
                Type::Function(FnType {
                    params: vec![t.clone()],
                    ret: Box::new(t.clone()),
                }),
            ),
            (
                "Option::is_some",
                option_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "Option::is_none",
                option_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "Option::map",
                vec![t_id, u_id],
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![t.clone()],
                        ret: Box::new(u.clone()),
                    })],
                    ret: Box::new(option_u.clone()),
                }),
            ),
            (
                "Option::and_then",
                vec![t_id, u_id],
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![t.clone()],
                        ret: Box::new(option_u.clone()),
                    })],
                    ret: Box::new(option_u.clone()),
                }),
            ),
            (
                "Option::unwrap_or_else",
                vec![t_id],
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![],
                        ret: Box::new(t.clone()),
                    })],
                    ret: Box::new(t.clone()),
                }),
            ),
            (
                "Option::ok",
                option_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Enum(Rc::from("Result"), vec![t.clone(), Type::Unit])),
                }),
            ),
            (
                "Option::zip",
                vec![t_id, u_id],
                Type::Function(FnType {
                    params: vec![Type::Enum(Rc::from("Option"), vec![u.clone()])],
                    ret: Box::new(Type::Enum(
                        Rc::from("Option"),
                        vec![Type::Tuple(vec![t.clone(), u.clone()])],
                    )),
                }),
            ),
        ];
        for (name, forall, fn_type) in option_methods {
            self.builtin_method_schemes.insert(
                name.to_string(),
                BuiltinMethodScheme {
                    forall,
                    self_type: option_ty.clone(),
                    fn_type,
                },
            );
        }

        let flatten_inner_id = self.next_var;
        self.next_var += 1;
        let flatten_inner = Type::Var(flatten_inner_id);
        let option_inner = Type::Enum(Rc::from("Option"), vec![flatten_inner.clone()]);
        self.builtin_method_schemes.insert(
            "Option::flatten".to_string(),
            BuiltinMethodScheme {
                forall: vec![flatten_inner_id],
                self_type: Type::Enum(Rc::from("Option"), vec![option_inner]),
                fn_type: Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Enum(Rc::from("Option"), vec![flatten_inner])),
                }),
            },
        );

        // impl Result<T, E>
        let rt_id = self.next_var;
        self.next_var += 1;
        let rt = Type::Var(rt_id);
        let re_id = self.next_var;
        self.next_var += 1;
        let re = Type::Var(re_id);
        let result_ty = Type::Enum(Rc::from("Result"), vec![rt.clone(), re.clone()]);
        let result_forall = vec![rt_id, re_id];

        let ru_id = self.next_var;
        self.next_var += 1;
        let ru = Type::Var(ru_id);
        let result_u = Type::Enum(Rc::from("Result"), vec![ru.clone(), re.clone()]);

        let rf_id = self.next_var;
        self.next_var += 1;
        let rf = Type::Var(rf_id);

        let result_methods: Vec<(&str, Vec<TypeVarId>, Type)> = vec![
            (
                "Result::unwrap",
                result_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(rt.clone()),
                }),
            ),
            (
                "Result::unwrap_or",
                result_forall.clone(),
                Type::Function(FnType {
                    params: vec![rt.clone()],
                    ret: Box::new(rt.clone()),
                }),
            ),
            (
                "Result::is_ok",
                result_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "Result::is_err",
                result_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Bool),
                }),
            ),
            (
                "Result::map",
                vec![rt_id, re_id, ru_id],
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![rt.clone()],
                        ret: Box::new(ru.clone()),
                    })],
                    ret: Box::new(result_u.clone()),
                }),
            ),
            (
                "Result::and_then",
                vec![rt_id, re_id, ru_id],
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![rt.clone()],
                        ret: Box::new(result_u.clone()),
                    })],
                    ret: Box::new(result_u),
                }),
            ),
            (
                "Result::map_err",
                vec![rt_id, re_id, rf_id],
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![re.clone()],
                        ret: Box::new(rf.clone()),
                    })],
                    ret: Box::new(Type::Enum(Rc::from("Result"), vec![rt.clone(), rf.clone()])),
                }),
            ),
            (
                "Result::unwrap_err",
                result_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(re.clone()),
                }),
            ),
            (
                "Result::unwrap_or_else",
                vec![rt_id, re_id],
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![re.clone()],
                        ret: Box::new(rt.clone()),
                    })],
                    ret: Box::new(rt.clone()),
                }),
            ),
            (
                "Result::ok",
                result_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Enum(Rc::from("Option"), vec![rt.clone()])),
                }),
            ),
            (
                "Result::err",
                result_forall.clone(),
                Type::Function(FnType {
                    params: vec![],
                    ret: Box::new(Type::Enum(Rc::from("Option"), vec![re.clone()])),
                }),
            ),
            (
                "Result::or_else",
                vec![rt_id, re_id, rf_id],
                Type::Function(FnType {
                    params: vec![Type::Function(FnType {
                        params: vec![re.clone()],
                        ret: Box::new(Type::Enum(Rc::from("Result"), vec![rt.clone(), rf.clone()])),
                    })],
                    ret: Box::new(Type::Enum(Rc::from("Result"), vec![rt, rf])),
                }),
            ),
        ];
        for (name, forall, fn_type) in result_methods {
            self.builtin_method_schemes.insert(
                name.to_string(),
                BuiltinMethodScheme {
                    forall,
                    self_type: result_ty.clone(),
                    fn_type,
                },
            );
        }
    }

    fn register_builtin_enums(&mut self) {
        self.register_enum(EnumDef {
            name: "Option".to_string(),
            generic_params: vec!["T".to_string()],
            variants: vec![
                VariantDef {
                    name: "Some".to_string(),
                    kind: VariantKind::Tuple(vec![Type::Generic(Rc::from("T"), vec![])]),
                },
                VariantDef {
                    name: "None".to_string(),
                    kind: VariantKind::Unit,
                },
            ],
        });
        self.register_enum(EnumDef {
            name: "Result".to_string(),
            generic_params: vec!["T".to_string(), "E".to_string()],
            variants: vec![
                VariantDef {
                    name: "Ok".to_string(),
                    kind: VariantKind::Tuple(vec![Type::Generic(Rc::from("T"), vec![])]),
                },
                VariantDef {
                    name: "Err".to_string(),
                    kind: VariantKind::Tuple(vec![Type::Generic(Rc::from("E"), vec![])]),
                },
            ],
        });
        self.register_enum(EnumDef {
            name: "ParseIntError".to_string(),
            generic_params: vec![],
            variants: vec![
                VariantDef {
                    name: "Empty".to_string(),
                    kind: VariantKind::Unit,
                },
                VariantDef {
                    name: "InvalidDigit".to_string(),
                    kind: VariantKind::Unit,
                },
                VariantDef {
                    name: "Overflow".to_string(),
                    kind: VariantKind::Unit,
                },
            ],
        });
    }

    pub fn register_struct(&mut self, def: StructDef) {
        self.structs.insert(def.name.clone(), def);
    }

    pub fn register_enum(&mut self, def: EnumDef) {
        self.enums.insert(def.name.clone(), def);
    }

    pub fn register_trait(&mut self, def: TraitDef) {
        self.traits.insert(def.name.clone(), def);
    }

    pub fn register_impl(&mut self, def: ImplDef) {
        self.impls.push(def);
    }

    pub fn lookup_struct(&self, name: &str) -> Option<&StructDef> {
        self.structs.get(name)
    }

    pub fn lookup_enum(&self, name: &str) -> Option<&EnumDef> {
        self.enums.get(name)
    }

    pub fn lookup_trait(&self, name: &str) -> Option<&TraitDef> {
        self.traits.get(name)
    }

    pub fn lookup_method(&self, self_type: &Type, method_name: &str) -> Option<&Type> {
        // Inherent impls first
        self.impls
            .iter()
            .filter(|imp| imp.trait_name.is_none() && imp.self_type == *self_type)
            .find_map(|imp| {
                imp.methods
                    .iter()
                    .find(|(name, _)| name == method_name)
                    .map(|(_, ty)| ty)
            })
            .or_else(|| {
                // Trait impls
                self.impls
                    .iter()
                    .filter(|imp| imp.trait_name.is_some() && imp.self_type == *self_type)
                    .find_map(|imp| {
                        imp.methods
                            .iter()
                            .find(|(name, _)| name == method_name)
                            .map(|(_, ty)| ty)
                    })
            })
    }

    pub fn lookup_method_by_name(&self, type_name: &str, method_name: &str) -> Option<&Type> {
        self.impls
            .iter()
            .filter(|imp| match &imp.self_type {
                Type::Struct(n, _) | Type::Enum(n, _) => n.as_ref() == type_name,
                _ => false,
            })
            .find_map(|imp| {
                imp.methods
                    .iter()
                    .find(|(name, _)| name == method_name)
                    .map(|(_, ty)| ty)
            })
    }

    pub fn instantiate_builtin_method(
        &mut self,
        receiver: &Type,
        method_name: &str,
        subst: &Substitution,
    ) -> Option<(Type, Type)> {
        let prefix = match receiver {
            Type::Vector(_) => "Vector",
            Type::Char => "char",
            Type::Str => "str",
            Type::Map(..) => "Map",
            Type::Set(_) => "Set",
            Type::Enum(name, _) if &**name == "Option" => "Option",
            Type::Enum(name, _) if &**name == "Result" => "Result",
            _ => return None,
        };
        let key = format!("{prefix}::{method_name}");
        let scheme = self.builtin_method_schemes.get(&key)?.clone();
        if scheme.forall.is_empty() {
            return Some((subst.apply(&scheme.self_type), subst.apply(&scheme.fn_type)));
        }
        let mut local_subst = subst.clone();
        for &var in &scheme.forall {
            let fresh = self.fresh_var();
            local_subst.bind(var, fresh);
        }
        Some((
            local_subst.apply(&scheme.self_type),
            local_subst.apply(&scheme.fn_type),
        ))
    }

    pub fn all_enums(&self) -> impl Iterator<Item = &EnumDef> {
        self.enums.values()
    }

    pub fn all_impls(&self) -> impl Iterator<Item = &ImplDef> {
        self.impls.iter()
    }

    pub fn required_trait_methods(&self, trait_name: &str) -> Vec<&MethodSig> {
        self.traits
            .get(trait_name)
            .map(|t| t.methods.iter().filter(|m| !m.has_default).collect())
            .unwrap_or_default()
    }

    pub fn resolve_type(&mut self, expr: &TypeExpr) -> Type {
        match expr {
            TypeExpr::Named(named) => self.resolve_named_type(&named.name),
            TypeExpr::Generic(name, args, _) => {
                let resolved_args = args
                    .iter()
                    .map(|a| self.resolve_type(a))
                    .collect::<Vec<Type>>();
                if name == "Vector" {
                    let elem = resolved_args
                        .into_iter()
                        .next()
                        .unwrap_or_else(|| self.fresh_var());
                    Type::Vector(Box::new(elem))
                } else if name == "Map" {
                    let mut args = resolved_args.into_iter();
                    let k = args.next().unwrap_or_else(|| self.fresh_var());
                    let v = args.next().unwrap_or_else(|| self.fresh_var());
                    Type::Map(Box::new(k), Box::new(v))
                } else if self.structs.contains_key(name) {
                    Type::Struct(Rc::from(name.as_str()), resolved_args)
                } else if self.enums.contains_key(name) {
                    Type::Enum(Rc::from(name.as_str()), resolved_args)
                } else {
                    Type::Generic(Rc::from(name.as_str()), vec![])
                }
            }
            _ => resolve_type_expr(expr),
        }
    }

    fn resolve_named_type(&mut self, name: &str) -> Type {
        if name == "Vector" {
            return Type::Vector(Box::new(self.fresh_var()));
        }
        if name == "Set" {
            return Type::Set(Box::new(self.fresh_var()));
        }
        if name == "Map" {
            return Type::Map(Box::new(self.fresh_var()), Box::new(self.fresh_var()));
        }
        if let Some(count) = self.structs.get(name).map(|d| d.generic_params.len()) {
            let args = (0..count).map(|_| self.fresh_var()).collect();
            Type::Struct(Rc::from(name), args)
        } else if let Some(count) = self.enums.get(name).map(|d| d.generic_params.len()) {
            let args = (0..count).map(|_| self.fresh_var()).collect();
            Type::Enum(Rc::from(name), args)
        } else {
            resolve_type_expr(&TypeExpr::Named(NamedType {
                name: name.to_string(),
                span: maat_span::Span::ZERO,
            }))
        }
    }

    pub fn fresh_var(&mut self) -> Type {
        let id = self.next_var;
        self.next_var += 1;
        Type::Var(id)
    }

    pub fn fresh_int_var(&mut self) -> Type {
        let id = self.next_var;
        self.next_var += 1;
        Type::IntVar(id)
    }

    pub fn define_var(&mut self, name: &str, ty: Type) {
        self.define_scheme(name, TypeScheme::monomorphic(ty));
    }

    pub fn define_scheme(&mut self, name: &str, scheme: TypeScheme) {
        let scope_start = self.scope_starts.last().copied().unwrap_or(0);
        if let Some(entry) = self.bindings[scope_start..]
            .iter_mut()
            .find(|(n, _)| n == name)
        {
            entry.1 = scheme;
        } else {
            self.bindings.push((name.to_string(), scheme));
        }
    }

    pub fn lookup_scheme(&self, name: &str) -> Option<&TypeScheme> {
        self.bindings
            .iter()
            .rev()
            .find(|(n, _)| n == name)
            .map(|(_, scheme)| scheme)
    }

    pub fn instantiate(&mut self, name: &str, subst: &Substitution) -> Option<Type> {
        let scheme = self.lookup_scheme(name)?.clone();
        Some(self.instantiate_scheme(&scheme, subst))
    }

    pub fn instantiate_scheme(&mut self, scheme: &TypeScheme, subst: &Substitution) -> Type {
        if scheme.forall.is_empty() {
            return subst.apply(&scheme.ty);
        }
        let mut local_subst = subst.clone();
        for &var in &scheme.forall {
            let fresh = self.fresh_var();
            local_subst.bind(var, fresh);
        }
        local_subst.apply(&scheme.ty)
    }

    pub fn generalize(&self, ty: &Type, subst: &Substitution) -> TypeScheme {
        let resolved = subst.apply(ty);
        let ty_vars = resolved.free_type_vars();
        let env_vars = self.free_env_vars(subst);
        let forall = ty_vars
            .difference(&env_vars)
            .copied()
            .collect::<Vec<TypeVarId>>();
        TypeScheme {
            forall,
            ty: resolved,
        }
    }

    fn free_env_vars(&self, subst: &Substitution) -> HashSet<TypeVarId> {
        let mut vars = HashSet::new();
        for (_, scheme) in &self.bindings {
            let resolved = subst.apply(&scheme.ty);
            let scheme_free = resolved.free_type_vars();
            let quantified = scheme
                .forall
                .iter()
                .copied()
                .collect::<HashSet<TypeVarId>>();
            vars.extend(scheme_free.difference(&quantified));
        }
        vars
    }

    pub fn push_scope(&mut self) {
        self.scope_starts.push(self.bindings.len());
    }

    pub fn pop_scope(&mut self) {
        if self.scope_starts.len() > 1 {
            let start = self
                .scope_starts
                .pop()
                .expect("scope_starts checked non-empty");
            self.bindings.truncate(start);
        }
    }
}

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

    #[test]
    fn define_and_lookup_var() {
        let mut env = TypeEnv::new();
        env.define_var("x", Type::I64);
        let subst = Substitution::new();
        let ty = env.instantiate("x", &subst);
        assert_eq!(ty, Some(Type::I64));
        assert_eq!(env.instantiate("y", &subst), None);
    }

    #[test]
    fn scope_shadowing() {
        let mut env = TypeEnv::new();
        let subst = Substitution::new();
        env.define_var("x", Type::I64);
        env.push_scope();
        env.define_var("x", Type::Bool);
        assert_eq!(env.instantiate("x", &subst), Some(Type::Bool));
        env.pop_scope();
        assert_eq!(env.instantiate("x", &subst), Some(Type::I64));
    }

    #[test]
    fn fresh_variables() {
        let mut env = TypeEnv::new();
        assert_eq!(env.fresh_var(), Type::Var(0));
        assert_eq!(env.fresh_var(), Type::Var(1));
        assert_eq!(env.fresh_var(), Type::Var(2));
    }

    #[test]
    fn generalize_and_instantiate() {
        let mut env = TypeEnv::new();
        let subst = Substitution::new();
        // Type: fn(?T0) -> ?T0  (identity function)
        let fn_ty = Type::Function(FnType {
            params: vec![Type::Var(0)],
            ret: Box::new(Type::Var(0)),
        });
        env.next_var = 1;
        let scheme = env.generalize(&fn_ty, &subst);
        assert_eq!(scheme.forall, vec![0]);
        let inst1 = env.instantiate_scheme(&scheme, &subst);
        let inst2 = env.instantiate_scheme(&scheme, &subst);
        match (&inst1, &inst2) {
            (Type::Function(f1), Type::Function(f2)) => {
                assert_ne!(f1.params[0], f2.params[0]);
            }
            _ => panic!("expected function types"),
        }
    }

    #[test]
    fn monomorphic_not_generalized() {
        let mut env = TypeEnv::new();
        let subst = Substitution::new();
        // Define x: ?T0 in the environment (simulates a lambda parameter)
        env.define_var("x", Type::Var(0));
        env.next_var = 1;
        // ?T0 is free in the env, so `generalize` should NOT quantify it
        let scheme = env.generalize(&Type::Var(0), &subst);
        assert!(scheme.forall.is_empty());
    }

    #[test]
    fn builtin_method_instantiation_produces_fresh_vars() {
        let mut env = TypeEnv::new();
        env.register_builtins();
        let subst = Substitution::new();
        let i64_vector = Type::Vector(Box::new(Type::I64));
        let str_vector = Type::Vector(Box::new(Type::Str));

        // Two separate instantiations should produce independent type variables.
        let (self1, fn1) = env
            .instantiate_builtin_method(&i64_vector, "first", &subst)
            .expect("Vector::first should exist");
        let (self2, fn2) = env
            .instantiate_builtin_method(&str_vector, "first", &subst)
            .expect("Vector::first should exist");
        match (&self1, &self2) {
            (Type::Vector(a), Type::Vector(b)) => {
                assert_ne!(
                    a, b,
                    "each instantiation must use independent type variables"
                );
            }
            _ => panic!("expected Vector types"),
        }
        match (&fn1, &fn2) {
            (Type::Function(f1), Type::Function(f2)) => {
                assert_ne!(
                    f1.ret, f2.ret,
                    "return types must use independent type variables"
                );
            }
            _ => panic!("expected Function types"),
        }
    }

    #[test]
    fn builtin_method_unification_resolves_element_type() {
        let mut env = TypeEnv::new();
        env.register_builtins();
        let mut subst = Substitution::new();
        let i64_vector = Type::Vector(Box::new(Type::I64));
        let (inst_self, inst_fn) = env
            .instantiate_builtin_method(&i64_vector, "first", &subst)
            .expect("Vector::first should exist");
        // Unify the instantiated self-type with the actual receiver.
        subst
            .unify(&inst_self, &i64_vector)
            .expect("unification should succeed");
        match subst.apply(&inst_fn) {
            Type::Function(fn_ty) => {
                let expected = Type::Enum(Rc::from("Option"), vec![Type::I64]);
                assert_eq!(
                    *fn_ty.ret, expected,
                    "return type should resolve to Option<i64>"
                );
            }
            _ => panic!("expected Function type"),
        }
    }
}
