//! 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"),
        }
    }
}
