brokk_bifrost_python/

imports.rs

//! Python import syntax, binding, and FQN resolution.
//!
//! Everything here is a free function over [`PythonSource`]; the
//! `ImportAnalysisProvider` impl that memoizes the per-file results stays on
//! `PythonAnalyzer` in `analyzer/python/imports.rs`.

use crate::bindings::python_direct_scope_bindings_bounded;
use crate::declarations::{parse_python_tree, py_node_text, python_module_name};
use crate::graph_support::{
    PythonSource, import_binder_from_imports, public_declarations_in_module,
    resolve_module_code_unit, resolve_module_code_units_batch,
};
use brokk_bifrost_core::analyzer::common::node_source_text;
use brokk_bifrost_core::analyzer::model::{
    ImportInfo, StructuredImportPath, StructuredImportPathKind,
};
use brokk_bifrost_core::analyzer::query_token::QueryToken;
use brokk_bifrost_core::analyzer::usages::model::{ExportEntry, ImportBinding, ImportKind};
use brokk_bifrost_core::analyzer::{CodeUnit, CodeUnitIndex, ProjectFile};
use brokk_bifrost_core::hash::{HashMap, HashSet};
use std::collections::VecDeque;
use tree_sitter::Node;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PythonModuleReplacement {
    pub target_module: String,
}

/// Return the absolute module named by the last direct assignment to `local`
/// that dominates `reference` when that assignment has the exact structured
/// shape `local = importlib_binding.import_module("module.name")`.
///
/// Assignments hidden in a branch do not dominate the reference. Any later
/// direct binding of `local` invalidates an earlier result. The caller proves
/// that the receiver of `import_module` is an in-scope namespace binding for
/// Python's `importlib` module; this helper only interprets the syntax and
/// statement order.
pub(crate) fn imported_module_assignment_at(
    reference: Node<'_>,
    local: &str,
    source: &str,
    mut is_importlib_binding: impl FnMut(&str) -> bool,
) -> Option<String> {
    let suite = enclosing_execution_suite(reference)?;
    let mut resolved = None;
    let mut cursor = suite.walk();
    for statement in suite.named_children(&mut cursor) {
        if statement.start_byte() >= reference.start_byte()
            || (statement.start_byte() <= reference.start_byte()
                && reference.end_byte() <= statement.end_byte())
        {
            break;
        }

        let expression = if statement.kind() == "expression_statement" {
            statement.named_child(0).unwrap_or(statement)
        } else {
            statement
        };
        if expression.kind() == "assignment"
            && expression.child_by_field_name("left").is_some_and(|left| {
                left.kind() == "identifier" && node_source_text(left, source) == local
            })
        {
            resolved = expression
                .child_by_field_name("right")
                .and_then(|right| importlib_call_module(right, source, &mut is_importlib_binding));
            continue;
        }

        let binds_local = python_direct_scope_bindings_bounded(statement, source, || true)
            .expect("unbounded binding collection cannot be cancelled")
            .into_iter()
            .any(|binding| node_source_text(binding.declaration, source) == local);
        if binds_local {
            resolved = None;
        }
    }
    resolved
}

