minni 0.1.1

use super::languages::SupportedLanguage;
use crate::db::SymbolEdge;
use chrono::Utc;
use tree_sitter::Tree;
use uuid::Uuid;

/// Extract import/use edges from a parsed tree-sitter `Tree`.
///
/// Supports Rust (`use` declarations) and Python (`import` / `from … import`).
/// For all other languages, returns an empty `Vec`.
///
/// `target_file` is always `None` in this implementation — resolution is deferred
/// to a future task.
pub fn extract_imports(
    tree: &Tree,
    content: &str,
    language: SupportedLanguage,
    source_file: &str,
) -> Vec<SymbolEdge> {
    match language {
        SupportedLanguage::Rust => extract_rust_imports(tree, content, source_file),
        SupportedLanguage::Python => extract_python_imports(tree, content, source_file),
        _ => Vec::new(),
    }
}

// ── Rust ──────────────────────────────────────────────────────────────────────

fn extract_rust_imports(tree: &Tree, content: &str, source_file: &str) -> Vec<SymbolEdge> {
    let mut edges = Vec::new();
    let root = tree.root_node();
    let timestamp = Utc::now().to_rfc3339();

    for i in 0..root.child_count() {
        let node = match root.child(i) {
            Some(n) => n,
            None => continue,
        };

        if node.kind() != "use_declaration" {
            continue;
        }

        // The use_declaration has a single child after the `use` keyword: the use tree.
        // Walk that subtree to collect all full paths.
        let mut targets: Vec<String> = Vec::new();
        collect_rust_use_targets(&node, content, &mut Vec::new(), &mut targets);

        for target in targets {
            edges.push(SymbolEdge {
                id: Uuid::new_v4().to_string(),
                source_file: source_file.to_string(),
                target_symbol: target,
                edge_type: "imports".to_string(),
                target_file: None,
                indexed_at: timestamp.clone(),
            });
        }
    }

    edges
}

/// Recursively walk a `use_declaration` child node, accumulating the path prefix
/// in `prefix_parts` and emitting complete target paths into `targets`.
fn collect_rust_use_targets(
    node: &tree_sitter::Node,
    content: &str,
    prefix_parts: &mut Vec<String>,
    targets: &mut Vec<String>,
) {
    let kind = node.kind();

    match kind {
        // use_declaration itself: descend into children to find the actual use tree
        "use_declaration" => {
            for i in 0..node.child_count() {
                if let Some(child) = node.child(i) {
                    // Skip keywords and punctuation; look for the use subtree
                    let ck = child.kind();
                    if ck != "use"
                        && ck != ";"
                        && ck != "pub"
                        && ck != "visibility_modifier"
                        && ck != "line_comment"
                        && ck != "block_comment"
                    {
                        collect_rust_use_targets(&child, content, prefix_parts, targets);
                    }
                }
            }
        }

        // scoped_identifier: foo::bar — treat left as prefix, right as final segment
        "scoped_identifier" => {
            // Collect segments for this scoped path
            let mut parts: Vec<String> = Vec::new();
            flatten_scoped_identifier(node, content, &mut parts);
            let combined = if prefix_parts.is_empty() {
                parts.join("::")
            } else {
                format!("{}::{}", prefix_parts.join("::"), parts.join("::"))
            };
            if !combined.is_empty() {
                targets.push(combined);
            }
        }

        // scoped_use_list: foo::{bar, baz}
        "scoped_use_list" => {
            // Left child is the prefix path (scoped_identifier or identifier)
            // Right child is the use_list
            let mut new_prefix = prefix_parts.clone();
            let child_count = node.child_count();
            let mut i = 0;
            while i < child_count {
                if let Some(child) = node.child(i) {
                    let ck = child.kind();
                    if ck != "::" && ck != "{" && ck != "}" {
                        if ck == "use_list" {
                            collect_rust_use_targets(&child, content, &mut new_prefix, targets);
                        } else {
                            // This is the prefix part (identifier or scoped_identifier)
                            let mut parts = Vec::new();
                            flatten_scoped_identifier(&child, content, &mut parts);
                            if !parts.is_empty() {
                                let seg = parts.join("::");
                                if !new_prefix.is_empty() {
                                    new_prefix.push(seg);
                                } else {
                                    new_prefix = parts;
                                }
                            }
                        }
                    }
                }
                i += 1;
            }
        }

        // use_list: {bar, baz, *}
        "use_list" => {
            for i in 0..node.child_count() {
                if let Some(child) = node.child(i) {
                    let ck = child.kind();
                    if ck != "{" && ck != "}" && ck != "," {
                        collect_rust_use_targets(&child, content, prefix_parts, targets);
                    }
                }
            }
        }

        // use_as_clause: foo as bar — use the original path (ignore alias)
        "use_as_clause" => {
            // First child is the original item (identifier or scoped_identifier)
            if let Some(orig) = node.child(0) {
                collect_rust_use_targets(&orig, content, prefix_parts, targets);
            }
        }

        // use_wildcard: the node text contains the full path including `::*`
        // (e.g. `std::io::*`).  When it is a bare `*` inside a grouped list the
        // prefix is stored in `prefix_parts`, so we fall back to that.
        "use_wildcard" => {
            let raw = node.utf8_text(content.as_bytes()).unwrap_or("").to_string();
            let path = if raw.contains("::") || prefix_parts.is_empty() {
                // The node already encodes the full path (e.g. "std::io::*")
                raw
            } else {
                format!("{}::*", prefix_parts.join("::"))
            };
            if !path.is_empty() {
                targets.push(path);
            }
        }

        // identifier: a bare name, combine with existing prefix.
        // Special case: `self` in a use-list (e.g. `use foo::{self, Bar}`) means
        // re-export / import of the module itself — emit just the prefix, not `foo::self`.
        "identifier" | "self" | "super" | "crate" => {
            let text = node.utf8_text(content.as_bytes()).unwrap_or("").to_string();
            if !text.is_empty() {
                let path = if text == "self" && !prefix_parts.is_empty() {
                    // `self` inside a grouped import → the module itself
                    prefix_parts.join("::")
                } else if prefix_parts.is_empty() {
                    text
                } else {
                    format!("{}::{}", prefix_parts.join("::"), text)
                };
                targets.push(path);
            }
        }

        _ => {
            // Unknown node — try to descend in case grammar changes
        }
    }
}

