lux/src/modules.rs

//! Module system for the Lux language
//!
//! Handles loading, parsing, and resolving module imports.

#![allow(dead_code)]

use crate::ast::{Declaration, ImportDecl, Program, Visibility};
use crate::parser::Parser;
use std::collections::{HashMap, HashSet};
use std::fs;
use std::path::{Path, PathBuf};

/// Error during module loading
#[derive(Debug, Clone)]
pub struct ModuleError {
    pub message: String,
    pub module_path: String,
}

impl std::fmt::Display for ModuleError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Module error in '{}': {}",
            self.module_path, self.message
        )
    }
}

impl std::error::Error for ModuleError {}

/// A loaded and parsed module
#[derive(Debug, Clone)]
pub struct Module {
    /// The module's canonical path (e.g., "std/list")
    pub path: String,
    /// The parsed program
    pub program: Program,
    /// Names exported by this module (public declarations)
    pub exports: HashSet<String>,
}

impl Module {
    /// Get all public declarations from this module
    pub fn public_declarations(&self) -> Vec<&Declaration> {
        self.program
            .declarations
            .iter()
            .filter(|d| {
                match d {
                    Declaration::Function(f) => f.visibility == Visibility::Public,
                    Declaration::Let(l) => l.visibility == Visibility::Public,
                    Declaration::Type(t) => t.visibility == Visibility::Public,
                    Declaration::Trait(t) => t.visibility == Visibility::Public,
                    // Effects, handlers, and impls are always public for now
                    Declaration::Effect(_) | Declaration::Handler(_) | Declaration::Impl(_) => true,
                }
            })
            .collect()
    }
}

/// Module loader and resolver
pub struct ModuleLoader {
    /// Base directories to search for modules
    search_paths: Vec<PathBuf>,
    /// Cache of loaded modules (path -> module)
    cache: HashMap<String, Module>,
    /// Modules currently being loaded (for circular dependency detection)
    loading: HashSet<String>,
}

impl ModuleLoader {
    pub fn new() -> Self {
        let mut loader = Self {
            search_paths: vec![PathBuf::from(".")],
            cache: HashMap::new(),
            loading: HashSet::new(),
        };

        // Add package paths if in a project with lux.toml
        loader.add_package_paths();

        loader
    }

    /// Create a loader with custom search paths
    pub fn with_paths(paths: Vec<PathBuf>) -> Self {
        let mut loader = Self {
            search_paths: paths,
            cache: HashMap::new(),
            loading: HashSet::new(),
        };

        // Add package paths if in a project with lux.toml
        loader.add_package_paths();

        loader
    }

    /// Add a search path
    pub fn add_search_path(&mut self, path: PathBuf) {
        self.search_paths.push(path);
    }

    /// Add package search paths from .lux_packages directory
    fn add_package_paths(&mut self) {
        // Find project root by looking for lux.toml
        if let Some(project_root) = Self::find_project_root() {
            let packages_dir = project_root.join(".lux_packages");

            if packages_dir.exists() {
                // Add the packages directory itself
                self.search_paths.push(packages_dir.clone());

                // Add each installed package directory
                if let Ok(entries) = fs::read_dir(&packages_dir) {
                    for entry in entries.flatten() {
                        let path = entry.path();
                        if path.is_dir() {
                            self.search_paths.push(path);
                        }
                    }
                }
            }
        }
    }

    /// Find the project root by looking for lux.toml
    fn find_project_root() -> Option<PathBuf> {
        let mut current = std::env::current_dir().ok()?;

        loop {
            if current.join("lux.toml").exists() {
                return Some(current);
            }

            if !current.pop() {
                return None;
            }
        }
    }

    /// Resolve a module path to a file path
    fn resolve_path(&self, module_path: &str) -> Option<PathBuf> {
        // Convert module path (e.g., "std/list") to file path (e.g., "std/list.lux")
        let relative_path = format!("{}.lux", module_path);

        for search_path in &self.search_paths {
            // Try direct path (e.g., search_path/module_path.lux)
            let full_path = search_path.join(&relative_path);
            if full_path.exists() {
                return Some(full_path);
            }

            // For packages, try module_path/lib.lux (e.g., .lux_packages/http/lib.lux)
            let lib_path = search_path.join(module_path).join("lib.lux");
            if lib_path.exists() {
                return Some(lib_path);
            }

            // Also try module_path/src/lib.lux for common package layout
            let src_lib_path = search_path.join(module_path).join("src").join("lib.lux");
            if src_lib_path.exists() {
                return Some(src_lib_path);
            }
        }

        None
    }