fn enclosing_execution_suite(mut node: Node<'_>) -> Option<Node<'_>> {
    while let Some(parent) = node.parent() {
        if matches!(parent.kind(), "function_definition" | "lambda") {
            return parent.child_by_field_name("body");
        }
        if parent.kind() == "module" {
            return Some(parent);
        }
        node = parent;
    }
    None
}

fn importlib_call_module(
    call: Node<'_>,
    source: &str,
    is_importlib_binding: &mut impl FnMut(&str) -> bool,
) -> Option<String> {
    if call.kind() != "call" {
        return None;
    }
    let function = call.child_by_field_name("function")?;
    if function.kind() != "attribute" {
        return None;
    }
    let (object, attribute) = (
        function.child_by_field_name("object")?,
        function.child_by_field_name("attribute")?,
    );
    if object.kind() != "identifier"
        || attribute.kind() != "identifier"
        || node_source_text(attribute, source) != "import_module"
        || !is_importlib_binding(node_source_text(object, source))
    {
        return None;
    }

    let arguments = call.child_by_field_name("arguments")?;
    if arguments.named_child_count() != 1 {
        return None;
    }
    let literal = arguments.named_child(0)?;
    if literal.kind() != "string"
        || literal
            .parent()
            .is_some_and(|parent| parent.kind() == "concatenated_string")
    {
        return None;
    }
    let mut content = None;
    for index in 0..literal.named_child_count() {
        let child = literal.named_child(index)?;
        match child.kind() {
            "string_start" | "string_end" => {}
            "string_content" if content.is_none() && child.named_child_count() == 0 => {
                content = Some(child);
            }
            _ => return None,
        }
    }
    let content = content?;
    let module = node_source_text(content, source);
    (!module.is_empty() && !module.starts_with('.')).then(|| module.to_string())
}

/// Collect absolute module names from structurally exact
/// `local = importlib.import_module("...")` assignments. This is a conservative
/// candidate-file index: dominance and local shadowing are checked later at the
/// reference site by [`imported_module_assignment_at`].
pub(crate) fn literal_importlib_modules(
    source: &str,
    bindings: &HashMap<String, ImportBinding>,
) -> HashSet<String> {
    let Some(tree) = parse_python_tree(source) else {
        return HashSet::default();
    };
    let mut modules = HashSet::default();
    let mut stack = vec![tree.root_node()];
    while let Some(node) = stack.pop() {
        if node.kind() == "assignment"
            && node
                .child_by_field_name("left")
                .is_some_and(|left| left.kind() == "identifier")
            && let Some(module) = node.child_by_field_name("right").and_then(|right| {
                importlib_call_module(right, source, &mut |local| {
                    bindings.get(local).is_some_and(|binding| {
                        binding.kind == ImportKind::Namespace
                            && binding
                                .namespace_imported_module
                                .as_deref()
                                .unwrap_or(&binding.module_specifier)
                                == "importlib"
                    })
                })
            })
        {
            modules.insert(module);
        }
        let mut cursor = node.walk();
        stack.extend(node.named_children(&mut cursor));
    }
    modules
}

/// Recognize the compatibility-shim idiom that replaces the current module
/// object with an imported workspace module:
///
/// `sys.modules[__name__] = imported_module`
///
/// Every component is resolved from tree-sitter fields and the import binder;
/// similarly spelled attributes or locals do not create an alias edge.
fn module_replacement_from_assignment(
    assignment: Node<'_>,
    source: &str,
    bindings: &HashMap<String, ImportBinding>,
) -> Option<PythonModuleReplacement> {
    let (left, right) = (
        assignment.child_by_field_name("left")?,
        assignment.child_by_field_name("right")?,
    );
    if left.kind() != "subscript" || right.kind() != "identifier" {
        return None;
    }
    let (value, subscript) = (
        left.child_by_field_name("value")?,
        left.child_by_field_name("subscript")?,
    );
    if value.kind() != "attribute"
        || subscript.kind() != "identifier"
        || node_source_text(subscript, source) != "__name__"
    {
        return None;
    }
    let (sys_local, modules) = (
        value.child_by_field_name("object")?,
        value.child_by_field_name("attribute")?,
    );
    if sys_local.kind() != "identifier"
        || modules.kind() != "identifier"
        || node_source_text(modules, source) != "modules"
    {
        return None;
    }
    let sys_binding = bindings.get(node_source_text(sys_local, source))?;
    let sys_module = sys_binding
        .namespace_imported_module
        .as_deref()
        .unwrap_or(&sys_binding.module_specifier);
    if sys_binding.kind != ImportKind::Namespace || sys_module != "sys" {
        return None;
    }

    let target_binding = bindings.get(node_source_text(right, source))?;
    if target_binding.kind != ImportKind::Namespace {
        return None;
    }
    Some(PythonModuleReplacement {
        target_module: target_binding
            .namespace_imported_module
            .clone()
            .unwrap_or_else(|| target_binding.module_specifier.clone()),
    })
}

fn remove_direct_scope_bindings(
    statement: Node<'_>,
    source: &str,
    bindings: &mut HashMap<String, ImportBinding>,
) {
    let mut stack = vec![statement];
    while let Some(node) = stack.pop() {
        for binding in python_direct_scope_bindings_bounded(node, source, || true)
            .expect("unbounded binding collection cannot be cancelled")
        {
            bindings.remove(node_source_text(binding.declaration, source));
        }
        if matches!(
            node.kind(),
            "function_definition" | "class_definition" | "lambda"
        ) {
            continue;
        }
        let mut cursor = node.walk();
        stack.extend(node.named_children(&mut cursor));
    }
}

/// Parse the structured import facts for a Python document without executing
/// it. LSP model binding uses the same AST-derived representation as the
/// analyzer so external APIs are never selected by a terminal-name scan.
pub fn parse_python_import_infos(source: &str) -> Vec<ImportInfo> {
    let mut parser = tree_sitter::Parser::new();
    parser
        .set_language(&tree_sitter_python::LANGUAGE.into())
        .expect("failed to load Python parser");
    let Some(tree) = parser.parse(source, None) else {
        return Vec::new();
    };
    let mut pending = vec![tree.root_node()];
    let mut imports = Vec::new();
    while let Some(node) = pending.pop() {
        if matches!(node.kind(), "import_statement" | "import_from_statement") {
            imports.extend(python_import_infos_from_node(node, source));
            continue;
        }
        let mut cursor = node.walk();
        pending.extend(node.named_children(&mut cursor));
    }
    imports
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PythonImportBinding {
    pub start_byte: usize,
    pub scope_start_byte: usize,
    pub scope_end_byte: usize,
    function_scoped: bool,
    pub local_name: String,
    pub qualified_name: String,
    pub consumed_attributes: usize,
}

impl PythonImportBinding {
    pub fn is_function_scoped(&self) -> bool {
        self.function_scoped
    }
}

/// Return source-ordered, parser-derived local bindings for explicit Python
/// imports. The byte offset lets callers select the last visible binding rather
/// than treating a whole document as one unordered import scope.
pub fn parse_python_import_bindings(source: &str) -> Vec<PythonImportBinding> {
    let mut parser = tree_sitter::Parser::new();
    parser
        .set_language(&tree_sitter_python::LANGUAGE.into())
        .expect("failed to load Python parser");
    let Some(tree) = parser.parse(source, None) else {
        return Vec::new();
    };
    python_import_bindings_from_tree(tree.root_node(), source)
}

/// Return structured import bindings from a tree the caller already owns.
/// Query paths use this form so one request never reparses a document merely
/// to recover function-local import scope.
pub fn python_import_bindings_from_tree(root: Node<'_>, source: &str) -> Vec<PythonImportBinding> {
    let mut pending = vec![root];
    let mut nodes = Vec::new();
    while let Some(node) = pending.pop() {
        if matches!(node.kind(), "import_statement" | "import_from_statement") {
            nodes.push(node);
            continue;
        }
        let mut cursor = node.walk();
        pending.extend(node.named_children(&mut cursor));
    }
    nodes.sort_by_key(Node::start_byte);
    nodes
        .into_iter()
        .flat_map(|node| {
            let (scope_start_byte, scope_end_byte, function_scoped) =
                python_import_binding_scope(node, source.len());
            python_import_infos_from_node(node, source)
                .into_iter()
                .filter_map(move |import| {
                    let path = import.path.as_ref()?;
                    let details = python_import_details(&import)?;
                    match details {
                        PythonImportDetails::Import { module, alias } => {
                            let consumed_attributes = if alias.is_some() {
                                0
                            } else {
                                path.segments.len().saturating_sub(1)
                            };
                            Some(PythonImportBinding {
                                start_byte: node.start_byte(),
                                scope_start_byte,
                                scope_end_byte,
                                function_scoped,
                                local_name: alias.or_else(|| path.segments.first().cloned())?,
                                qualified_name: module,
                                consumed_attributes,
                            })
                        }
                        PythonImportDetails::FromImport {
                            module,
                            name,
                            alias,
                            wildcard: false,
                        } => Some(PythonImportBinding {
                            start_byte: node.start_byte(),
                            scope_start_byte,
                            scope_end_byte,
                            function_scoped,
                            local_name: alias.unwrap_or(name.clone()),
                            qualified_name: format!("{module}.{name}"),
                            consumed_attributes: 0,
                        }),
                        PythonImportDetails::FromImport { wildcard: true, .. } => None,
                    }
                })
        })
        .collect()
}

/// Build import bindings from the canonical generic import projection.
///
/// The producer has already interpreted each path and retained lexical scope
/// byte ranges in the import facts. This adapter only maps those ranges back
/// to the caller's prepared tree to recover the nearest function or lambda
/// scope; it never reparses source or reconstructs import syntax.
pub fn python_import_bindings_from_imports(
    imports: &[ImportInfo],
    root: Node<'_>,
    source_len: usize,
) -> Vec<PythonImportBinding> {
    let mut bindings = imports
        .iter()
        .filter_map(|import| {
            let path = import.path.as_ref()?;
            let details = python_import_details(import)?;
            let mut binding_scope = None;
            for scope in path.lexical_scopes.iter().rev() {
                let node = root.descendant_for_byte_range(scope.start_byte, scope.end_byte)?;
                if node.start_byte() != scope.start_byte || node.end_byte() != scope.end_byte {
                    return None;
                }
                if matches!(node.kind(), "function_definition" | "lambda") {
                    binding_scope = Some((scope.start_byte, scope.end_byte, true));
                    break;
                }
            }
            let (scope_start_byte, scope_end_byte, function_scoped) =
                binding_scope.unwrap_or((0, source_len, false));
            let start_byte = path.declaration_start_byte;
            match details {
                PythonImportDetails::Import { module, alias } => {
                    let consumed_attributes = if alias.is_some() {
                        0
                    } else {
                        path.segments.len().saturating_sub(1)
                    };
                    Some(PythonImportBinding {
                        start_byte,
                        scope_start_byte,
                        scope_end_byte,
                        function_scoped,
                        local_name: alias.or_else(|| path.segments.first().cloned())?,
                        qualified_name: module,
                        consumed_attributes,
                    })
                }
                PythonImportDetails::FromImport {
                    module,
                    name,
                    alias,
                    wildcard: false,
                } => Some(PythonImportBinding {
                    start_byte,
                    scope_start_byte,
                    scope_end_byte,
                    function_scoped,
                    local_name: alias.unwrap_or_else(|| name.clone()),
                    qualified_name: format!("{module}.{name}"),
                    consumed_attributes: 0,
                }),
                PythonImportDetails::FromImport { wildcard: true, .. } => None,
            }
        })
        .collect::<Vec<_>>();
    bindings.sort_by_key(|binding| binding.start_byte);
    bindings
}

fn python_import_binding_scope(node: Node<'_>, source_len: usize) -> (usize, usize, bool) {
    let mut parent = node.parent();
    while let Some(scope) = parent {
        if matches!(scope.kind(), "function_definition" | "lambda") {
            return (scope.start_byte(), scope.end_byte(), true);
        }
        parent = scope.parent();
    }
    (0, source_len, false)
}

#[cfg(test)]
mod tests {
    use super::*;
    use brokk_bifrost_core::analyzer::usages::model::ImportBinder;

    fn importlib_assignment_for(source: &str, importlib_local: &str) -> Option<String> {
        let tree = parse_python_tree(source).expect("valid Python fixture");
        let reference_start = source.rfind("canonical.setup").expect("reference site");
        let mut stack = vec![tree.root_node()];
        let mut reference = None;
        while let Some(node) = stack.pop() {
            if node.kind() == "identifier"
                && node.start_byte() == reference_start
                && node_source_text(node, source) == "canonical"
            {
                reference = Some(node);
                break;
            }
            let mut cursor = node.walk();
            stack.extend(node.named_children(&mut cursor));
        }
        imported_module_assignment_at(
            reference.expect("canonical reference node"),
            "canonical",
            source,
            |local| local == importlib_local,
        )
    }

    #[test]
    fn existing_tree_import_bindings_preserve_function_scope_and_relative_identity() {
        let source = r#"def register():
    from .workers.special import Worker
    return Worker()
"#;
        let tree = parse_python_tree(source).expect("Python tree");
        let bindings = python_import_bindings_from_tree(tree.root_node(), source);

        assert_eq!(bindings.len(), 1, "{bindings:#?}");
        assert!(bindings[0].is_function_scoped(), "{bindings:#?}");
        assert_eq!(bindings[0].local_name, "Worker");
        assert_eq!(bindings[0].qualified_name, ".workers.special.Worker");
        assert!(bindings[0].scope_start_byte <= bindings[0].start_byte);
        assert!(bindings[0].scope_end_byte >= source.rfind("Worker()").unwrap());
    }

    #[test]
    fn importlib_assignment_requires_a_direct_literal_dominating_write() {
        let valid = r#"import importlib as loader

def test():
    canonical = loader.import_module("routes.contacts.contacts_routes")
    assert canonical.setup_contacts_routes
"#;
        assert_eq!(
            importlib_assignment_for(valid, "loader"),
            Some("routes.contacts.contacts_routes".to_string())
        );

        for near_miss in [
            r#"import importlib as loader
def test(module_name):
    canonical = loader.import_module(module_name)
    assert canonical.setup_contacts_routes
"#,
            r#"import importlib as loader
def test():
    if enabled:
        canonical = loader.import_module("routes.contacts.contacts_routes")
    assert canonical.setup_contacts_routes
"#,
            r#"import importlib as loader
def test():
    canonical = loader.import_module("routes.contacts.contacts_routes")
    canonical = object()
    assert canonical.setup_contacts_routes
"#,
        ] {
            assert_eq!(
                importlib_assignment_for(near_miss, "loader"),
                None,
                "near miss must not infer a module receiver: {near_miss:?}"
            );
        }
        assert_eq!(
            importlib_assignment_for(valid, "shadowed_loader"),
            None,
            "a similarly shaped non-importlib receiver must not resolve"
        );
    }

    fn replacement_for(source: &str, binder: &ImportBinder) -> Option<PythonModuleReplacement> {
        let tree = parse_python_tree(source).expect("valid Python fixture");
        let root = tree.root_node();
        let mut replacement = None;
        let mut cursor = root.walk();
        for statement in root.named_children(&mut cursor) {
            let statement = if statement.kind() == "expression_statement" {
                statement.named_child(0).expect("fixture expression")
            } else {
                statement
            };
            if statement.kind() == "assignment"
                && let Some(next) =
                    module_replacement_from_assignment(statement, source, &binder.bindings)
                && replacement.replace(next).is_some()
            {
                return None;
            }
        }
        replacement
    }

    #[test]
    fn module_replacement_requires_exact_structured_import_bindings() {
        let source = r#"import sys as _sys
from routes.contacts import contacts_routes as _canonical

_sys.modules[__name__] = _canonical
"#;
        let mut binder = ImportBinder::empty();
        binder.bindings.insert(
            "_sys".to_string(),
            ImportBinding {
                module_specifier: "sys".to_string(),
                namespace_imported_module: Some("sys".to_string()),
                kind: ImportKind::Namespace,
                imported_name: None,
            },
        );
        binder.bindings.insert(
            "_canonical".to_string(),
            ImportBinding {
                module_specifier: "routes.contacts.contacts_routes".to_string(),
                namespace_imported_module: None,
                kind: ImportKind::Namespace,
                imported_name: None,
            },
        );

        assert_eq!(
            replacement_for(source, &binder),
            Some(PythonModuleReplacement {
                target_module: "routes.contacts.contacts_routes".to_string(),
            })
        );

        for near_miss in [
            "cache.modules[__name__] = _canonical\n",
            "_sys.modules[module_name] = _canonical\n",
            "_sys.modules[__name__] = build()\n",
            "def replace():\n    _sys.modules[__name__] = _canonical\n",
        ] {
            assert_eq!(
                replacement_for(near_miss, &binder),
                None,
                "near miss must not replace module identity: {near_miss:?}"
            );
        }
    }
}

/// Borrow each statement's contiguous canonical leaves without reinterpreting
/// target or alias syntax in executable binding timelines.
pub(crate) fn python_imports_by_statement(imports: &[ImportInfo]) -> HashMap<usize, &[ImportInfo]> {
    let mut statement_imports = HashMap::default();
    for group in imports.chunk_by(|left, right| {
        left.path.as_ref().map(|path| path.declaration_start_byte)
            == right.path.as_ref().map(|path| path.declaration_start_byte)
    }) {
        if let Some(path) = &group[0].path {
            assert!(
                statement_imports
                    .insert(path.declaration_start_byte, group)
                    .is_none(),
                "canonical Python import leaves must remain in statement order"
            );
        }
    }
    statement_imports
}

pub fn module_replacement_of(
    python: &dyn PythonSource,
    file: &ProjectFile,
    source: &str,
    imports: &[ImportInfo],
) -> Option<PythonModuleReplacement> {
    let tree = parse_python_tree(source)?;
    let root = tree.root_node();
    let mut bindings: HashMap<String, ImportBinding> = HashMap::default();
    let statement_imports = python_imports_by_statement(imports);
    let mut replacement = None;
    let mut cursor = root.walk();
    for statement in root.named_children(&mut cursor) {
        let statement = if statement.kind() == "expression_statement" {
            let Some(expression) = statement.named_child(0) else {
                continue;
            };
            expression
        } else {
            statement
        };
        match statement.kind() {
            "import_statement" | "import_from_statement" => {
                if let Some(imports) = statement_imports.get(&statement.start_byte()) {
                    bindings.extend(import_binder_from_imports(python, file, imports).bindings);
                }
            }
            "assignment" => {
                if let Some(next) = module_replacement_from_assignment(statement, source, &bindings)
                    && replacement.replace(next).is_some()
                {
                    return None;
                }
                remove_direct_scope_bindings(statement, source, &mut bindings);
            }
            _ => remove_direct_scope_bindings(statement, source, &mut bindings),
        }
    }
    replacement
}

pub fn resolve_import_bindings(
    python: &dyn PythonSource,
    token: QueryToken<'_>,
    file: &ProjectFile,
) -> HashMap<String, CodeUnit> {
    let imports = python.import_info_of(token, file);
    let mut bindings = HashMap::default();
    for resolved in resolve_imports_batched(python, file, &imports) {
        for (binding, code_unit) in resolved {
            bindings.insert(binding, code_unit);
        }
    }
    bindings
}

/// Resolves every import in `imports` (`file`'s own imports), batching each import's primary
/// module FQN lookup (see `primary_module_fqn`) into one store transaction instead of one per
/// import. Shared by `resolve_import_bindings` and `resolve_import_target_files`, the two per-file
/// "resolve everything" entry points -- both are called once per candidate file by the usages
/// candidate walker, so unbatched resolution here means one store transaction per import times
/// every file in the workspace.
pub fn resolve_imports_batched(
    python: &dyn PythonSource,
    file: &ProjectFile,
    imports: &[ImportInfo],
) -> Vec<Vec<(String, CodeUnit)>> {
    let primary_fqns: Vec<Option<String>> = imports
        .iter()
        .map(|import| primary_module_fqn(file, import))
        .collect();
    let to_resolve: Vec<String> = primary_fqns.iter().flatten().cloned().collect();
    let mut batch_results =
        resolve_module_code_units_batch(python, Some(file), &to_resolve).into_iter();

    imports
        .iter()
        .zip(primary_fqns.iter())
        .map(|(import, primary_fqn)| {
            let hint = primary_fqn.as_ref().map(|_| batch_results.next().unwrap());
            resolve_import_with_hint(python, file, import, hint.as_ref())
        })
        .collect()
}

/// Resolve Python imports to the workspace files whose modules they execute.
///
/// A file dependency graph does not need declaration binding: for
/// `from pkg import member`, the dependency on `pkg` exists whether `member`
/// is a local declaration, a re-export, or an optional external name. When
/// `member` is itself the workspace submodule `pkg.member`, that file is a
/// dependency too. Resolve both module spellings from the path-backed module
/// index in one batch and deliberately avoid the exact-definition fallbacks
/// used by [`resolve_imports_batched`].
pub fn resolve_import_files_batched(
    python: &dyn PythonSource,
    file: &ProjectFile,
    imports: &[ImportInfo],
) -> HashSet<ProjectFile> {
    let mut module_fqns = Vec::new();
    for import in imports {
        match python_import_details(import) {
            Some(PythonImportDetails::Import { module, .. }) => module_fqns.push(module),
            Some(PythonImportDetails::FromImport {
                module,
                name,
                wildcard,
                ..
            }) => {
                let resolved_module = if module.starts_with('.') {
                    resolve_python_relative_module(file, &module)
                } else {
                    Some(module)
                };
                let Some(resolved_module) = resolved_module else {
                    continue;
                };
                if !wildcard {
                    module_fqns.push(format!("{resolved_module}.{name}"));
                }
                module_fqns.push(resolved_module);
            }
            None => {}
        }
    }
    module_fqns.sort();
    module_fqns.dedup();

    python
        .path_module_fqns_batch(&module_fqns)
        .into_iter()
        .flatten()
        .flatten()
        .filter(CodeUnit::is_module)
        .map(|unit| unit.source().clone())
        .collect()
}

/// The module FQN `resolve_import`'s fast path checks first, if any -- must stay in sync with the
/// two `resolve_module_code_unit` call sites in `resolve_import_with_hint` below, since it's what
/// lets `resolve_import_target_files` batch-resolve them ahead of the serial fallback logic.
fn primary_module_fqn(file: &ProjectFile, import: &ImportInfo) -> Option<String> {
    match python_import_details(import)? {
        PythonImportDetails::Import { module, alias } => Some(python_namespace_binding_module(
            import,
            alias.as_deref(),
            &module,
        )),
        PythonImportDetails::FromImport {
            module,
            name,
            wildcard,
            ..
        } => {
            if wildcard {
                return None;
            }
            let resolved_module = if module.starts_with('.') {
                resolve_python_relative_module(file, &module)
            } else {
                Some(module)
            };
            resolved_module.map(|resolved_module| format!("{resolved_module}.{name}"))
        }
    }
}

pub fn resolve_import(
    python: &dyn PythonSource,
    file: &ProjectFile,
    import: &ImportInfo,
) -> Vec<(String, CodeUnit)> {
    resolve_import_with_hint(python, file, import, None)
}

/// `primary_hint`, when `Some`, is the already-resolved result of this import's primary module FQN
/// (see `primary_module_fqn`) so the batched caller doesn't pay for a second lookup of the same FQN.
fn resolve_import_with_hint(
    python: &dyn PythonSource,
    file: &ProjectFile,
    import: &ImportInfo,
    primary_hint: Option<&Option<CodeUnit>>,
) -> Vec<(String, CodeUnit)> {
    if let Some(details) = python_import_details(import) {
        match details {
            PythonImportDetails::Import { module, alias } => {
                let binding = python_namespace_binding_name(import, alias.as_deref(), &module);
                let bound_module =
                    python_namespace_binding_module(import, alias.as_deref(), &module);
                let resolved = match primary_hint {
                    Some(hint) => hint.clone(),
                    None => resolve_module_code_unit(python, Some(file), &bound_module),
                };
                if let Some(module_code_unit) = resolved {
                    return vec![(binding, module_code_unit)];
                }
            }
            PythonImportDetails::FromImport {
                module,
                name,
                alias,
                wildcard,
            } => {
                let resolved_module = if module.starts_with('.') {
                    resolve_python_relative_module(file, &module)
                } else {
                    Some(module)
                };
                let Some(resolved_module) = resolved_module else {
                    return Vec::new();
                };
                if wildcard {
                    return public_declarations_in_module(python, file, &resolved_module)
                        .into_iter()
                        .map(|code_unit| (code_unit.identifier().to_string(), code_unit))
                        .collect();
                }

                let binding = alias.clone().unwrap_or_else(|| name.clone());
                let module_candidate = format!("{resolved_module}.{name}");
                let resolved = match primary_hint {
                    Some(hint) => hint.clone(),
                    None => resolve_module_code_unit(python, Some(file), &module_candidate),
                };
                if let Some(code_unit) = resolved {
                    return vec![(binding, code_unit)];
                }
                let exported =
                    resolve_exported_name_from_module(python, Some(file), &resolved_module, &name);
                if !exported.is_empty() {
                    return exported
                        .into_iter()
                        .map(|code_unit| (binding.clone(), code_unit))
                        .collect();
                }
                let definitions: Vec<_> = python.definitions(&module_candidate).collect();
                if !definitions.is_empty() {
                    return definitions
                        .into_iter()
                        .map(|code_unit| (binding.clone(), code_unit))
                        .collect();
                }
                let package_candidate: Vec<_> = python
                    .definitions(&format!("{resolved_module}.{name}"))
                    .collect();
                if !package_candidate.is_empty() {
                    return package_candidate
                        .into_iter()
                        .map(|code_unit| (binding.clone(), code_unit))
                        .collect();
                }
            }
        }
    }
    Vec::new()
}

/// `importer`, when known, is the file whose import statement asked for `fqn`;
/// its own import root is what decides between two same-named modules. `None`
/// is for an FQN-level question no import statement asked.
pub fn resolve_exported_fqn(
    python: &dyn PythonSource,
    importer: Option<&ProjectFile>,
    fqn: &str,
) -> Vec<CodeUnit> {
    let Some((module, name)) = fqn.rsplit_once('.') else {
        return Vec::new();
    };
    resolve_exported_name_from_module(python, importer, module, name)
}

/// Resolve an unambiguous chain of explicit named reexports without
/// constructing export indexes for each intermediate module. Star exports,
/// shadowing, and every other ambiguous shape return `None` so callers can
/// use the complete, source-order-aware export resolver below.
///
/// Each hop carries its own importer: the first is the file that asked, and
/// every later one is the module whose reexport named the next specifier, so
/// a facade chain stays inside the roots that actually wrote it.
fn resolve_direct_named_exported_fqn(
    python: &dyn PythonSource,
    importer: Option<&ProjectFile>,
    fqn: &str,
) -> Option<Vec<CodeUnit>> {
    let (module, name) = fqn.rsplit_once('.')?;
    let mut results = Vec::new();
    let mut queue = VecDeque::from([(importer.cloned(), module.to_string(), name.to_string())]);
    let mut visited = HashSet::default();

    while let Some((importer, module, export_name)) = queue.pop_front() {
        if !visited.insert((module.clone(), export_name.clone())) {
            continue;
        }
        let module_unit = resolve_module_code_unit(python, importer.as_ref(), &module)?;
        let file = module_unit.source();
        let local = local_export_declarations(python, file, &export_name);
        let binder = python.import_binder_of(file);
        let binding = binder.bindings.get(&export_name);
        if !local.is_empty() && binding.is_some() {
            return None;
        }
        if !local.is_empty() {
            results.extend(local);
            continue;
        }
        let binding = binding?;
        if binding.kind != ImportKind::Named {
            return None;
        }
        let imported_name = binding.imported_name.as_ref()?;
        queue.push_back((
            Some(file.clone()),
            binding.module_specifier.clone(),
            imported_name.clone(),
        ));
    }

    results.sort_by(|left, right| {
        left.source()
            .cmp(right.source())
            .then_with(|| left.fq_name().cmp(&right.fq_name()))
    });
    results.dedup();
    (!results.is_empty()).then_some(results)
}

/// Resolve a Python FQN with the cheapest semantically complete tier that
/// can answer it. The direct reexport walk handles only proven,
/// collision-free chains; ambiguous shapes use the ordered export index,
/// and the exact lookup remains the final fallback for non-export symbols.
///
/// The first two tiers resolve a module specifier, so they take `importer`
/// for the same reason the import resolver does. The exact lookup is a
/// workspace-wide question about a dotted name and stays root-blind.
pub fn resolve_fqn_candidates(
    python: &dyn PythonSource,
    importer: Option<&ProjectFile>,
    fqn: &str,
    exact: impl FnOnce(&str) -> Vec<CodeUnit>,
) -> Vec<CodeUnit> {
    if let Some(candidates) = resolve_direct_named_exported_fqn(python, importer, fqn) {
        return candidates;
    }
    let candidates = resolve_exported_fqn(python, importer, fqn);
    if !candidates.is_empty() {
        return candidates;
    }
    exact(fqn)
}

fn resolve_exported_name_from_module(
    python: &dyn PythonSource,
    importer: Option<&ProjectFile>,
    module: &str,
    name: &str,
) -> Vec<CodeUnit> {
    let Some(module_unit) = resolve_module_code_unit(python, importer, module) else {
        return Vec::new();
    };
    resolve_exported_name(python, module_unit.source(), name)
}

fn resolve_exported_name(
    python: &dyn PythonSource,
    module_file: &ProjectFile,
    name: &str,
) -> Vec<CodeUnit> {
    let mut results = Vec::new();
    let mut queue = VecDeque::from([(module_file.clone(), name.to_string())]);
    let mut visited = HashSet::default();

    while let Some((file, export_name)) = queue.pop_front() {
        if !visited.insert((file.clone(), export_name.clone())) {
            continue;
        }

        let index = python.export_index_of(&file);
        if let Some(entry) = index.exports_by_name.get(&export_name) {
            match entry {
                ExportEntry::Local { local_name } => {
                    results.extend(local_export_declarations(python, &file, local_name));
                }
                ExportEntry::ReexportedNamed {
                    module_specifier,
                    imported_name,
                } => {
                    for target_file in
                        resolve_module_files_for_export(python, &file, module_specifier)
                    {
                        queue.push_back((target_file, imported_name.clone()));
                    }
                }
                ExportEntry::ReexportedModule { module_specifier } => {
                    // Terminal: the export *is* the module, so the walk stops
                    // here instead of looking the name up inside it.
                    results.extend(resolve_module_code_unit(
                        python,
                        Some(&file),
                        module_specifier,
                    ));
                }
                ExportEntry::Default { local_name } => {
                    if let Some(local_name) = local_name {
                        results.extend(local_export_declarations(python, &file, local_name));
                    }
                }
            }
            continue;
        }

        if !export_name.starts_with('_') {
            for star in &index.reexport_stars {
                for target_file in
                    resolve_module_files_for_export(python, &file, &star.module_specifier)
                {
                    queue.push_back((target_file, export_name.clone()));
                }
            }
        }
    }

    results.sort_by(|left, right| {
        left.source()
            .cmp(right.source())
            .then_with(|| left.fq_name().cmp(&right.fq_name()))
    });
    results.dedup();
    results
}

fn local_export_declarations(
    index: &dyn CodeUnitIndex,
    file: &ProjectFile,
    local_name: &str,
) -> Vec<CodeUnit> {
    index
        .top_level_declarations(file)
        .into_iter()
        .filter(|unit| {
            unit.identifier() == local_name
                && index
                    .parent_of(unit)
                    .is_some_and(|parent| parent.is_module() && parent.source() == file)
        })
        .collect()
}

fn resolve_module_files_for_export(
    python: &dyn PythonSource,
    importing_file: &ProjectFile,
    module_specifier: &str,
) -> Vec<ProjectFile> {
    let resolved_module = if module_specifier.starts_with('.') {
        resolve_python_relative_module(importing_file, module_specifier)
    } else {
        Some(module_specifier.to_string())
    };
    let Some(resolved_module) = resolved_module else {
        return Vec::new();
    };
    // Tree-sitter tells us the import syntax, but module-to-file resolution
    // is analyzer state. Use the prebuilt module code-unit map here instead
    // of the usage index so interactive definition lookup stays lightweight.
    resolve_module_code_unit(python, Some(importing_file), &resolved_module)
        .map(|unit| vec![unit.source().clone()])
        .unwrap_or_default()
}

pub fn extract_package_from_python_wildcard(import: &ImportInfo) -> Option<String> {
    let details = python_import_details(import)?;
    match details {
        PythonImportDetails::FromImport {
            module, wildcard, ..
        } if wildcard => Some(module),
        _ => None,
    }
}

#[derive(Debug, Clone)]
pub enum PythonImportDetails {
    Import {
        module: String,
        alias: Option<String>,
    },
    FromImport {
        module: String,
        name: String,
        alias: Option<String>,
        wildcard: bool,
    },
}

/// One import leaf interpreted directly from the Python AST.
///
/// The generic [`ImportInfo`] remains the display/query projection. The
/// source nodes are retained beside it so the coordinated primary producer can
/// intern declaration, binder, and alias occurrences without matching names or
/// byte ranges after the tree has been walked.
#[derive(Debug, Clone)]
pub struct PythonImportSyntax<'tree> {
    pub info: ImportInfo,
    pub declaration: Node<'tree>,
    pub target: Option<Node<'tree>>,
    pub alias: Option<Node<'tree>>,
}

pub fn python_import_infos_from_node(node: Node<'_>, source: &str) -> Vec<ImportInfo> {
    python_import_syntaxes_from_node(node, source)
        .into_iter()
        .map(|syntax| syntax.info)
        .collect()
}

/// Interpret every import leaf in one statement and retain its exact AST
/// binder nodes. Malformed leaves are omitted with the same admission rules as
/// the historical generic import projection.
pub fn python_import_syntaxes_from_node<'tree>(
    node: Node<'tree>,
    source: &str,
) -> Vec<PythonImportSyntax<'tree>> {
    match node.kind() {
        "import_statement" => python_namespace_import_syntaxes(node, source),
        "import_from_statement" => python_from_import_syntaxes(node, source),
        _ => Vec::new(),
    }
}