/// Flatten a `scoped_identifier` (or plain `identifier`) into a list of segments.
fn flatten_scoped_identifier(node: &tree_sitter::Node, content: &str, parts: &mut Vec<String>) {
    match node.kind() {
        "scoped_identifier" => {
            for i in 0..node.child_count() {
                if let Some(child) = node.child(i) {
                    if child.kind() != "::" {
                        flatten_scoped_identifier(&child, content, parts);
                    }
                }
            }
        }
        "identifier" | "self" | "super" | "crate" => {
            let text = node.utf8_text(content.as_bytes()).unwrap_or("").to_string();
            if !text.is_empty() {
                parts.push(text);
            }
        }
        _ => {
            // e.g. use_wildcard inside a scoped path — emit as-is
            let text = node.utf8_text(content.as_bytes()).unwrap_or("").to_string();
            if !text.is_empty() {
                parts.push(text);
            }
        }
    }
}

// ── Python ─────────────────────────────────────────────────────────────────────

fn extract_python_imports(tree: &Tree, content: &str, source_file: &str) -> Vec<SymbolEdge> {
    let mut edges = Vec::new();
    let root = tree.root_node();
    let timestamp = Utc::now().to_rfc3339();
    let content_bytes = content.as_bytes();

    for i in 0..root.child_count() {
        let node = match root.child(i) {
            Some(n) => n,
            None => continue,
        };

        match node.kind() {
            // import os  /  import os.path  /  import os as o
            "import_statement" => {
                // named children: name nodes (dotted_name or aliased_import)
                for j in 0..node.child_count() {
                    if let Some(child) = node.child(j) {
                        let ck = child.kind();
                        if ck == "dotted_name" {
                            let text = child.utf8_text(content_bytes).unwrap_or("").to_string();
                            if !text.is_empty() {
                                edges.push(make_edge(source_file, &text, &timestamp));
                            }
                        } else if ck == "aliased_import" {
                            // aliased_import: dotted_name as identifier
                            // use the original name (first child)
                            if let Some(orig) = child.child(0) {
                                let text = orig.utf8_text(content_bytes).unwrap_or("").to_string();
                                if !text.is_empty() {
                                    edges.push(make_edge(source_file, &text, &timestamp));
                                }
                            }
                        }
                    }
                }
            }

            // from os.path import join, exists
            // from . import utils
            // from ..models import Foo
            "import_from_statement" => {
                // Structure: "from" <module> "import" <names...>
                // module can be: dotted_name, relative_import (leading dots), or None (from . import)
                let module_prefix = extract_python_from_module(&node, content_bytes);

                // Collect imported names (after "import")
                let mut past_import_kw = false;
                for j in 0..node.child_count() {
                    if let Some(child) = node.child(j) {
                        let ck = child.kind();

                        if ck == "import" {
                            past_import_kw = true;
                            continue;
                        }
                        if !past_import_kw {
                            continue;
                        }

                        // Dispatch on node kind; recurse into container nodes
                        // (e.g. `import_list` / parenthesized wrappers in `from x import (a, b)`).
                        collect_python_import_names(
                            &child,
                            content_bytes,
                            &module_prefix,
                            source_file,
                            &timestamp,
                            &mut edges,
                        );
                    }
                }
            }

            _ => {}
        }
    }

    edges
}