    /// Load a module by its import path
    pub fn load_module(&mut self, module_path: &str) -> Result<&Module, ModuleError> {
        // Check if already cached
        if self.cache.contains_key(module_path) {
            return Ok(self.cache.get(module_path).unwrap());
        }

        // Check for circular dependency
        if self.loading.contains(module_path) {
            return Err(ModuleError {
                message: "Circular dependency detected".to_string(),
                module_path: module_path.to_string(),
            });
        }

        // Mark as loading
        self.loading.insert(module_path.to_string());

        // Resolve to file path
        let file_path = self.resolve_path(module_path).ok_or_else(|| ModuleError {
            message: format!("Module not found. Searched in: {:?}", self.search_paths),
            module_path: module_path.to_string(),
        })?;

        // Load the module
        let module = self.load_file(&file_path, module_path)?;

        // Remove from loading set
        self.loading.remove(module_path);

        // Cache the module
        self.cache.insert(module_path.to_string(), module);

        Ok(self.cache.get(module_path).unwrap())
    }

    /// Load a module from a file path
    fn load_file(&mut self, file_path: &Path, module_path: &str) -> Result<Module, ModuleError> {
        // Read the file
        let source = fs::read_to_string(file_path).map_err(|e| ModuleError {
            message: format!("Failed to read file: {}", e),
            module_path: module_path.to_string(),
        })?;

        // Parse the source
        let program = Parser::parse_source(&source).map_err(|e| ModuleError {
            message: format!("Parse error: {}", e),
            module_path: module_path.to_string(),
        })?;

        // Load any imports this module has
        for import in &program.imports {
            let import_path = import.path.to_string();
            self.load_module(&import_path)?;
        }

        // Collect exports
        let exports = self.collect_exports(&program);

        Ok(Module {
            path: module_path.to_string(),
            program,
            exports,
        })
    }

    /// Load a program from source (for REPL or direct execution)
    pub fn load_source(
        &mut self,
        source: &str,
        base_path: Option<&Path>,
    ) -> Result<Program, ModuleError> {
        // Add base path to search paths if provided
        if let Some(base) = base_path {
            if let Some(parent) = base.parent() {
                if !self.search_paths.contains(&parent.to_path_buf()) {
                    self.search_paths.push(parent.to_path_buf());
                }
            }
        }

        // Parse the source
        let program = Parser::parse_source(source).map_err(|e| ModuleError {
            message: format!("Parse error: {}", e),
            module_path: "<main>".to_string(),
        })?;

        // Load any imports
        for import in &program.imports {
            let import_path = import.path.to_string();
            self.load_module(&import_path)?;
        }

        Ok(program)
    }

    /// Collect exported names from a program
    fn collect_exports(&self, program: &Program) -> HashSet<String> {
        let mut exports = HashSet::new();

        for decl in &program.declarations {
            match decl {
                Declaration::Function(f) if f.visibility == Visibility::Public => {
                    exports.insert(f.name.name.clone());
                }
                Declaration::Let(l) if l.visibility == Visibility::Public => {
                    exports.insert(l.name.name.clone());
                }
                Declaration::Type(t) if t.visibility == Visibility::Public => {
                    exports.insert(t.name.name.clone());
                }
                Declaration::Effect(e) => {
                    // Effects are always exported
                    exports.insert(e.name.name.clone());
                }
                Declaration::Handler(h) => {
                    // Handlers are always exported
                    exports.insert(h.name.name.clone());
                }
                _ => {}
            }
        }

        exports
    }

    /// Get a cached module
    pub fn get_module(&self, module_path: &str) -> Option<&Module> {
        self.cache.get(module_path)
    }

    /// Get all loaded modules
    pub fn loaded_modules(&self) -> impl Iterator<Item = (&String, &Module)> {
        self.cache.iter()
    }

    /// Clear the module cache
    pub fn clear_cache(&mut self) {
        self.cache.clear();
    }