pub fn python_import_details(import: &ImportInfo) -> Option<PythonImportDetails> {
    let path = import.path.as_ref()?;
    match path.kind? {
        StructuredImportPathKind::Namespace => Some(PythonImportDetails::Import {
            module: join_python_import_segments(&path.segments),
            alias: import.alias.clone(),
        }),
        // These declaration forms belong to Java and Rust, not Python.
        StructuredImportPathKind::StaticMember | StructuredImportPathKind::ExternCrate => None,
        StructuredImportPathKind::ImportFrom => {
            let (name, module_segments) = if import.is_wildcard {
                ("*".to_string(), path.segments.as_slice())
            } else {
                let (name, module_segments) = path.segments.split_last()?;
                (name.clone(), module_segments)
            };
            Some(PythonImportDetails::FromImport {
                module: join_python_import_segments(module_segments),
                name,
                alias: import.alias.clone(),
                wildcard: import.is_wildcard,
            })
        }
    }
}

fn python_namespace_import_syntaxes<'tree>(
    node: Node<'tree>,
    source: &str,
) -> Vec<PythonImportSyntax<'tree>> {
    let mut infos = Vec::new();
    let mut cursor = node.walk();
    for imported in node.children_by_field_name("name", &mut cursor) {
        let (module_node, alias_node) = if imported.kind() == "aliased_import" {
            let Some(name) = imported.child_by_field_name("name") else {
                continue;
            };
            (name, imported.child_by_field_name("alias"))
        } else {
            (imported, None)
        };
        let alias = alias_node
            .map(|alias| py_node_text(alias, source).trim().to_string())
            .filter(|alias| !alias.is_empty());
        let segments = python_path_segments(module_node, source);
        if segments.is_empty() {
            continue;
        }
        // `import a.b` binds `a`: the first segment's own token. A renamed
        // import binds its alias token instead.
        let target = alias
            .is_none()
            .then(|| python_first_segment_node(module_node))
            .flatten();
        let binder_span = alias_node
            .or(target)
            .map(brokk_bifrost_core::analyzer::common::node_span);
        let module = join_python_import_segments(&segments);
        let identifier = alias.clone().or_else(|| segments.first().cloned());
        let info = ImportInfo {
            raw_snippet: if let Some(alias) = &alias {
                format!("import {module} as {alias}")
            } else {
                format!("import {module}")
            },
            is_wildcard: false,
            is_global: false,
            identifier,
            alias,
            path: Some(StructuredImportPath {
                segments,
                kind: Some(StructuredImportPathKind::Namespace),
                lexical_prefixes: Vec::new(),
                lexical_scopes: Vec::new(),
                declaration_start_byte: node.start_byte(),
            }),
            binder_span,
        };
        infos.push(PythonImportSyntax {
            info,
            declaration: node,
            target,
            alias: alias_node,
        });
    }
    infos
}