/// Emit edges for a single child node of an `import_from_statement` that appears
/// after the `import` keyword.  Handles:
/// - `wildcard_import` → `module.*`
/// - `identifier` / `dotted_name` → `module.name`
/// - `aliased_import` → `module.original_name`
/// - any container-like node (e.g. `import_list`, parenthesized forms) → recurse
fn collect_python_import_names(
    node: &tree_sitter::Node,
    content_bytes: &[u8],
    module_prefix: &str,
    source_file: &str,
    timestamp: &str,
    edges: &mut Vec<SymbolEdge>,
) {
    let ck = node.kind();
    if ck == "wildcard_import" {
        let target = join_python_module(module_prefix, "*");
        edges.push(make_edge(source_file, &target, timestamp));
    } else if ck == "dotted_name" || ck == "identifier" {
        let name = node.utf8_text(content_bytes).unwrap_or("").to_string();
        if !name.is_empty() {
            let target = join_python_module(module_prefix, &name);
            edges.push(make_edge(source_file, &target, timestamp));
        }
    } else if ck == "aliased_import" {
        // use original name (first child), ignore alias
        if let Some(orig) = node.child(0) {
            let name = orig.utf8_text(content_bytes).unwrap_or("").to_string();
            if !name.is_empty() {
                let target = join_python_module(module_prefix, &name);
                edges.push(make_edge(source_file, &target, timestamp));
            }
        }
    } else if ck != "," && ck != "(" && ck != ")" {
        // Potential container node (import_list, parenthesized_expression, etc.) — recurse
        for i in 0..node.child_count() {
            if let Some(child) = node.child(i) {
                collect_python_import_names(
                    &child,
                    content_bytes,
                    module_prefix,
                    source_file,
                    timestamp,
                    edges,
                );
            }
        }
    }
}

/// Extract the module prefix string from a `import_from_statement` node.
///
/// Handles:
/// - `from os.path import …`  → `"os.path"`
/// - `from . import …`        → `"."`
/// - `from ..models import …` → `"..models"`
fn extract_python_from_module(node: &tree_sitter::Node, content_bytes: &[u8]) -> String {
    let mut dots = String::new();
    let mut module_name = String::new();

    for i in 0..node.child_count() {
        if let Some(child) = node.child(i) {
            match child.kind() {
                // Stop processing once we reach the "import" keyword
                "import" => break,
                "from" => {}
                "relative_import" => {
                    // relative_import: one or more "." followed by optional dotted_name
                    for j in 0..child.child_count() {
                        if let Some(rc) = child.child(j) {
                            match rc.kind() {
                                "import_prefix" => {
                                    // import_prefix holds the leading dots
                                    dots = rc.utf8_text(content_bytes).unwrap_or("").to_string();
                                }
                                "dotted_name" => {
                                    module_name =
                                        rc.utf8_text(content_bytes).unwrap_or("").to_string();
                                }
                                _ => {}
                            }
                        }
                    }
                }
                "dotted_name" => {
                    module_name = child.utf8_text(content_bytes).unwrap_or("").to_string();
                }
                _ => {}
            }
        }
    }

    if dots.is_empty() {
        module_name
    } else if module_name.is_empty() {
        // e.g. `from . import utils` → prefix is "."
        dots
    } else {
        // e.g. `from ..models import Foo` → prefix is "..models"
        format!("{}{}", dots, module_name)
    }
}