    /// Resolve imports for a program and return the names to be imported
    pub fn resolve_imports(
        &self,
        imports: &[ImportDecl],
    ) -> Result<HashMap<String, ResolvedImport>, ModuleError> {
        let mut resolved = HashMap::new();

        for import in imports {
            let module_path = import.path.to_string();
            let module = self.get_module(&module_path).ok_or_else(|| ModuleError {
                message: "Module not loaded".to_string(),
                module_path: module_path.clone(),
            })?;

            let import_name = if let Some(ref alias) = import.alias {
                // import foo/bar as Baz -> use "Baz" as the name
                alias.name.clone()
            } else {
                // import foo/bar -> use "bar" as the name (last segment)
                import
                    .path
                    .segments
                    .last()
                    .map(|s| s.name.clone())
                    .unwrap_or_else(|| module_path.clone())
            };

            if import.wildcard {
                // import foo.* -> import all exports directly
                for export in &module.exports {
                    resolved.insert(
                        export.clone(),
                        ResolvedImport {
                            module_path: module_path.clone(),
                            name: export.clone(),
                            kind: ImportKind::Direct,
                        },
                    );
                }
            } else if let Some(ref items) = import.items {
                // import foo.{a, b, c} -> import specific items
                for item in items {
                    if !module.exports.contains(&item.name) {
                        return Err(ModuleError {
                            message: format!("'{}' is not exported from module", item.name),
                            module_path: module_path.clone(),
                        });
                    }
                    resolved.insert(
                        item.name.clone(),
                        ResolvedImport {
                            module_path: module_path.clone(),
                            name: item.name.clone(),
                            kind: ImportKind::Direct,
                        },
                    );
                }
            } else {
                // import foo/bar -> import as module object
                resolved.insert(
                    import_name,
                    ResolvedImport {
                        module_path: module_path.clone(),
                        name: module_path.clone(),
                        kind: ImportKind::Module,
                    },
                );
            }
        }

        Ok(resolved)
    }
}

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

/// A resolved import
#[derive(Debug, Clone)]
pub struct ResolvedImport {
    /// The module path this import comes from
    pub module_path: String,
    /// The name being imported
    pub name: String,
    /// What kind of import this is
    pub kind: ImportKind,
}