fn python_from_import_syntaxes<'tree>(
    node: Node<'tree>,
    source: &str,
) -> Vec<PythonImportSyntax<'tree>> {
    let Some(module_node) = node.child_by_field_name("module_name") else {
        return Vec::new();
    };
    let module_segments = python_module_segments(module_node, source);
    if module_segments.is_empty() {
        return Vec::new();
    }

    let mut infos = Vec::new();
    let has_wildcard_import = {
        let mut cursor = node.walk();
        node.named_children(&mut cursor)
            .any(|child| child.kind() == "wildcard_import")
    };
    let mut cursor = node.walk();
    let imported_names: Vec<_> = node.children_by_field_name("name", &mut cursor).collect();
    if has_wildcard_import {
        let module = join_python_import_segments(&module_segments);
        let info = ImportInfo {
            raw_snippet: format!("from {module} import *"),
            is_wildcard: true,
            is_global: false,
            identifier: None,
            alias: None,
            path: Some(StructuredImportPath {
                segments: module_segments,
                kind: Some(StructuredImportPathKind::ImportFrom),
                lexical_prefixes: Vec::new(),
                lexical_scopes: Vec::new(),
                declaration_start_byte: node.start_byte(),
            }),
            binder_span: None,
        };
        infos.push(PythonImportSyntax {
            info,
            declaration: node,
            target: None,
            alias: None,
        });
        return infos;
    }
    if imported_names.is_empty() {
        return infos;
    }

    for imported in imported_names {
        let (name_node, alias_node) = if imported.kind() == "aliased_import" {
            let Some(name) = imported.child_by_field_name("name") else {
                continue;
            };
            (name, imported.child_by_field_name("alias"))
        } else {
            (imported, None)
        };
        let alias = alias_node
            .map(|alias| py_node_text(alias, source).trim().to_string())
            .filter(|alias| !alias.is_empty());
        let name_segments = python_path_segments(name_node, source);
        if name_segments.is_empty() {
            continue;
        }
        // `from m import x` binds `x`'s own token; a rename binds the alias
        // token. A multi-segment imported name binds no single token.
        let target = alias
            .is_none()
            .then(|| {
                (name_segments.len() == 1)
                    .then(|| python_first_segment_node(name_node))
                    .flatten()
            })
            .flatten();
        let binder_span = alias_node
            .or(target)
            .map(brokk_bifrost_core::analyzer::common::node_span);
        let imported_name = join_python_import_segments(&name_segments);
        let mut segments = module_segments.clone();
        segments.extend(name_segments);
        let module = join_python_import_segments(&module_segments);
        let info = ImportInfo {
            raw_snippet: if let Some(alias) = &alias {
                format!("from {module} import {imported_name} as {alias}")
            } else {
                format!("from {module} import {imported_name}")
            },
            is_wildcard: false,
            is_global: false,
            identifier: Some(alias.clone().unwrap_or_else(|| imported_name.clone())),
            alias,
            path: Some(StructuredImportPath {
                segments,
                kind: Some(StructuredImportPathKind::ImportFrom),
                lexical_prefixes: Vec::new(),
                lexical_scopes: Vec::new(),
                declaration_start_byte: node.start_byte(),
            }),
            binder_span,
        };
        infos.push(PythonImportSyntax {
            info,
            declaration: node,
            target,
            alias: alias_node,
        });
    }
    infos
}

