astmap-core 0.0.2

use std::collections::HashMap;
use std::path::Path;

use tracing::{debug, warn};

use crate::error::ScanError;
use crate::model::{DepLevel, DepType, Symbol};
use crate::port::{LanguageRegistry, LanguageResolver, ScanStore};

use super::types::{ImportResolutionMap, PendingImport};

/// Resolve imports to target file IDs and symbol-level cross-file deps.
///
/// Returns (resolution_map, dependency_count).
pub(super) fn resolve(
    project_dir: &Path,
    pending_imports: &[PendingImport],
    db: &dyn ScanStore,
    lang: &dyn LanguageRegistry,
) -> Result<(ImportResolutionMap, i64), ScanError> {
    let file_path_to_id = db.get_all_file_paths()?;

    // Build reverse map for O(1) file ID → path lookups
    let file_id_to_path: HashMap<i64, String> = file_path_to_id
        .iter()
        .map(|(path, &fid)| (fid, path.clone()))
        .collect();

    let mut resolvers: HashMap<String, Box<dyn LanguageResolver>> = HashMap::new();
    let mut import_resolution: ImportResolutionMap = HashMap::new();
    let mut dep_count = 0i64;

    for imp in pending_imports {
        // Get or create resolver (cached, necessarily stateful)
        let resolver = match resolvers.entry(imp.source_extension.clone()) {
            std::collections::hash_map::Entry::Occupied(e) => e.into_mut(),
            std::collections::hash_map::Entry::Vacant(e) => {
                match lang.resolver_for(&imp.source_extension, project_dir) {
                    Some(r) => e.insert(r),
                    None => {
                        warn!("no resolver for extension: {}", imp.source_extension);
                        continue;
                    }
                }
            }
        };

        // Pure: resolve import path to target file ID
        let target_file_id = resolver
            .resolve_import(&imp.import_path, &imp.source_rel_path, project_dir)
            .and_then(|abs_path| {
                let rel = abs_path
                    .strip_prefix(project_dir)
                    .unwrap_or(&abs_path)
                    .to_string_lossy()
                    .to_string();
                file_path_to_id.get(&rel).copied()
            });

        // Pure: resolve imported symbol names
        let resolved_names = target_file_id
            .map(|target_fid| {
                resolve_symbol_names(
                    imp,
                    target_fid,
                    &file_path_to_id,
                    &file_id_to_path,
                    db,
                    lang,
                )
            })
            .unwrap_or_default();

        // Effect: insert file-level import dependency
        db.insert_dependency(
            imp.source_file_id,
            None,
            target_file_id,
            None,
            DepType::Imports,
            DepLevel::File,
            Some(&imp.import_path),
        )?;
        dep_count += 1;

        // Collect resolution result (value-based, no mutation of shared state)
        if let Some(target_fid) = target_file_id {
            if !resolved_names.is_empty() {
                import_resolution.insert(
                    (imp.source_file_id, imp.import_path.clone()),
                    (target_fid, resolved_names),
                );
            }
        }
    }

    Ok((import_resolution, dep_count))
}

/// Resolve imported symbol names to their DB IDs by searching the target file
/// and its sibling files. Write-free — reads DB but does not write.
fn resolve_symbol_names(
    imp: &PendingImport,
    target_fid: i64,
    file_path_to_id: &HashMap<String, i64>,
    file_id_to_path: &HashMap<i64, String>,
    db: &dyn ScanStore,
    lang: &dyn LanguageRegistry,
) -> HashMap<String, i64> {
    let candidate_file_ids =
        collect_candidate_files(target_fid, file_path_to_id, file_id_to_path, lang);
    let is_wildcard = imp.imported_symbols.is_empty() && imp.import_path.ends_with("::*");

    // Fetch symbols from DB (read-only), then merge in pure function
    let symbol_lists: Vec<(i64, Vec<Symbol>)> = candidate_file_ids
        .iter()
        .filter_map(|&fid| db.get_symbols_for_file(fid).ok().map(|syms| (fid, syms)))
        .collect();

    if is_wildcard {
        let resolved = merge_symbol_maps(&symbol_lists);
        debug!(
            "wildcard import '{}' expanded to {} symbols",
            imp.import_path,
            resolved.len()
        );
        resolved
    } else if !imp.imported_symbols.is_empty() {
        let candidate_symbols = merge_symbol_maps(&symbol_lists);
        imp.imported_symbols
            .iter()
            .filter_map(|name| candidate_symbols.get(name).map(|&id| (name.clone(), id)))
            .collect()
    } else {
        HashMap::new()
    }
}