/// Join a Python module prefix with a name using `.`, taking care not to
/// double-up dots when the prefix itself ends with one (e.g. relative imports
/// like `"."` + `"utils"` → `".utils"`, not `"..utils"`).
fn join_python_module(prefix: &str, name: &str) -> String {
    if prefix.is_empty() {
        name.to_string()
    } else if prefix.ends_with('.') {
        format!("{}{}", prefix, name)
    } else {
        format!("{}.{}", prefix, name)
    }
}

fn make_edge(source_file: &str, target_symbol: &str, timestamp: &str) -> SymbolEdge {
    SymbolEdge {
        id: Uuid::new_v4().to_string(),
        source_file: source_file.to_string(),
        target_symbol: target_symbol.to_string(),
        edge_type: "imports".to_string(),
        target_file: None,
        indexed_at: timestamp.to_string(),
    }
}

// ── Tests ──────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use tree_sitter::Parser;

    fn parse_rust(src: &str) -> Tree {
        let mut parser = Parser::new();
        parser
            .set_language(&tree_sitter_rust::language())
            .expect("failed to set Rust language");
        parser.parse(src, None).expect("failed to parse Rust")
    }

    fn parse_python(src: &str) -> Tree {
        let mut parser = Parser::new();
        parser
            .set_language(&tree_sitter_python::language())
            .expect("failed to set Python language");
        parser.parse(src, None).expect("failed to parse Python")
    }

    fn symbols(edges: &[SymbolEdge]) -> Vec<&str> {
        let mut v: Vec<&str> = edges.iter().map(|e| e.target_symbol.as_str()).collect();
        v.sort_unstable();
        v
    }

    // ── Rust tests ─────────────────────────────────────────────────────────────

    #[test]
    fn test_rust_simple_use() {
        let src = "use std::fs::File;\nfn main() {}";
        let tree = parse_rust(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Rust, "src/main.rs");
        assert_eq!(symbols(&edges), vec!["std::fs::File"]);
        assert!(edges.iter().all(|e| e.edge_type == "imports"));
        assert!(edges.iter().all(|e| e.target_file.is_none()));
        assert!(edges.iter().all(|e| e.source_file == "src/main.rs"));
    }

    #[test]
    fn test_rust_grouped_use() {
        let src = "use crate::db::{CodeChunk, Database};\nfn main() {}";
        let tree = parse_rust(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Rust, "src/lib.rs");
        assert_eq!(
            symbols(&edges),
            vec!["crate::db::CodeChunk", "crate::db::Database"]
        );
    }

    #[test]
    fn test_rust_self_in_grouped_use() {
        // `use crate::db::{self, Database}` — `self` means re-import the module itself.
        let src = "use crate::db::{self, Database};\nfn main() {}";
        let tree = parse_rust(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Rust, "src/lib.rs");
        // Should produce exactly `crate::db` (not `crate::db::self`) and `crate::db::Database`
        assert_eq!(symbols(&edges), vec!["crate::db", "crate::db::Database"]);
    }

    #[test]
    fn test_rust_glob_use() {
        let src = "use std::io::*;\nfn main() {}";
        let tree = parse_rust(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Rust, "src/main.rs");
        assert_eq!(symbols(&edges), vec!["std::io::*"]);
    }

    #[test]
    fn test_rust_use_as_alias() {
        let src = "use std::collections::HashMap as Map;\nfn main() {}";
        let tree = parse_rust(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Rust, "src/main.rs");
        // Should emit the original path, not the alias
        assert_eq!(symbols(&edges), vec!["std::collections::HashMap"]);
    }

    #[test]
    fn test_rust_multiple_uses() {
        let src = r#"
use anyhow::Result;
use std::path::Path;
use super::languages::SupportedLanguage;
fn main() {}
"#;
        let tree = parse_rust(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Rust, "src/indexer/mod.rs");
        let syms = symbols(&edges);
        assert!(syms.contains(&"anyhow::Result"), "{:?}", syms);
        assert!(syms.contains(&"std::path::Path"), "{:?}", syms);
        assert!(
            syms.contains(&"super::languages::SupportedLanguage"),
            "{:?}",
            syms
        );
    }

    #[test]
    fn test_rust_empty_file() {
        let src = "";
        let tree = parse_rust(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Rust, "src/empty.rs");
        assert!(edges.is_empty());
    }

    #[test]
    fn test_rust_no_imports() {
        let src = "fn main() { println!(\"hello\"); }";
        let tree = parse_rust(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Rust, "src/main.rs");
        assert!(edges.is_empty());
    }

    // ── Python tests ───────────────────────────────────────────────────────────

    #[test]
    fn test_python_simple_import() {
        let src = "import os\n";
        let tree = parse_python(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Python, "app/main.py");
        assert_eq!(symbols(&edges), vec!["os"]);
    }

    #[test]
    fn test_python_dotted_import() {
        let src = "import os.path\n";
        let tree = parse_python(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Python, "app/main.py");
        assert_eq!(symbols(&edges), vec!["os.path"]);
    }

    #[test]
    fn test_python_from_import_single() {
        let src = "from os.path import join\n";
        let tree = parse_python(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Python, "app/main.py");
        assert_eq!(symbols(&edges), vec!["os.path.join"]);
    }

    #[test]
    fn test_python_from_import_multiple() {
        let src = "from os.path import join, exists\n";
        let tree = parse_python(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Python, "app/main.py");
        assert_eq!(symbols(&edges), vec!["os.path.exists", "os.path.join"]);
    }

    #[test]
    fn test_python_from_import_parenthesized() {
        let src = "from os.path import (join, exists)\n";
        let tree = parse_python(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Python, "app/main.py");
        assert_eq!(symbols(&edges), vec!["os.path.exists", "os.path.join"]);
    }

    #[test]
    fn test_python_from_import_aliased() {
        let src = "from os.path import join as j\n";
        let tree = parse_python(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Python, "app/main.py");
        // Should emit the original name, not the alias
        assert_eq!(symbols(&edges), vec!["os.path.join"]);
    }

    #[test]
    fn test_python_relative_import_single_dot() {
        let src = "from . import utils\n";
        let tree = parse_python(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Python, "app/sub/mod.py");
        assert_eq!(symbols(&edges), vec![".utils"]);
    }

    #[test]
    fn test_python_relative_import_double_dot() {
        let src = "from ..models import Foo\n";
        let tree = parse_python(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Python, "app/sub/mod.py");
        assert_eq!(symbols(&edges), vec!["..models.Foo"]);
    }

    #[test]
    fn test_python_empty_file() {
        let src = "";
        let tree = parse_python(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Python, "app/empty.py");
        assert!(edges.is_empty());
    }

    #[test]
    fn test_python_no_imports() {
        let src = "def foo():\n    pass\n";
        let tree = parse_python(src);
        let edges = extract_imports(&tree, src, SupportedLanguage::Python, "app/main.py");
        assert!(edges.is_empty());
    }

    // ── Unsupported language ───────────────────────────────────────────────────

    #[test]
    fn test_unsupported_language_returns_empty() {
        // Use Go as an example unsupported language for import extraction
        let src = "package main\nimport \"fmt\"\n";
        // We can only parse this if we have a Go parser — use Rust parser as a proxy
        // and pass Go as language to exercise the match arm.
        let mut parser = Parser::new();
        parser
            .set_language(&tree_sitter_rust::language())
            .expect("set language");
        // Parse as Rust (will still return a tree even if content is invalid)
        let tree = parser.parse(src, None).expect("parse");
        let edges = extract_imports(&tree, src, SupportedLanguage::Go, "main.go");
        assert!(edges.is_empty(), "Go should return empty edges");
    }

    #[test]
    fn test_unsupported_languages_return_empty() {
        let src = "fn main() {}";
        let tree = parse_rust(src);
        for lang in [
            SupportedLanguage::JavaScript,
            SupportedLanguage::TypeScript,
            SupportedLanguage::Go,
            SupportedLanguage::C,
            SupportedLanguage::Cpp,
            SupportedLanguage::Java,
            SupportedLanguage::CSharp,
        ] {
            let edges = extract_imports(&tree, src, lang, "file");
            assert!(
                edges.is_empty(),
                "{} should return empty edges",
                lang.name()
            );
        }
    }
}