use crate::ast::AstNode;
use crate::diagnostic::{ColorConfig, DiagnosticEmitter, Files, StandardEmitter, ToDiagnostic};
use crate::embedded_std::embedded_std_sources;
use crate::lexer::Lexer;
use crate::parser::Parser;
use crate::source::Source;
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};

pub struct ModuleResolver {
    base_path: PathBuf,
    compiled_modules: HashMap<String, Vec<AstNode>>,
    import_stack: Vec<String>,
    being_imported: HashSet<String>,
    canonical_cache: HashMap<PathBuf, String>, // canonical path -> module_path
    embedded_sources: &'static HashMap<String, String>,
}

impl ModuleResolver {
    pub fn new(base_path: PathBuf) -> Self {
        Self {
            base_path,
            compiled_modules: HashMap::new(),
            import_stack: Vec::new(),
            being_imported: HashSet::new(),
            canonical_cache: HashMap::new(),
            embedded_sources: embedded_std_sources(),
        }
    }

    fn resolve_embedded_key<'a>(&'a self, module_path: &'a str) -> Option<&'a str> {
        if self.embedded_sources.contains_key(module_path) {
            return Some(module_path);
        }

        module_path
            .strip_prefix("std.")
            .filter(|short| self.embedded_sources.contains_key(*short))
    }

    fn normalize_module_key<'a>(&'a self, module_path: &'a str) -> &'a str {
        self.resolve_embedded_key(module_path)
            .unwrap_or(module_path)
    }

    pub fn has_embedded_module(&self, module_path: &str) -> bool {
        self.resolve_embedded_key(module_path).is_some()
    }

    fn resolve_embedded_module_if_any(
        &mut self,
        module_path: &str,
        files: &mut Files,
    ) -> Result<Option<Vec<AstNode>>, String> {
        if let Some(embedded_key) = self.resolve_embedded_key(module_path) {
            let embedded_key = embedded_key.to_string();
            let cache_key = self.normalize_module_key(module_path).to_string();

            if let Some(nodes) = self.compiled_modules.get(&cache_key) {
                return Ok(Some(nodes.clone()));
            }

            if self.being_imported.contains(&cache_key) {
                return Err(format!(
                    "Circular import detected: {} -> {}",
                    self.import_stack.join(" -> "),
                    cache_key
                ));
            }

            self.being_imported.insert(cache_key.clone());
            self.import_stack.push(cache_key.clone());

            let source = self
                .embedded_sources
                .get(&embedded_key)
                .ok_or_else(|| format!("Embedded module not found: {}", module_path))?;

            let virtual_path =
                PathBuf::from(format!("<embedded>/{}.mux", embedded_key.replace('.', "/")));
            let nodes = self.parse_module(&virtual_path, files, Some(source.as_str()))?;
            return Ok(Some(nodes));
        }
        Ok(None)
    }

    fn determine_file_path(
        &self,
        module_path: &str,
        current_file: Option<&Path>,
    ) -> Result<PathBuf, String> {
        if module_path.starts_with("./") || module_path.starts_with("../") {
            // Relative import - resolve relative to current file
            let current_dir = current_file
                .and_then(|p| p.parent())
                .ok_or("Cannot resolve relative import: no current file")?;

            let relative_path = module_path.trim_start_matches("./");
            let mut path = current_dir.to_path_buf();

            // Handle ../ parts
            for part in relative_path.split('/') {
                if part == ".." {
                    path.pop();
                } else if !part.is_empty() {
                    path.push(part);
                }
            }
            path.set_extension("mux");
            Ok(path)
        } else if module_path.starts_with('/') {
            // Absolute import
            let mut path = PathBuf::from(module_path);
            path.set_extension("mux");
            Ok(path)
        } else {
            // Project-relative import (utils.logger)
            self.module_path_to_file(module_path)
        }
    }

    // Resolve import path relative to current file (for relative imports)
    pub fn resolve_import_path(
        &mut self,
        module_path: &str,
        current_file: Option<&Path>,
        files: &mut Files,
    ) -> Result<Vec<AstNode>, String> {
        if let Some(nodes) = self.resolve_embedded_module_if_any(module_path, files)? {
            return Ok(nodes);
        }

        let file_path = self.determine_file_path(module_path, current_file)?;

        // Canonicalize for cache key
        let canonical_path = file_path
            .canonicalize()
            .map_err(|e| format!("Cannot resolve module path {}: {}", module_path, e))?;

        // Check cache by canonical path
        if let Some(cached_module_path) = self.canonical_cache.get(&canonical_path)
            && let Some(nodes) = self.compiled_modules.get(cached_module_path)
        {
            return Ok(nodes.clone());
        }

        // Check circular imports
        if self.being_imported.contains(module_path) {
            return Err(format!(
                "Circular import detected: {} -> {}",
                self.import_stack.join(" -> "),
                module_path
            ));
        }

        // Mark as being imported
        self.being_imported.insert(module_path.to_string());
        self.import_stack.push(module_path.to_string());

        // Parse module
        let nodes = self.parse_module(&canonical_path, files, None)?;

        // Cache the canonical path
        self.canonical_cache
            .insert(canonical_path, module_path.to_string());

        Ok(nodes)
    }

