//! Type environment for type inference

use crate::middleend::types::{MonoType, TyVarGenerator, TypeScheme};

use std::collections::HashMap;

/// Type environment mapping identifiers to type schemes

#[derive(Debug, Clone)]

pub struct TypeEnv {

    bindings: HashMap<String, TypeScheme>,

}

impl TypeEnv {

    #[must_use]

    pub fn new() -> Self {

        TypeEnv {

            bindings: HashMap::new(),

        }

    }

    /// Create a standard environment with built-in functions

    #[must_use]

    pub fn standard() -> Self {

        let mut env = Self::new();

        // Arithmetic functions

        env.bind(

            "add",

            TypeScheme::mono(MonoType::Function(

                Box::new(MonoType::Int),

                Box::new(MonoType::Function(

                    Box::new(MonoType::Int),

                    Box::new(MonoType::Int),

                )),

            )),

        );

        // IO functions

        env.bind(

            "print",

            TypeScheme::mono(MonoType::Function(

                Box::new(MonoType::String),

                Box::new(MonoType::Unit),

            )),

        );

        env.bind(

            "println",

            TypeScheme::mono(MonoType::Function(

                Box::new(MonoType::String),

                Box::new(MonoType::Unit),

            )),

        );

        // Comparison functions

        env.bind(

            "eq",

            TypeScheme::mono(MonoType::Function(

                Box::new(MonoType::Int),

                Box::new(MonoType::Function(

                    Box::new(MonoType::Int),

                    Box::new(MonoType::Bool),

                )),

            )),

        );

        env

    }

    /// Bind a name to a type scheme

    pub fn bind(&mut self, name: impl Into<String>, scheme: TypeScheme) {

        self.bindings.insert(name.into(), scheme);

    }

    /// Look up a name in the environment

    #[must_use]

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

        self.bindings.get(name)

    }

    /// Extend the environment with a new binding (functional style)

    #[must_use]

    pub fn extend(&self, name: impl Into<String>, scheme: TypeScheme) -> Self {

        let mut new_env = self.clone();

        new_env.bind(name, scheme);

        new_env

    }

    /// Get free type variables in the environment

    #[must_use]

    pub fn free_vars(&self) -> Vec<crate::middleend::types::TyVar> {

        let mut vars = Vec::new();

        for scheme in self.bindings.values() {

            // Only collect free variables not bound by the scheme

            let scheme_free = scheme.ty.free_vars();

            for var in scheme_free {

                if !scheme.vars.contains(&var) {

                    vars.push(var);

                }

            }

        }

        vars

    }

    /// Generalize a monomorphic type to a type scheme

    #[must_use]

    pub fn generalize(&self, ty: MonoType) -> TypeScheme {

        let ty_vars = ty.free_vars();

        let env_vars = self.free_vars();

        // Variables to generalize are those in ty but not in env

        let gen_vars: Vec<_> = ty_vars

            .into_iter()

            .filter(|v| !env_vars.contains(v))

            .collect();

        TypeScheme { vars: gen_vars, ty }

    }

    /// Instantiate a type scheme with fresh variables

    pub fn instantiate(&self, scheme: &TypeScheme, gen: &mut TyVarGenerator) -> MonoType {

        scheme.instantiate(gen)

    }

}

impl Default for TypeEnv {

    fn default() -> Self {

        Self::new()

    }

}

#[cfg(test)]

#[allow(clippy::unwrap_used, clippy::panic)]

mod tests {

    use super::*;

    use crate::middleend::types::TyVar;

    #[test]

    fn test_env_lookup() {

        let mut env = TypeEnv::new();

        env.bind("x", TypeScheme::mono(MonoType::Int));

        assert!(env.lookup("x").is_some());

        assert!(env.lookup("y").is_none());

    }

    #[test]

    fn test_env_extend() {

        let env = TypeEnv::new();

        let env2 = env.extend("x", TypeScheme::mono(MonoType::Bool));

        assert!(env.lookup("x").is_none());

        assert!(env2.lookup("x").is_some());

    }

    #[test]

    fn test_generalization() {

        let env = TypeEnv::new();

        let var = TyVar(0);

        // A type with a free variable

        let ty = MonoType::Function(

            Box::new(MonoType::Var(var.clone())),

            Box::new(MonoType::Var(var.clone())),

        );

        let scheme = env.generalize(ty);

        // The variable should be generalized

        assert_eq!(scheme.vars.len(), 1);

        assert!(scheme.vars.contains(&var));

    }

    #[test]

    fn test_no_generalization_with_env_vars() {

        let mut env = TypeEnv::new();

        let var = TyVar(0);

        // Add a binding with the same variable to the environment

        env.bind("y", TypeScheme::mono(MonoType::Var(var.clone())));

        // Try to generalize a type with the same variable

        let ty = MonoType::Function(Box::new(MonoType::Var(var)), Box::new(MonoType::Int));

        let scheme = env.generalize(ty);

        // The variable should NOT be generalized (it's in the env)

        assert_eq!(scheme.vars.len(), 0);

    }

    #[test]

    fn test_standard_env() {

        let env = TypeEnv::standard();

        assert!(env.lookup("println").is_some());

        assert!(env.lookup("print").is_some());

        assert!(env.lookup("add").is_some());

        assert!(env.lookup("eq").is_some());

    }