/// Merge pre-fetched symbol lists into a name → id map (first wins).
///
/// Pure function — no DB access.
fn merge_symbol_maps(symbol_lists: &[(i64, Vec<Symbol>)]) -> HashMap<String, i64> {
    symbol_lists
        .iter()
        .flat_map(|(_, syms)| syms.iter())
        .fold(HashMap::new(), |mut acc, sym| {
            acc.entry(sym.name.clone()).or_insert(sym.id);
            acc
        })
}

/// Collect candidate file IDs: the target file + any sibling files
/// as determined by the language registry.
///
/// Pure function — no DB access.
fn collect_candidate_files(
    target_fid: i64,
    file_path_to_id: &HashMap<String, i64>,
    file_id_to_path: &HashMap<i64, String>,
    lang: &dyn LanguageRegistry,
) -> Vec<i64> {
    let rel_path = match file_id_to_path.get(&target_fid) {
        Some(p) => p.as_str(),
        None => return vec![target_fid],
    };

    let mut ids = vec![target_fid];

    if let Some(expansion) = lang.sibling_expansion(rel_path) {
        ids.extend(
            file_path_to_id
                .iter()
                .filter(|(path, &fid)| {
                    if fid == target_fid {
                        return false;
                    }
                    // Extension filter (empty means any)
                    if !expansion.extension.is_empty()
                        && !path.ends_with(&format!(".{}", expansion.extension))
                    {
                        return false;
                    }
                    // Prefix / root-level filter
                    if expansion.root_only {
                        !path.contains('/')
                    } else {
                        path.starts_with(&expansion.prefix)
                    }
                })
                .map(|(_, &fid)| fid),
        );
    }

    ids
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::port::SiblingExpansion;
    use crate::SymbolKind;

    /// Mock LanguageRegistry that provides sibling expansion rules for testing.
    struct MockLangRegistry;

    impl crate::port::LanguageParser for MockLangRegistry {
        fn parse(&self, _: &str, _: &std::path::Path) -> crate::model::ParseResult {
            unimplemented!()
        }
        fn language_name(&self) -> &str {
            "mock"
        }
    }

    impl LanguageRegistry for MockLangRegistry {
        fn parser_for(&self, _: &str) -> Option<&dyn crate::port::LanguageParser> {
            None
        }
        fn resolver_for(
            &self,
            _: &str,
            _: &std::path::Path,
        ) -> Option<Box<dyn crate::port::LanguageResolver>> {
            None
        }
        fn supported_extensions(&self) -> &[&str] {
            &["rs", "py", "go"]
        }
        fn config_files(&self) -> &[&str] {
            &["tsconfig.json", "pyproject.toml", "go.mod"]
        }
        fn sibling_expansion(&self, rel_path: &str) -> Option<SiblingExpansion> {
            if rel_path.ends_with(".rs") {
                let stem = rel_path.strip_suffix(".rs")?;
                Some(SiblingExpansion {
                    prefix: format!("{}/", stem),
                    extension: String::new(),
                    root_only: false,
                })
            } else if rel_path.ends_with("__init__.py") {
                let dir = rel_path
                    .strip_suffix("__init__.py")
                    .filter(|p| !p.is_empty())?;
                Some(SiblingExpansion {
                    prefix: dir.to_string(),
                    extension: "py".to_string(),
                    root_only: false,
                })
            } else if rel_path.ends_with(".go") {
                match rel_path.rsplit_once('/') {
                    Some((dir, _)) => Some(SiblingExpansion {
                        prefix: format!("{}/", dir),
                        extension: "go".to_string(),
                        root_only: false,
                    }),
                    None => Some(SiblingExpansion {
                        prefix: String::new(),
                        extension: "go".to_string(),
                        root_only: true,
                    }),
                }
            } else {
                None
            }
        }
    }

    // --- Pure function tests ---

    #[test]
    fn test_merge_symbol_maps_first_wins() {
        let lists = vec![
            (
                1,
                vec![Symbol {
                    id: 10,
                    file_id: 1,
                    parent_id: None,
                    name: "Foo".to_string(),
                    kind: SymbolKind::Struct,
                    signature: None,
                    summary: None,
                    start_line: 1,
                    end_line: 5,
                    start_byte: 0,
                    end_byte: 100,
                }],
            ),
            (
                2,
                vec![Symbol {
                    id: 20,
                    file_id: 2,
                    parent_id: None,
                    name: "Foo".to_string(),
                    kind: SymbolKind::Struct,
                    signature: None,
                    summary: None,
                    start_line: 1,
                    end_line: 5,
                    start_byte: 0,
                    end_byte: 100,
                }],
            ),
        ];

        let result = merge_symbol_maps(&lists);
        assert_eq!(result.get("Foo"), Some(&10), "first occurrence should win");
    }

    #[test]
    fn test_merge_symbol_maps_empty() {
        let lists: Vec<(i64, Vec<Symbol>)> = vec![];
        let result = merge_symbol_maps(&lists);
        assert!(result.is_empty());
    }

    #[test]
    fn test_collect_candidate_files_with_submodules() {
        let mut file_path_to_id = HashMap::new();
        file_path_to_id.insert("src/scanner.rs".to_string(), 1);
        file_path_to_id.insert("src/scanner/parser.rs".to_string(), 2);
        file_path_to_id.insert("src/scanner/registry.rs".to_string(), 3);
        file_path_to_id.insert("src/other.rs".to_string(), 4);

        let file_id_to_path: HashMap<i64, String> = file_path_to_id
            .iter()
            .map(|(p, &id)| (id, p.clone()))
            .collect();

        let mut ids =
            collect_candidate_files(1, &file_path_to_id, &file_id_to_path, &MockLangRegistry);
        ids.sort();
        assert!(ids.contains(&1), "target file itself");
        assert!(ids.contains(&2), "submodule parser.rs");
        assert!(ids.contains(&3), "submodule registry.rs");
        assert!(!ids.contains(&4), "other.rs is not a submodule");
    }

    #[test]
    fn test_collect_candidate_files_no_submodules() {
        let mut file_path_to_id = HashMap::new();
        file_path_to_id.insert("src/main.rs".to_string(), 1);
        file_path_to_id.insert("src/other.rs".to_string(), 2);

        let file_id_to_path: HashMap<i64, String> = file_path_to_id
            .iter()
            .map(|(p, &id)| (id, p.clone()))
            .collect();

        let ids = collect_candidate_files(1, &file_path_to_id, &file_id_to_path, &MockLangRegistry);
        // main.rs has no submodule directory "src/main/"
        assert_eq!(ids, vec![1]);
    }

    // --- Additional coverage tests ---

    #[test]
    fn test_collect_candidate_files_unknown_fid() {
        // When target_fid is not in file_id_to_path, should return just the fid
        let file_path_to_id = HashMap::new();
        let file_id_to_path = HashMap::new();

        let ids =
            collect_candidate_files(999, &file_path_to_id, &file_id_to_path, &MockLangRegistry);
        assert_eq!(ids, vec![999], "unknown fid should return just itself");
    }

    #[test]
    fn test_collect_candidate_files_python_package() {
        // Python package: pkg/__init__.py should also search pkg/*.py siblings
        let mut file_path_to_id = HashMap::new();
        file_path_to_id.insert("pkg/__init__.py".to_string(), 1);
        file_path_to_id.insert("pkg/utils.py".to_string(), 2);
        file_path_to_id.insert("pkg/models.py".to_string(), 3);
        file_path_to_id.insert("other/stuff.py".to_string(), 4);

        let file_id_to_path: HashMap<i64, String> = file_path_to_id
            .iter()
            .map(|(p, &id)| (id, p.clone()))
            .collect();

        let mut ids =
            collect_candidate_files(1, &file_path_to_id, &file_id_to_path, &MockLangRegistry);
        ids.sort();
        assert!(ids.contains(&1), "target file itself");
        assert!(ids.contains(&2), "pkg/utils.py sibling");
        assert!(ids.contains(&3), "pkg/models.py sibling");
        assert!(!ids.contains(&4), "other/stuff.py is not in pkg/");
    }

    #[test]
    fn test_collect_candidate_files_python_no_double_count() {
        // __init__.py should not appear twice (once as target, once as sibling)
        let mut file_path_to_id = HashMap::new();
        file_path_to_id.insert("pkg/__init__.py".to_string(), 1);
        file_path_to_id.insert("pkg/helper.py".to_string(), 2);

        let file_id_to_path: HashMap<i64, String> = file_path_to_id
            .iter()
            .map(|(p, &id)| (id, p.clone()))
            .collect();

        let ids = collect_candidate_files(1, &file_path_to_id, &file_id_to_path, &MockLangRegistry);
        let count_of_1 = ids.iter().filter(|&&id| id == 1).count();
        assert_eq!(
            count_of_1, 1,
            "__init__.py should appear only once, ids: {:?}",
            ids
        );
    }

    #[test]
    fn test_collect_candidate_files_go_package() {
        // Go: all .go files in same directory are the same package
        let mut file_path_to_id = HashMap::new();
        file_path_to_id.insert("pkg/types.go".to_string(), 1);
        file_path_to_id.insert("pkg/service.go".to_string(), 2);
        file_path_to_id.insert("pkg/helpers.go".to_string(), 3);
        file_path_to_id.insert("other/main.go".to_string(), 4);

        let file_id_to_path: HashMap<i64, String> = file_path_to_id
            .iter()
            .map(|(p, &id)| (id, p.clone()))
            .collect();

        let mut ids =
            collect_candidate_files(1, &file_path_to_id, &file_id_to_path, &MockLangRegistry);
        ids.sort();
        assert!(ids.contains(&1), "target file itself");
        assert!(ids.contains(&2), "pkg/service.go sibling");
        assert!(ids.contains(&3), "pkg/helpers.go sibling");
        assert!(!ids.contains(&4), "other/main.go is different package");
    }

    #[test]
    fn test_collect_candidate_files_go_root_level() {
        // Go: root-level .go files (no directory) should also expand
        let mut file_path_to_id = HashMap::new();
        file_path_to_id.insert("main.go".to_string(), 1);
        file_path_to_id.insert("helpers.go".to_string(), 2);
        file_path_to_id.insert("utils.go".to_string(), 3);
        file_path_to_id.insert("pkg/other.go".to_string(), 4);

        let file_id_to_path: HashMap<i64, String> = file_path_to_id
            .iter()
            .map(|(p, &id)| (id, p.clone()))
            .collect();

        let mut ids =
            collect_candidate_files(1, &file_path_to_id, &file_id_to_path, &MockLangRegistry);
        ids.sort();
        assert!(ids.contains(&1), "target file itself");
        assert!(ids.contains(&2), "helpers.go root sibling");
        assert!(ids.contains(&3), "utils.go root sibling");
        assert!(!ids.contains(&4), "pkg/other.go is in subdir");
    }

    #[test]
    fn test_merge_symbol_maps_multiple_files() {
        let lists = vec![
            (
                1,
                vec![
                    Symbol {
                        id: 10,
                        file_id: 1,
                        parent_id: None,
                        name: "Alpha".to_string(),
                        kind: SymbolKind::Function,
                        signature: None,
                        summary: None,
                        start_line: 1,
                        end_line: 5,
                        start_byte: 0,
                        end_byte: 100,
                    },
                    Symbol {
                        id: 11,
                        file_id: 1,
                        parent_id: None,
                        name: "Beta".to_string(),
                        kind: SymbolKind::Function,
                        signature: None,
                        summary: None,
                        start_line: 6,
                        end_line: 10,
                        start_byte: 101,
                        end_byte: 200,
                    },
                ],
            ),
            (
                2,
                vec![Symbol {
                    id: 20,
                    file_id: 2,
                    parent_id: None,
                    name: "Gamma".to_string(),
                    kind: SymbolKind::Struct,
                    signature: None,
                    summary: None,
                    start_line: 1,
                    end_line: 3,
                    start_byte: 0,
                    end_byte: 50,
                }],
            ),
        ];

        let result = merge_symbol_maps(&lists);
        assert_eq!(result.len(), 3);
        assert_eq!(result.get("Alpha"), Some(&10));
        assert_eq!(result.get("Beta"), Some(&11));
        assert_eq!(result.get("Gamma"), Some(&20));
    }
}