/// Kind of import
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ImportKind {
    /// Import as a module object (import foo/bar)
    Module,
    /// Direct import of a name (import foo.{bar} or import foo.*)
    Direct,
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::TempDir;

    fn create_test_module(dir: &Path, name: &str, content: &str) -> PathBuf {
        let path = dir.join(format!("{}.lux", name));
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).unwrap();
        }
        let mut file = fs::File::create(&path).unwrap();
        file.write_all(content.as_bytes()).unwrap();
        path
    }

    #[test]
    fn test_load_simple_module() {
        let dir = TempDir::new().unwrap();
        create_test_module(
            dir.path(),
            "math",
            r#"
            pub fn add(a: Int, b: Int): Int = a + b
            pub fn sub(a: Int, b: Int): Int = a - b
            fn private_fn(): Int = 42
        "#,
        );

        let mut loader = ModuleLoader::with_paths(vec![dir.path().to_path_buf()]);
        let module = loader.load_module("math").unwrap();

        assert_eq!(module.path, "math");
        assert!(module.exports.contains("add"));
        assert!(module.exports.contains("sub"));
        assert!(!module.exports.contains("private_fn"));
    }

    #[test]
    fn test_load_nested_module() {
        let dir = TempDir::new().unwrap();
        create_test_module(
            dir.path(),
            "std/list",
            r#"
            pub fn length(list: List<Int>): Int = 0
        "#,
        );

        let mut loader = ModuleLoader::with_paths(vec![dir.path().to_path_buf()]);
        let module = loader.load_module("std/list").unwrap();

        assert_eq!(module.path, "std/list");
        assert!(module.exports.contains("length"));
    }

    #[test]
    fn test_module_not_found() {
        let dir = TempDir::new().unwrap();
        let mut loader = ModuleLoader::with_paths(vec![dir.path().to_path_buf()]);

        let result = loader.load_module("nonexistent");
        assert!(result.is_err());
        assert!(result.unwrap_err().message.contains("not found"));
    }

    #[test]
    fn test_circular_dependency_detection() {
        let dir = TempDir::new().unwrap();
        create_test_module(
            dir.path(),
            "a",
            r#"
            import b
            pub fn foo(): Int = 1
        "#,
        );
        create_test_module(
            dir.path(),
            "b",
            r#"
            import a
            pub fn bar(): Int = 2
        "#,
        );

        let mut loader = ModuleLoader::with_paths(vec![dir.path().to_path_buf()]);
        let result = loader.load_module("a");

        assert!(result.is_err());
        assert!(result.unwrap_err().message.contains("Circular"));
    }

    #[test]
    fn test_module_caching() {
        let dir = TempDir::new().unwrap();
        create_test_module(
            dir.path(),
            "cached",
            r#"
            pub fn foo(): Int = 42
        "#,
        );

        let mut loader = ModuleLoader::with_paths(vec![dir.path().to_path_buf()]);

        // Load twice
        loader.load_module("cached").unwrap();
        loader.load_module("cached").unwrap();

        // Should only be in cache once
        assert_eq!(loader.cache.len(), 1);
    }

    #[test]
    fn test_end_to_end_module_import() {
        use crate::interpreter::Interpreter;
        use crate::typechecker::TypeChecker;

        let dir = TempDir::new().unwrap();

        // Create a utility module with public functions
        create_test_module(
            dir.path(),
            "utils",
            r#"
            pub fn double(x: Int): Int = x * 2
            pub fn square(x: Int): Int = x * x
            fn private_helper(): Int = 0
        "#,
        );

        // Create a main program that imports and uses the module
        let main_source = r#"
            import utils

            let result = utils.double(21)
        "#;

        // Set up module loader
        let mut loader = ModuleLoader::with_paths(vec![dir.path().to_path_buf()]);

        // Load and parse the main program
        let main_path = dir.path().join("main.lux");
        let program = loader.load_source(main_source, Some(&main_path)).unwrap();

        // Type check with module support
        let mut checker = TypeChecker::new();
        checker
            .check_program_with_modules(&program, &loader)
            .unwrap();

        // Run with module support
        let mut interp = Interpreter::new();
        let result = interp.run_with_modules(&program, &loader).unwrap();

        // Should evaluate to 42
        assert_eq!(format!("{}", result), "42");
    }

    #[test]
    fn test_selective_import() {
        use crate::interpreter::Interpreter;
        use crate::typechecker::TypeChecker;

        let dir = TempDir::new().unwrap();

        // Create a module with multiple exports
        create_test_module(
            dir.path(),
            "math",
            r#"
            pub fn add(a: Int, b: Int): Int = a + b
            pub fn mul(a: Int, b: Int): Int = a * b
        "#,
        );

        // Import only the add function
        let main_source = r#"
            import math.{add}

            let result = add(10, 5)
        "#;

        let mut loader = ModuleLoader::with_paths(vec![dir.path().to_path_buf()]);
        let main_path = dir.path().join("main.lux");
        let program = loader.load_source(main_source, Some(&main_path)).unwrap();

        let mut checker = TypeChecker::new();
        checker
            .check_program_with_modules(&program, &loader)
            .unwrap();

        let mut interp = Interpreter::new();
        let result = interp.run_with_modules(&program, &loader).unwrap();

        assert_eq!(format!("{}", result), "15");
    }

    #[test]
    fn test_module_with_alias() {
        use crate::interpreter::Interpreter;
        use crate::typechecker::TypeChecker;

        let dir = TempDir::new().unwrap();

        // Create a nested module
        create_test_module(
            dir.path(),
            "lib/helpers",
            r#"
            pub fn greet(): String = "hello"
        "#,
        );

        // Import with alias
        let main_source = r#"
            import lib/helpers as h

            let result = h.greet()
        "#;

        let mut loader = ModuleLoader::with_paths(vec![dir.path().to_path_buf()]);
        let main_path = dir.path().join("main.lux");
        let program = loader.load_source(main_source, Some(&main_path)).unwrap();

        let mut checker = TypeChecker::new();
        checker
            .check_program_with_modules(&program, &loader)
            .unwrap();

        let mut interp = Interpreter::new();
        let result = interp.run_with_modules(&program, &loader).unwrap();

        assert_eq!(format!("{}", result), "\"hello\"");
    }

    #[test]
    fn test_transitive_imports() {
        use crate::interpreter::Interpreter;
        use crate::typechecker::TypeChecker;

        let dir = TempDir::new().unwrap();

        // Create base module
        create_test_module(
            dir.path(),
            "base",
            r#"
            pub fn value(): Int = 100
        "#,
        );

        // Create mid module that imports base
        create_test_module(
            dir.path(),
            "mid",
            r#"
            import base

            pub fn doubled(): Int = base.value() * 2
        "#,
        );

        // Create main that imports mid
        let main_source = r#"
            import mid

            let result = mid.doubled()
        "#;

        let mut loader = ModuleLoader::with_paths(vec![dir.path().to_path_buf()]);
        let main_path = dir.path().join("main.lux");
        let program = loader.load_source(main_source, Some(&main_path)).unwrap();

        let mut checker = TypeChecker::new();
        checker
            .check_program_with_modules(&program, &loader)
            .unwrap();

        let mut interp = Interpreter::new();
        let result = interp.run_with_modules(&program, &loader).unwrap();

        assert_eq!(format!("{}", result), "200");
    }

    #[test]
    fn test_package_lib_lux_entry_point() {
        // Test that packages with lib.lux are correctly resolved
        let dir = TempDir::new().unwrap();

        // Create a package structure: mypackage/lib.lux
        let pkg_dir = dir.path().join("mypackage");
        fs::create_dir_all(&pkg_dir).unwrap();

        let lib_content = r#"
            pub fn getValue(): Int = 42
        "#;
        fs::write(pkg_dir.join("lib.lux"), lib_content).unwrap();

        // The loader should find mypackage via lib.lux
        let mut loader = ModuleLoader::with_paths(vec![dir.path().to_path_buf()]);
        let module = loader.load_module("mypackage");

        assert!(module.is_ok(), "Should find package via lib.lux");
        let module = module.unwrap();
        assert!(module.exports.contains("getValue"));
    }

    #[test]
    fn test_package_src_lib_lux_entry_point() {
        // Test that packages with src/lib.lux are correctly resolved
        let dir = TempDir::new().unwrap();

        // Create a package structure: mypackage/src/lib.lux
        let src_dir = dir.path().join("anotherpackage").join("src");
        fs::create_dir_all(&src_dir).unwrap();

        let lib_content = r#"
            pub fn compute(): Int = 100
        "#;
        fs::write(src_dir.join("lib.lux"), lib_content).unwrap();

        // The loader should find anotherpackage via src/lib.lux
        let mut loader = ModuleLoader::with_paths(vec![dir.path().to_path_buf()]);
        let module = loader.load_module("anotherpackage");

        assert!(module.is_ok(), "Should find package via src/lib.lux");
        let module = module.unwrap();
        assert!(module.exports.contains("compute"));
    }

    #[test]
    fn test_lux_packages_directory_resolution() {
        use crate::interpreter::Interpreter;
        use crate::typechecker::TypeChecker;

        let dir = TempDir::new().unwrap();

        // Create a lux.toml to make this a "project"
        fs::write(
            dir.path().join("lux.toml"),
            "[project]\nname = \"testproj\"\nversion = \"0.1.0\"\n[dependencies]\n",
        )
        .unwrap();

        // Create .lux_packages/mylib/lib.lux
        let pkg_dir = dir.path().join(".lux_packages").join("mylib");
        fs::create_dir_all(&pkg_dir).unwrap();
        fs::write(
            pkg_dir.join("lib.lux"),
            "pub fn helper(): Int = 999\n",
        )
        .unwrap();

        // Create main.lux that imports mylib
        let main_source = r#"
            import mylib

            let result = mylib.helper()
        "#;

        // Change to the project directory to test find_project_root
        let original_dir = std::env::current_dir().unwrap();
        std::env::set_current_dir(dir.path()).unwrap();

        let mut loader = ModuleLoader::new();
        let main_path = dir.path().join("main.lux");
        let program = loader.load_source(main_source, Some(&main_path)).unwrap();

        let mut checker = TypeChecker::new();
        checker
            .check_program_with_modules(&program, &loader)
            .unwrap();

        let mut interp = Interpreter::new();
        let result = interp.run_with_modules(&program, &loader).unwrap();

        // Restore original directory
        std::env::set_current_dir(original_dir).unwrap();

        assert_eq!(format!("{}", result), "999");
    }
}