fn python_module_segments(module: Node<'_>, source: &str) -> Vec<String> {
    if module.kind() == "relative_import" {
        let mut cursor = module.walk();
        let mut prefix = String::new();
        let mut path_node = None;
        for child in module.named_children(&mut cursor) {
            match child.kind() {
                "import_prefix" if prefix.is_empty() => {
                    prefix = py_node_text(child, source).trim().to_string();
                }
                "dotted_name" if path_node.is_none() => {
                    path_node = Some(child);
                }
                _ => {}
            }
        }
        let mut segments = path_node
            .map(|path| python_path_segments(path, source))
            .unwrap_or_default();
        if !prefix.is_empty() {
            if let Some(first) = segments.first_mut() {
                first.insert_str(0, &prefix);
            } else {
                segments.push(prefix);
            }
        }
        return segments;
    }
    python_path_segments(module, source)
}

/// The token that spells a path's first segment: the identifier itself, or a
/// dotted name's first identifier. `None` when the shape has no leading
/// identifier token of its own (e.g. a relative-import prefix).
fn python_first_segment_node(node: Node<'_>) -> Option<Node<'_>> {
    match node.kind() {
        "identifier" => Some(node),
        "dotted_name" => {
            let mut cursor = node.walk();
            node.named_children(&mut cursor)
                .find(|child| child.kind() == "identifier")
        }
        _ => None,
    }
}

fn python_path_segments(node: Node<'_>, source: &str) -> Vec<String> {
    match node.kind() {
        "identifier" => vec![py_node_text(node, source).trim().to_string()],
        "dotted_name" => {
            let mut segments = Vec::new();
            let mut cursor = node.walk();
            for child in node.named_children(&mut cursor) {
                segments.extend(python_path_segments(child, source));
            }
            segments
        }
        _ => {
            let mut segments = Vec::new();
            let mut cursor = node.walk();
            for child in node.named_children(&mut cursor) {
                segments.extend(python_path_segments(child, source));
            }
            segments
        }
    }
}

fn join_python_import_segments(segments: &[String]) -> String {
    let Some((first, rest)) = segments.split_first() else {
        return String::new();
    };
    if first.starts_with('.') && !rest.is_empty() {
        format!("{first}.{}", rest.join("."))
    } else {
        segments.join(".")
    }
}

pub fn python_namespace_binding_name(
    import: &ImportInfo,
    alias: Option<&str>,
    module: &str,
) -> String {
    import
        .identifier
        .clone()
        .or_else(|| alias.map(str::to_string))
        .unwrap_or_else(|| module.to_string())
}

pub fn python_namespace_binding_module(
    import: &ImportInfo,
    alias: Option<&str>,
    module: &str,
) -> String {
    if alias.is_some() {
        return module.to_string();
    }
    import
        .path
        .as_ref()
        .and_then(|path| path.segments.first().cloned())
        .unwrap_or_else(|| module.to_string())
}

pub fn resolve_python_relative_module(
    source_file: &ProjectFile,
    module_expr: &str,
) -> Option<String> {
    resolve_python_relative_module_from_package(&python_current_package(source_file), module_expr)
}

/// Resolve a structured Python module expression against an already-known
/// package identity. Dependency-pack producers know module identities without
/// owning a workspace [`ProjectFile`], so they use this entry point.
pub fn resolve_python_relative_module_from_package(
    current_package: &str,
    module_expr: &str,
) -> Option<String> {
    let level = module_expr.chars().take_while(|ch| *ch == '.').count();
    let suffix = module_expr[level..].trim_matches('.');
    let mut parts: Vec<_> = current_package
        .split('.')
        .filter(|part| !part.is_empty())
        .map(str::to_string)
        .collect();
    if level == 0 {
        return Some(module_expr.to_string());
    }
    if level > 0 {
        if level - 1 > parts.len() {
            return None;
        }
        parts.truncate(parts.len() - (level - 1));
    }
    if !suffix.is_empty() {
        parts.extend(suffix.split('.').map(str::to_string));
    }
    Some(parts.join("."))
}

fn python_current_package(source_file: &ProjectFile) -> String {
    let module = python_module_name(source_file);
    if source_file
        .rel_path()
        .file_name()
        .and_then(|name| name.to_str())
        == Some("__init__.py")
    {
        module
    } else {
        module
            .rsplit_once('.')
            .map(|(package, _)| package.to_string())
            .unwrap_or_default()
    }
}