    pub fn finish_import(&mut self, module_path: &str) {
        let cache_key = self.normalize_module_key(module_path).to_string();
        self.import_stack.pop();
        self.being_imported.remove(&cache_key);
    }

    pub fn cache_module(&mut self, module_path: &str, nodes: Vec<AstNode>) {
        let cache_key = self.normalize_module_key(module_path).to_string();
        self.compiled_modules.insert(cache_key, nodes);
    }

    /// Check if a module path resolves to a file, directory, both, or neither.
    /// Returns (has_file, has_directory) tuple.
    pub fn check_module_path(&self, module_path: &str) -> (bool, bool) {
        if self.has_embedded_module(module_path) {
            return (true, false);
        }

        let embedded_prefix = format!("{}.", module_path);
        let has_embedded_directory = self
            .embedded_sources
            .keys()
            .any(|key| key.starts_with(&embedded_prefix));

        let mut file_path = self.base_path.clone();
        for part in module_path.split('.') {
            file_path.push(part);
        }

        let mux_file = file_path.with_extension("mux");
        let dir_path = file_path;

        let has_file = mux_file.exists() && mux_file.is_file();
        let has_directory = dir_path.exists() && dir_path.is_dir();

        (has_file, has_directory || has_embedded_directory)
    }

    /// Get all .mux files in a directory module
    pub fn get_submodules(&self, module_path: &str) -> Result<Vec<String>, String> {
        let mut submodules = HashSet::new();
        let embedded_prefix = format!("{}.", module_path);
        for key in self.embedded_sources.keys() {
            if let Some(rest) = key.strip_prefix(&embedded_prefix)
                && !rest.is_empty()
                && let Some(child) = rest.split('.').next()
            {
                submodules.insert(child.to_string());
            }
        }

        let mut dir_path = self.base_path.clone();
        for part in module_path.split('.') {
            dir_path.push(part);
        }

        if dir_path.exists()
            && dir_path.is_dir()
            && let Ok(entries) = std::fs::read_dir(&dir_path)
        {
            for entry in entries.flatten() {
                let path = entry.path();
                if path.is_file()
                    && path.extension().is_some_and(|ext| ext == "mux")
                    && let Some(stem) = path.file_stem().and_then(|s| s.to_str())
                {
                    submodules.insert(stem.to_string());
                }
            }
        }

        if submodules.is_empty() {
            return Err(format!("Module directory not found: {}", module_path));
        }

        let mut result: Vec<String> = submodules.into_iter().collect();
        result.sort();
        Ok(result)
    }

    fn module_path_to_file(&self, module_path: &str) -> Result<PathBuf, String> {
        let mut path = self.base_path.clone();
        for part in module_path.split('.') {
            path.push(part);
        }
        path.set_extension("mux");

        if !path.exists() {
            return Err(format!(
                "Module not found: {} (looked for {:?})",
                module_path, path
            ));
        }

        Ok(path)
    }

    fn parse_module(
        &self,
        file_path: &Path,
        files: &mut Files,
        source_override: Option<&str>,
    ) -> Result<Vec<AstNode>, String> {
        let source_str = if let Some(src) = source_override {
            src.to_string()
        } else {
            std::fs::read_to_string(file_path)
                .map_err(|e| format!("Failed to open module: {}", e))?
        };

        let file_id = files.add(file_path, source_str.clone());
        let mut src = Source::from_string(source_str);

        let mut lex = Lexer::new(&mut src);
        let tokens = match lex.lex_all() {
            Ok(t) => t,
            Err(e) => {
                let emitter = StandardEmitter::new(ColorConfig::Auto);
                emitter.emit(&e.to_diagnostic(file_id), files);
                return Err(format!("Lexer error in module {}", file_path.display()));
            }
        };

        let mut parser = Parser::new(&tokens);
        match parser.parse() {
            Ok(nodes) => Ok(nodes),
            Err((_, errors)) => {
                let emitter = StandardEmitter::new(ColorConfig::Auto);
                let diagnostics: Vec<_> = errors.iter().map(|e| e.to_diagnostic(file_id)).collect();
                emitter.emit_batch(&diagnostics, files);
                Err(format!("Parse error in module {}", file_path.display()))
            }
        }
    }
}