    #[test]

    fn test_default_env() {

        let env = TypeEnv::default();

        assert!(env.lookup("nonexistent").is_none());

        assert_eq!(env.bindings.len(), 0);

    }

    #[test]

    fn test_multiple_bindings() {

        let mut env = TypeEnv::new();

        env.bind("x", TypeScheme::mono(MonoType::Int));

        env.bind("y", TypeScheme::mono(MonoType::Bool));

        env.bind("z", TypeScheme::mono(MonoType::String));

        assert!(env.lookup("x").is_some());

        assert!(env.lookup("y").is_some());

        assert!(env.lookup("z").is_some());

        assert!(env.lookup("w").is_none());

    }

    #[test]

    fn test_bind_overwrites() {

        let mut env = TypeEnv::new();

        env.bind("x", TypeScheme::mono(MonoType::Int));

        env.bind("x", TypeScheme::mono(MonoType::Bool));

        let scheme = env.lookup("x").unwrap();

        match &scheme.ty {

            MonoType::Bool => {}, // Expected

            _ => panic!("Expected Bool type after overwrite"),

        }

    }

    #[test]

    fn test_env_clone() {

        let mut env1 = TypeEnv::new();

        env1.bind("x", TypeScheme::mono(MonoType::Int));

        let env2 = env1.clone();

        assert!(env2.lookup("x").is_some());

    }

    #[test]

    fn test_free_vars_empty() {

        let env = TypeEnv::new();

        assert!(env.free_vars().is_empty());

    }

    #[test]

    fn test_free_vars_with_schemes() {

        let mut env = TypeEnv::new();

        let var1 = TyVar(1);

        let var2 = TyVar(2);

        // Add a scheme with a free variable

        let scheme1 = TypeScheme {

            vars: vec![],

            ty: MonoType::Var(var1.clone()),

        };

        env.bind("x", scheme1);

        // Add a scheme with a bound variable

        let scheme2 = TypeScheme {

            vars: vec![var2.clone()],

            ty: MonoType::Var(var2),

        };

        env.bind("y", scheme2);

        let free_vars = env.free_vars();

        assert!(free_vars.contains(&var1));

        assert!(!free_vars.contains(&TyVar(2))); // var2 is bound

    }

    #[test]

    fn test_generalize_empty_env() {

        let env = TypeEnv::new();

        let var = TyVar(5);

        let ty = MonoType::Var(var.clone());

        let scheme = env.generalize(ty);

        assert_eq!(scheme.vars.len(), 1);

        assert!(scheme.vars.contains(&var));

    }

    #[test]

    fn test_generalize_complex_type() {

        let env = TypeEnv::new();

        let var1 = TyVar(10);

        let var2 = TyVar(11);

        let ty = MonoType::Function(

            Box::new(MonoType::Var(var1.clone())),

            Box::new(MonoType::Function(

                Box::new(MonoType::Var(var2.clone())),

                Box::new(MonoType::Int),

            )),

        );

        let scheme = env.generalize(ty);

        assert_eq!(scheme.vars.len(), 2);

        assert!(scheme.vars.contains(&var1));

        assert!(scheme.vars.contains(&var2));

    }

    #[test]

    fn test_instantiate_scheme() {

        let env = TypeEnv::new();

        let mut gen = TyVarGenerator::new();

        // Create a polymorphic scheme: forall a. a -> a

        let var = TyVar(20);

        let scheme = TypeScheme {

            vars: vec![var.clone()],

            ty: MonoType::Function(

                Box::new(MonoType::Var(var.clone())),

                Box::new(MonoType::Var(var)),

            ),

        };

        let instance = env.instantiate(&scheme, &mut gen);

        // Should get fresh variables

        match instance {

            MonoType::Function(arg, ret) => {

                match (*arg, *ret) {

                    (MonoType::Var(v1), MonoType::Var(v2)) => {

                        assert_eq!(v1, v2); // Same fresh variable

                        assert_ne!(v1, TyVar(20)); // Different from original

                    }

                    _ => panic!("Expected function with variable types"),

                }

            }

            _ => panic!("Expected function type"),

        }

    }

    #[test]

    fn test_standard_env_function_types() {

        let env = TypeEnv::standard();

        // Test add function type

        let add_scheme = env.lookup("add").unwrap();

        match &add_scheme.ty {

            MonoType::Function(arg1, rest) => {

                assert!(matches!(**arg1, MonoType::Int));

                match rest.as_ref() {

                    MonoType::Function(arg2, ret_type) => {

                        assert!(matches!(**arg2, MonoType::Int));

                        assert!(matches!(**ret_type, MonoType::Int));

                    }

                    _ => panic!("Expected curried function type"),

                }

            }

            _ => panic!("Expected function type for add"),

        }

        // Test print function type

        let print_scheme = env.lookup("print").unwrap();

        match &print_scheme.ty {

            MonoType::Function(arg, ret) => {

                assert!(matches!(**arg, MonoType::String));

                assert!(matches!(**ret, MonoType::Unit));

            }

            _ => panic!("Expected function type for print"),

        }

    }

}