Skip to main content

aft/
imports.rs

1//! Import analysis engine: parsing, grouping, deduplication, and insertion.
2//!
3//! Provides per-language import handling dispatched by `LangId`. Each language
4//! implementation extracts imports from tree-sitter ASTs, classifies them into
5//! groups, and generates import text.
6//!
7//! Currently supports: TypeScript, TSX, JavaScript.
8
9use std::ops::Range;
10
11use tree_sitter::{Node, Parser, Tree};
12
13use crate::parser::{grammar_for, LangId};
14
15// ---------------------------------------------------------------------------
16// Shared types
17// ---------------------------------------------------------------------------
18
19/// What kind of import this is.
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum ImportKind {
22    /// `import { X } from 'y'` or `import X from 'y'`
23    Value,
24    /// `import type { X } from 'y'`
25    Type,
26    /// `import './side-effect'`
27    SideEffect,
28}
29
30/// Which logical group an import belongs to (language-specific).
31///
32/// Ordering matches conventional import group sorting:
33///   Stdlib (first) < External < Internal (last)
34///
35/// Language mapping:
36///   - TS/JS/TSX: External (no `.` prefix), Internal (`.`/`..` prefix)
37///   - Python:    Stdlib, External (third-party), Internal (relative `.`/`..`)
38///   - Rust:      Stdlib (std/core/alloc), External (crates), Internal (crate/self/super)
39///   - Go:        Stdlib (no dots in path), External (dots in path)
40#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
41pub enum ImportGroup {
42    /// Standard library (Python stdlib, Rust std/core/alloc, Go stdlib).
43    /// TS/JS don't use this group.
44    Stdlib,
45    /// External/third-party packages.
46    External,
47    /// Internal/relative imports (TS relative, Python local, Rust crate/self/super).
48    Internal,
49}
50
51impl ImportGroup {
52    /// Human-readable label for the group.
53    pub fn label(&self) -> &'static str {
54        match self {
55            ImportGroup::Stdlib => "stdlib",
56            ImportGroup::External => "external",
57            ImportGroup::Internal => "internal",
58        }
59    }
60}
61
62/// A single parsed import statement.
63#[derive(Debug, Clone)]
64pub struct ImportStatement {
65    /// The module path (e.g., `react`, `./utils`, `../config`).
66    pub module_path: String,
67    /// Named imports (e.g., `["useState", "useEffect"]`).
68    pub names: Vec<String>,
69    /// Default import name (e.g., `React` from `import React from 'react'`).
70    pub default_import: Option<String>,
71    /// Namespace import name (e.g., `path` from `import * as path from 'path'`).
72    pub namespace_import: Option<String>,
73    /// What kind: value, type, or side-effect.
74    pub kind: ImportKind,
75    /// Which group this import belongs to.
76    pub group: ImportGroup,
77    /// Byte range in the original source.
78    pub byte_range: Range<usize>,
79    /// Raw text of the import statement.
80    pub raw_text: String,
81}
82
83/// A block of parsed imports from a file.
84#[derive(Debug, Clone)]
85pub struct ImportBlock {
86    /// All parsed import statements, in source order.
87    pub imports: Vec<ImportStatement>,
88    /// Overall byte range covering all import statements (start of first to end of last).
89    /// `None` if no imports found.
90    pub byte_range: Option<Range<usize>>,
91}
92
93impl ImportBlock {
94    pub fn empty() -> Self {
95        ImportBlock {
96            imports: Vec::new(),
97            byte_range: None,
98        }
99    }
100}
101
102// ---------------------------------------------------------------------------
103// Core API
104// ---------------------------------------------------------------------------
105
106/// Parse imports from source using the provided tree-sitter tree.
107pub fn parse_imports(source: &str, tree: &Tree, lang: LangId) -> ImportBlock {
108    match lang {
109        LangId::TypeScript | LangId::Tsx | LangId::JavaScript => parse_ts_imports(source, tree),
110        LangId::Python => parse_py_imports(source, tree),
111        LangId::Rust => parse_rs_imports(source, tree),
112        LangId::Go => parse_go_imports(source, tree),
113        LangId::Markdown => ImportBlock::empty(),
114    }
115}
116
117/// Check if an import with the given module + name combination already exists.
118///
119/// For dedup: same module path AND (same named import OR same default import).
120/// Side-effect imports match on module path alone.
121pub fn is_duplicate(
122    block: &ImportBlock,
123    module_path: &str,
124    names: &[String],
125    default_import: Option<&str>,
126    type_only: bool,
127) -> bool {
128    let target_kind = if type_only {
129        ImportKind::Type
130    } else {
131        ImportKind::Value
132    };
133
134    for imp in &block.imports {
135        if imp.module_path != module_path {
136            continue;
137        }
138
139        // For side-effect imports or whole-module imports (no names, no default):
140        // module path match alone is sufficient.
141        if names.is_empty()
142            && default_import.is_none()
143            && imp.names.is_empty()
144            && imp.default_import.is_none()
145        {
146            return true;
147        }
148
149        // For side-effect imports specifically (TS/JS): module match is enough
150        if names.is_empty() && default_import.is_none() && imp.kind == ImportKind::SideEffect {
151            return true;
152        }
153
154        // Kind must match for dedup (value imports don't dedup against type imports)
155        if imp.kind != target_kind && imp.kind != ImportKind::SideEffect {
156            continue;
157        }
158
159        // Check default import match
160        if let Some(def) = default_import {
161            if imp.default_import.as_deref() == Some(def) {
162                return true;
163            }
164        }
165
166        // Check named imports — if ALL requested names already exist
167        if !names.is_empty() && names.iter().all(|n| imp.names.contains(n)) {
168            return true;
169        }
170    }
171
172    false
173}
174
175/// Find the byte offset where a new import should be inserted.
176///
177/// Strategy:
178/// - Find all existing imports in the same group.
179/// - Within that group, find the alphabetical position by module path.
180/// - Type imports sort after value imports within the same group and module-sort position.
181/// - If no imports exist in the target group, insert after the last import of the
182///   nearest preceding group (or before the first import of the nearest following
183///   group, or at file start if no groups exist).
184/// - Returns (byte_offset, needs_newline_before, needs_newline_after)
185pub fn find_insertion_point(
186    source: &str,
187    block: &ImportBlock,
188    group: ImportGroup,
189    module_path: &str,
190    type_only: bool,
191) -> (usize, bool, bool) {
192    if block.imports.is_empty() {
193        // No imports at all — insert at start of file
194        return (0, false, source.is_empty().then_some(false).unwrap_or(true));
195    }
196
197    let target_kind = if type_only {
198        ImportKind::Type
199    } else {
200        ImportKind::Value
201    };
202
203    // Collect imports in the target group
204    let group_imports: Vec<&ImportStatement> =
205        block.imports.iter().filter(|i| i.group == group).collect();
206
207    if group_imports.is_empty() {
208        // No imports in this group yet — find nearest neighbor group
209        // Try preceding groups (lower ordinal) first
210        let preceding_last = block.imports.iter().filter(|i| i.group < group).last();
211
212        if let Some(last) = preceding_last {
213            let end = last.byte_range.end;
214            let insert_at = skip_newline(source, end);
215            return (insert_at, true, true);
216        }
217
218        // No preceding group — try following groups (higher ordinal)
219        let following_first = block.imports.iter().find(|i| i.group > group);
220
221        if let Some(first) = following_first {
222            return (first.byte_range.start, false, true);
223        }
224
225        // Shouldn't reach here if block is non-empty, but handle gracefully
226        let first_byte = block.imports.first().unwrap().byte_range.start;
227        return (first_byte, false, true);
228    }
229
230    // Find position within the group (alphabetical by module path, type after value)
231    for imp in &group_imports {
232        let cmp = module_path.cmp(&imp.module_path);
233        match cmp {
234            std::cmp::Ordering::Less => {
235                // Insert before this import
236                return (imp.byte_range.start, false, false);
237            }
238            std::cmp::Ordering::Equal => {
239                // Same module — type imports go after value imports
240                if target_kind == ImportKind::Type && imp.kind == ImportKind::Value {
241                    // Insert after this value import
242                    let end = imp.byte_range.end;
243                    let insert_at = skip_newline(source, end);
244                    return (insert_at, false, false);
245                }
246                // Insert before (or it's a duplicate, caller should have checked)
247                return (imp.byte_range.start, false, false);
248            }
249            std::cmp::Ordering::Greater => continue,
250        }
251    }
252
253    // Module path sorts after all existing imports in this group — insert at end
254    let last = group_imports.last().unwrap();
255    let end = last.byte_range.end;
256    let insert_at = skip_newline(source, end);
257    (insert_at, false, false)
258}
259
260/// Generate an import line for the given language.
261pub fn generate_import_line(
262    lang: LangId,
263    module_path: &str,
264    names: &[String],
265    default_import: Option<&str>,
266    type_only: bool,
267) -> String {
268    match lang {
269        LangId::TypeScript | LangId::Tsx | LangId::JavaScript => {
270            generate_ts_import_line(module_path, names, default_import, type_only)
271        }
272        LangId::Python => generate_py_import_line(module_path, names, default_import),
273        LangId::Rust => generate_rs_import_line(module_path, names, type_only),
274        LangId::Go => generate_go_import_line(module_path, default_import, false),
275        LangId::Markdown => String::new(),
276    }
277}
278
279/// Check if the given language is supported by the import engine.
280pub fn is_supported(lang: LangId) -> bool {
281    matches!(
282        lang,
283        LangId::TypeScript
284            | LangId::Tsx
285            | LangId::JavaScript
286            | LangId::Python
287            | LangId::Rust
288            | LangId::Go
289    )
290}
291
292/// Classify a module path into a group for TS/JS/TSX.
293pub fn classify_group_ts(module_path: &str) -> ImportGroup {
294    if module_path.starts_with('.') {
295        ImportGroup::Internal
296    } else {
297        ImportGroup::External
298    }
299}
300
301/// Classify a module path into a group for the given language.
302pub fn classify_group(lang: LangId, module_path: &str) -> ImportGroup {
303    match lang {
304        LangId::TypeScript | LangId::Tsx | LangId::JavaScript => classify_group_ts(module_path),
305        LangId::Python => classify_group_py(module_path),
306        LangId::Rust => classify_group_rs(module_path),
307        LangId::Go => classify_group_go(module_path),
308        LangId::Markdown => ImportGroup::External,
309    }
310}
311
312/// Parse a file from disk and return its import block.
313/// Convenience wrapper that handles parsing.
314pub fn parse_file_imports(
315    path: &std::path::Path,
316    lang: LangId,
317) -> Result<(String, Tree, ImportBlock), crate::error::AftError> {
318    let source =
319        std::fs::read_to_string(path).map_err(|e| crate::error::AftError::FileNotFound {
320            path: format!("{}: {}", path.display(), e),
321        })?;
322
323    let grammar = grammar_for(lang);
324    let mut parser = Parser::new();
325    parser
326        .set_language(&grammar)
327        .map_err(|e| crate::error::AftError::ParseError {
328            message: format!("grammar init failed for {:?}: {}", lang, e),
329        })?;
330
331    let tree = parser
332        .parse(&source, None)
333        .ok_or_else(|| crate::error::AftError::ParseError {
334            message: format!("tree-sitter parse returned None for {}", path.display()),
335        })?;
336
337    let block = parse_imports(&source, &tree, lang);
338    Ok((source, tree, block))
339}
340
341// ---------------------------------------------------------------------------
342// TS/JS/TSX implementation
343// ---------------------------------------------------------------------------
344
345/// Parse imports from a TS/JS/TSX file.
346///
347/// Walks the AST root's direct children looking for `import_statement` nodes (D041).
348fn parse_ts_imports(source: &str, tree: &Tree) -> ImportBlock {
349    let root = tree.root_node();
350    let mut imports = Vec::new();
351
352    let mut cursor = root.walk();
353    if !cursor.goto_first_child() {
354        return ImportBlock::empty();
355    }
356
357    loop {
358        let node = cursor.node();
359        if node.kind() == "import_statement" {
360            if let Some(imp) = parse_single_ts_import(source, &node) {
361                imports.push(imp);
362            }
363        }
364        if !cursor.goto_next_sibling() {
365            break;
366        }
367    }
368
369    let byte_range = if imports.is_empty() {
370        None
371    } else {
372        let start = imports.first().unwrap().byte_range.start;
373        let end = imports.last().unwrap().byte_range.end;
374        Some(start..end)
375    };
376
377    ImportBlock {
378        imports,
379        byte_range,
380    }
381}
382
383/// Parse a single `import_statement` node into an `ImportStatement`.
384fn parse_single_ts_import(source: &str, node: &Node) -> Option<ImportStatement> {
385    let raw_text = source[node.byte_range()].to_string();
386    let byte_range = node.byte_range();
387
388    // Find the source module (string/string_fragment child of the import)
389    let module_path = extract_module_path(source, node)?;
390
391    // Determine if this is a type-only import: `import type ...`
392    let is_type_only = has_type_keyword(node);
393
394    // Extract import clause details
395    let mut names = Vec::new();
396    let mut default_import = None;
397    let mut namespace_import = None;
398
399    let mut child_cursor = node.walk();
400    if child_cursor.goto_first_child() {
401        loop {
402            let child = child_cursor.node();
403            match child.kind() {
404                "import_clause" => {
405                    extract_import_clause(
406                        source,
407                        &child,
408                        &mut names,
409                        &mut default_import,
410                        &mut namespace_import,
411                    );
412                }
413                // In some grammars, the default import is a direct identifier child
414                "identifier" => {
415                    let text = &source[child.byte_range()];
416                    if text != "import" && text != "from" && text != "type" {
417                        default_import = Some(text.to_string());
418                    }
419                }
420                _ => {}
421            }
422            if !child_cursor.goto_next_sibling() {
423                break;
424            }
425        }
426    }
427
428    // Classify kind
429    let kind = if names.is_empty() && default_import.is_none() && namespace_import.is_none() {
430        ImportKind::SideEffect
431    } else if is_type_only {
432        ImportKind::Type
433    } else {
434        ImportKind::Value
435    };
436
437    let group = classify_group_ts(&module_path);
438
439    Some(ImportStatement {
440        module_path,
441        names,
442        default_import,
443        namespace_import,
444        kind,
445        group,
446        byte_range,
447        raw_text,
448    })
449}
450
451/// Extract the module path string from an import_statement node.
452///
453/// Looks for a `string` child node and extracts the content without quotes.
454fn extract_module_path(source: &str, node: &Node) -> Option<String> {
455    let mut cursor = node.walk();
456    if !cursor.goto_first_child() {
457        return None;
458    }
459
460    loop {
461        let child = cursor.node();
462        if child.kind() == "string" {
463            // Get the text and strip quotes
464            let text = &source[child.byte_range()];
465            let stripped = text
466                .trim_start_matches(|c| c == '\'' || c == '"')
467                .trim_end_matches(|c| c == '\'' || c == '"');
468            return Some(stripped.to_string());
469        }
470        if !cursor.goto_next_sibling() {
471            break;
472        }
473    }
474    None
475}
476
477/// Check if the import_statement has a `type` keyword (import type ...).
478///
479/// In tree-sitter-typescript, `import type { X } from 'y'` produces a `type`
480/// node as a direct child of `import_statement`, between `import` and `import_clause`.
481fn has_type_keyword(node: &Node) -> bool {
482    let mut cursor = node.walk();
483    if !cursor.goto_first_child() {
484        return false;
485    }
486
487    loop {
488        let child = cursor.node();
489        if child.kind() == "type" {
490            return true;
491        }
492        if !cursor.goto_next_sibling() {
493            break;
494        }
495    }
496
497    false
498}
499
500/// Extract named imports, default import, and namespace import from an import_clause.
501fn extract_import_clause(
502    source: &str,
503    node: &Node,
504    names: &mut Vec<String>,
505    default_import: &mut Option<String>,
506    namespace_import: &mut Option<String>,
507) {
508    let mut cursor = node.walk();
509    if !cursor.goto_first_child() {
510        return;
511    }
512
513    loop {
514        let child = cursor.node();
515        match child.kind() {
516            "identifier" => {
517                // This is a default import: `import Foo from 'bar'`
518                let text = &source[child.byte_range()];
519                if text != "type" {
520                    *default_import = Some(text.to_string());
521                }
522            }
523            "named_imports" => {
524                // `{ name1, name2 }`
525                extract_named_imports(source, &child, names);
526            }
527            "namespace_import" => {
528                // `* as name`
529                extract_namespace_import(source, &child, namespace_import);
530            }
531            _ => {}
532        }
533        if !cursor.goto_next_sibling() {
534            break;
535        }
536    }
537}
538
539/// Extract individual names from a named_imports node (`{ a, b, c }`).
540fn extract_named_imports(source: &str, node: &Node, names: &mut Vec<String>) {
541    let mut cursor = node.walk();
542    if !cursor.goto_first_child() {
543        return;
544    }
545
546    loop {
547        let child = cursor.node();
548        if child.kind() == "import_specifier" {
549            // import_specifier can have `name` (the imported name) and optional `alias`
550            if let Some(name_node) = child.child_by_field_name("name") {
551                names.push(source[name_node.byte_range()].to_string());
552            } else {
553                // Fallback: first identifier child
554                let mut spec_cursor = child.walk();
555                if spec_cursor.goto_first_child() {
556                    loop {
557                        if spec_cursor.node().kind() == "identifier"
558                            || spec_cursor.node().kind() == "type_identifier"
559                        {
560                            names.push(source[spec_cursor.node().byte_range()].to_string());
561                            break;
562                        }
563                        if !spec_cursor.goto_next_sibling() {
564                            break;
565                        }
566                    }
567                }
568            }
569        }
570        if !cursor.goto_next_sibling() {
571            break;
572        }
573    }
574}
575
576/// Extract the alias name from a namespace_import node (`* as name`).
577fn extract_namespace_import(source: &str, node: &Node, namespace_import: &mut Option<String>) {
578    let mut cursor = node.walk();
579    if !cursor.goto_first_child() {
580        return;
581    }
582
583    loop {
584        let child = cursor.node();
585        if child.kind() == "identifier" {
586            *namespace_import = Some(source[child.byte_range()].to_string());
587            return;
588        }
589        if !cursor.goto_next_sibling() {
590            break;
591        }
592    }
593}
594
595/// Generate an import line for TS/JS/TSX.
596fn generate_ts_import_line(
597    module_path: &str,
598    names: &[String],
599    default_import: Option<&str>,
600    type_only: bool,
601) -> String {
602    let type_prefix = if type_only { "type " } else { "" };
603
604    // Side-effect import
605    if names.is_empty() && default_import.is_none() {
606        return format!("import '{module_path}';");
607    }
608
609    // Default import only
610    if names.is_empty() {
611        if let Some(def) = default_import {
612            return format!("import {type_prefix}{def} from '{module_path}';");
613        }
614    }
615
616    // Named imports only
617    if default_import.is_none() {
618        let mut sorted_names = names.to_vec();
619        sorted_names.sort();
620        let names_str = sorted_names.join(", ");
621        return format!("import {type_prefix}{{ {names_str} }} from '{module_path}';");
622    }
623
624    // Both default and named imports
625    if let Some(def) = default_import {
626        let mut sorted_names = names.to_vec();
627        sorted_names.sort();
628        let names_str = sorted_names.join(", ");
629        return format!("import {type_prefix}{def}, {{ {names_str} }} from '{module_path}';");
630    }
631
632    // Shouldn't reach here, but handle gracefully
633    format!("import '{module_path}';")
634}
635
636// ---------------------------------------------------------------------------
637// Python implementation
638// ---------------------------------------------------------------------------
639
640/// Python 3.x standard library module names (top-level modules).
641/// Used for import group classification. Covers the commonly-used modules;
642/// unknown modules are assumed third-party.
643const PYTHON_STDLIB: &[&str] = &[
644    "__future__",
645    "_thread",
646    "abc",
647    "aifc",
648    "argparse",
649    "array",
650    "ast",
651    "asynchat",
652    "asyncio",
653    "asyncore",
654    "atexit",
655    "audioop",
656    "base64",
657    "bdb",
658    "binascii",
659    "bisect",
660    "builtins",
661    "bz2",
662    "calendar",
663    "cgi",
664    "cgitb",
665    "chunk",
666    "cmath",
667    "cmd",
668    "code",
669    "codecs",
670    "codeop",
671    "collections",
672    "colorsys",
673    "compileall",
674    "concurrent",
675    "configparser",
676    "contextlib",
677    "contextvars",
678    "copy",
679    "copyreg",
680    "cProfile",
681    "crypt",
682    "csv",
683    "ctypes",
684    "curses",
685    "dataclasses",
686    "datetime",
687    "dbm",
688    "decimal",
689    "difflib",
690    "dis",
691    "distutils",
692    "doctest",
693    "email",
694    "encodings",
695    "enum",
696    "errno",
697    "faulthandler",
698    "fcntl",
699    "filecmp",
700    "fileinput",
701    "fnmatch",
702    "fractions",
703    "ftplib",
704    "functools",
705    "gc",
706    "getopt",
707    "getpass",
708    "gettext",
709    "glob",
710    "grp",
711    "gzip",
712    "hashlib",
713    "heapq",
714    "hmac",
715    "html",
716    "http",
717    "idlelib",
718    "imaplib",
719    "imghdr",
720    "importlib",
721    "inspect",
722    "io",
723    "ipaddress",
724    "itertools",
725    "json",
726    "keyword",
727    "lib2to3",
728    "linecache",
729    "locale",
730    "logging",
731    "lzma",
732    "mailbox",
733    "mailcap",
734    "marshal",
735    "math",
736    "mimetypes",
737    "mmap",
738    "modulefinder",
739    "multiprocessing",
740    "netrc",
741    "numbers",
742    "operator",
743    "optparse",
744    "os",
745    "pathlib",
746    "pdb",
747    "pickle",
748    "pickletools",
749    "pipes",
750    "pkgutil",
751    "platform",
752    "plistlib",
753    "poplib",
754    "posixpath",
755    "pprint",
756    "profile",
757    "pstats",
758    "pty",
759    "pwd",
760    "py_compile",
761    "pyclbr",
762    "pydoc",
763    "queue",
764    "quopri",
765    "random",
766    "re",
767    "readline",
768    "reprlib",
769    "resource",
770    "rlcompleter",
771    "runpy",
772    "sched",
773    "secrets",
774    "select",
775    "selectors",
776    "shelve",
777    "shlex",
778    "shutil",
779    "signal",
780    "site",
781    "smtplib",
782    "sndhdr",
783    "socket",
784    "socketserver",
785    "sqlite3",
786    "ssl",
787    "stat",
788    "statistics",
789    "string",
790    "stringprep",
791    "struct",
792    "subprocess",
793    "symtable",
794    "sys",
795    "sysconfig",
796    "syslog",
797    "tabnanny",
798    "tarfile",
799    "tempfile",
800    "termios",
801    "textwrap",
802    "threading",
803    "time",
804    "timeit",
805    "tkinter",
806    "token",
807    "tokenize",
808    "tomllib",
809    "trace",
810    "traceback",
811    "tracemalloc",
812    "tty",
813    "turtle",
814    "types",
815    "typing",
816    "unicodedata",
817    "unittest",
818    "urllib",
819    "uuid",
820    "venv",
821    "warnings",
822    "wave",
823    "weakref",
824    "webbrowser",
825    "wsgiref",
826    "xml",
827    "xmlrpc",
828    "zipapp",
829    "zipfile",
830    "zipimport",
831    "zlib",
832];
833
834/// Classify a Python import into a group.
835pub fn classify_group_py(module_path: &str) -> ImportGroup {
836    // Relative imports start with '.'
837    if module_path.starts_with('.') {
838        return ImportGroup::Internal;
839    }
840    // Check stdlib: use the top-level module name (before first '.')
841    let top_module = module_path.split('.').next().unwrap_or(module_path);
842    if PYTHON_STDLIB.contains(&top_module) {
843        ImportGroup::Stdlib
844    } else {
845        ImportGroup::External
846    }
847}
848
849/// Parse imports from a Python file.
850fn parse_py_imports(source: &str, tree: &Tree) -> ImportBlock {
851    let root = tree.root_node();
852    let mut imports = Vec::new();
853
854    let mut cursor = root.walk();
855    if !cursor.goto_first_child() {
856        return ImportBlock::empty();
857    }
858
859    loop {
860        let node = cursor.node();
861        match node.kind() {
862            "import_statement" => {
863                if let Some(imp) = parse_py_import_statement(source, &node) {
864                    imports.push(imp);
865                }
866            }
867            "import_from_statement" => {
868                if let Some(imp) = parse_py_import_from_statement(source, &node) {
869                    imports.push(imp);
870                }
871            }
872            _ => {}
873        }
874        if !cursor.goto_next_sibling() {
875            break;
876        }
877    }
878
879    let byte_range = if imports.is_empty() {
880        None
881    } else {
882        let start = imports.first().unwrap().byte_range.start;
883        let end = imports.last().unwrap().byte_range.end;
884        Some(start..end)
885    };
886
887    ImportBlock {
888        imports,
889        byte_range,
890    }
891}
892
893/// Parse `import X` or `import X.Y` Python statements.
894fn parse_py_import_statement(source: &str, node: &Node) -> Option<ImportStatement> {
895    let raw_text = source[node.byte_range()].to_string();
896    let byte_range = node.byte_range();
897
898    // Find the dotted_name child (the module name)
899    let mut module_path = String::new();
900    let mut c = node.walk();
901    if c.goto_first_child() {
902        loop {
903            if c.node().kind() == "dotted_name" {
904                module_path = source[c.node().byte_range()].to_string();
905                break;
906            }
907            if !c.goto_next_sibling() {
908                break;
909            }
910        }
911    }
912    if module_path.is_empty() {
913        return None;
914    }
915
916    let group = classify_group_py(&module_path);
917
918    Some(ImportStatement {
919        module_path,
920        names: Vec::new(),
921        default_import: None,
922        namespace_import: None,
923        kind: ImportKind::Value,
924        group,
925        byte_range,
926        raw_text,
927    })
928}
929
930/// Parse `from X import Y, Z` or `from . import Y` Python statements.
931fn parse_py_import_from_statement(source: &str, node: &Node) -> Option<ImportStatement> {
932    let raw_text = source[node.byte_range()].to_string();
933    let byte_range = node.byte_range();
934
935    let mut module_path = String::new();
936    let mut names = Vec::new();
937
938    let mut c = node.walk();
939    if c.goto_first_child() {
940        loop {
941            let child = c.node();
942            match child.kind() {
943                "dotted_name" => {
944                    // Could be the module name or an imported name
945                    // The module name comes right after `from`, imported names come after `import`
946                    // Use position: if we haven't set module_path yet and this comes
947                    // before the `import` keyword, it's the module.
948                    if module_path.is_empty()
949                        && !has_seen_import_keyword(source, node, child.start_byte())
950                    {
951                        module_path = source[child.byte_range()].to_string();
952                    } else {
953                        // It's an imported name
954                        names.push(source[child.byte_range()].to_string());
955                    }
956                }
957                "relative_import" => {
958                    // from . import X or from ..module import X
959                    module_path = source[child.byte_range()].to_string();
960                }
961                _ => {}
962            }
963            if !c.goto_next_sibling() {
964                break;
965            }
966        }
967    }
968
969    // module_path must be non-empty for a valid import
970    if module_path.is_empty() {
971        return None;
972    }
973
974    let group = classify_group_py(&module_path);
975
976    Some(ImportStatement {
977        module_path,
978        names,
979        default_import: None,
980        namespace_import: None,
981        kind: ImportKind::Value,
982        group,
983        byte_range,
984        raw_text,
985    })
986}
987
988/// Check if the `import` keyword appears before the given byte position in a from...import node.
989fn has_seen_import_keyword(_source: &str, parent: &Node, before_byte: usize) -> bool {
990    let mut c = parent.walk();
991    if c.goto_first_child() {
992        loop {
993            let child = c.node();
994            if child.kind() == "import" && child.start_byte() < before_byte {
995                return true;
996            }
997            if child.start_byte() >= before_byte {
998                return false;
999            }
1000            if !c.goto_next_sibling() {
1001                break;
1002            }
1003        }
1004    }
1005    false
1006}
1007
1008/// Generate a Python import line.
1009fn generate_py_import_line(
1010    module_path: &str,
1011    names: &[String],
1012    _default_import: Option<&str>,
1013) -> String {
1014    if names.is_empty() {
1015        // `import module`
1016        format!("import {module_path}")
1017    } else {
1018        // `from module import name1, name2`
1019        let mut sorted = names.to_vec();
1020        sorted.sort();
1021        let names_str = sorted.join(", ");
1022        format!("from {module_path} import {names_str}")
1023    }
1024}
1025
1026// ---------------------------------------------------------------------------
1027// Rust implementation
1028// ---------------------------------------------------------------------------
1029
1030/// Classify a Rust use path into a group.
1031pub fn classify_group_rs(module_path: &str) -> ImportGroup {
1032    // Extract the first path segment (before ::)
1033    let first_seg = module_path.split("::").next().unwrap_or(module_path);
1034    match first_seg {
1035        "std" | "core" | "alloc" => ImportGroup::Stdlib,
1036        "crate" | "self" | "super" => ImportGroup::Internal,
1037        _ => ImportGroup::External,
1038    }
1039}
1040
1041/// Parse imports from a Rust file.
1042fn parse_rs_imports(source: &str, tree: &Tree) -> ImportBlock {
1043    let root = tree.root_node();
1044    let mut imports = Vec::new();
1045
1046    let mut cursor = root.walk();
1047    if !cursor.goto_first_child() {
1048        return ImportBlock::empty();
1049    }
1050
1051    loop {
1052        let node = cursor.node();
1053        if node.kind() == "use_declaration" {
1054            if let Some(imp) = parse_rs_use_declaration(source, &node) {
1055                imports.push(imp);
1056            }
1057        }
1058        if !cursor.goto_next_sibling() {
1059            break;
1060        }
1061    }
1062
1063    let byte_range = if imports.is_empty() {
1064        None
1065    } else {
1066        let start = imports.first().unwrap().byte_range.start;
1067        let end = imports.last().unwrap().byte_range.end;
1068        Some(start..end)
1069    };
1070
1071    ImportBlock {
1072        imports,
1073        byte_range,
1074    }
1075}
1076
1077/// Parse a single `use` declaration from Rust.
1078fn parse_rs_use_declaration(source: &str, node: &Node) -> Option<ImportStatement> {
1079    let raw_text = source[node.byte_range()].to_string();
1080    let byte_range = node.byte_range();
1081
1082    // Check for `pub` visibility modifier
1083    let mut has_pub = false;
1084    let mut use_path = String::new();
1085    let mut names = Vec::new();
1086
1087    let mut c = node.walk();
1088    if c.goto_first_child() {
1089        loop {
1090            let child = c.node();
1091            match child.kind() {
1092                "visibility_modifier" => {
1093                    has_pub = true;
1094                }
1095                "scoped_identifier" | "identifier" | "use_as_clause" => {
1096                    // Full path like `std::collections::HashMap` or just `serde`
1097                    use_path = source[child.byte_range()].to_string();
1098                }
1099                "scoped_use_list" => {
1100                    // e.g. `serde::{Deserialize, Serialize}`
1101                    use_path = source[child.byte_range()].to_string();
1102                    // Also extract the individual names from the use_list
1103                    extract_rs_use_list_names(source, &child, &mut names);
1104                }
1105                _ => {}
1106            }
1107            if !c.goto_next_sibling() {
1108                break;
1109            }
1110        }
1111    }
1112
1113    if use_path.is_empty() {
1114        return None;
1115    }
1116
1117    let group = classify_group_rs(&use_path);
1118
1119    Some(ImportStatement {
1120        module_path: use_path,
1121        names,
1122        default_import: if has_pub {
1123            Some("pub".to_string())
1124        } else {
1125            None
1126        },
1127        namespace_import: None,
1128        kind: ImportKind::Value,
1129        group,
1130        byte_range,
1131        raw_text,
1132    })
1133}
1134
1135/// Extract individual names from a Rust `scoped_use_list` node.
1136fn extract_rs_use_list_names(source: &str, node: &Node, names: &mut Vec<String>) {
1137    let mut c = node.walk();
1138    if c.goto_first_child() {
1139        loop {
1140            let child = c.node();
1141            if child.kind() == "use_list" {
1142                // Walk into the use_list to find identifiers
1143                let mut lc = child.walk();
1144                if lc.goto_first_child() {
1145                    loop {
1146                        let lchild = lc.node();
1147                        if lchild.kind() == "identifier" || lchild.kind() == "scoped_identifier" {
1148                            names.push(source[lchild.byte_range()].to_string());
1149                        }
1150                        if !lc.goto_next_sibling() {
1151                            break;
1152                        }
1153                    }
1154                }
1155            }
1156            if !c.goto_next_sibling() {
1157                break;
1158            }
1159        }
1160    }
1161}
1162
1163/// Generate a Rust import line.
1164fn generate_rs_import_line(module_path: &str, names: &[String], _type_only: bool) -> String {
1165    if names.is_empty() {
1166        format!("use {module_path};")
1167    } else {
1168        // If names are provided, generate `use prefix::{names};`
1169        // But the caller may pass module_path as the full path including the item,
1170        // e.g., "serde::Deserialize". For simple cases, just use the module_path directly.
1171        format!("use {module_path};")
1172    }
1173}
1174
1175// ---------------------------------------------------------------------------
1176// Go implementation
1177// ---------------------------------------------------------------------------
1178
1179/// Classify a Go import path into a group.
1180pub fn classify_group_go(module_path: &str) -> ImportGroup {
1181    // stdlib paths don't contain dots (e.g., "fmt", "os", "net/http")
1182    // external paths contain dots (e.g., "github.com/pkg/errors")
1183    if module_path.contains('.') {
1184        ImportGroup::External
1185    } else {
1186        ImportGroup::Stdlib
1187    }
1188}
1189
1190/// Parse imports from a Go file.
1191fn parse_go_imports(source: &str, tree: &Tree) -> ImportBlock {
1192    let root = tree.root_node();
1193    let mut imports = Vec::new();
1194
1195    let mut cursor = root.walk();
1196    if !cursor.goto_first_child() {
1197        return ImportBlock::empty();
1198    }
1199
1200    loop {
1201        let node = cursor.node();
1202        if node.kind() == "import_declaration" {
1203            parse_go_import_declaration(source, &node, &mut imports);
1204        }
1205        if !cursor.goto_next_sibling() {
1206            break;
1207        }
1208    }
1209
1210    let byte_range = if imports.is_empty() {
1211        None
1212    } else {
1213        let start = imports.first().unwrap().byte_range.start;
1214        let end = imports.last().unwrap().byte_range.end;
1215        Some(start..end)
1216    };
1217
1218    ImportBlock {
1219        imports,
1220        byte_range,
1221    }
1222}
1223
1224/// Parse a single Go import_declaration (may contain one or multiple specs).
1225fn parse_go_import_declaration(source: &str, node: &Node, imports: &mut Vec<ImportStatement>) {
1226    let mut c = node.walk();
1227    if c.goto_first_child() {
1228        loop {
1229            let child = c.node();
1230            match child.kind() {
1231                "import_spec" => {
1232                    if let Some(imp) = parse_go_import_spec(source, &child) {
1233                        imports.push(imp);
1234                    }
1235                }
1236                "import_spec_list" => {
1237                    // Grouped imports: walk into the list
1238                    let mut lc = child.walk();
1239                    if lc.goto_first_child() {
1240                        loop {
1241                            if lc.node().kind() == "import_spec" {
1242                                if let Some(imp) = parse_go_import_spec(source, &lc.node()) {
1243                                    imports.push(imp);
1244                                }
1245                            }
1246                            if !lc.goto_next_sibling() {
1247                                break;
1248                            }
1249                        }
1250                    }
1251                }
1252                _ => {}
1253            }
1254            if !c.goto_next_sibling() {
1255                break;
1256            }
1257        }
1258    }
1259}
1260
1261/// Parse a single Go import_spec node.
1262fn parse_go_import_spec(source: &str, node: &Node) -> Option<ImportStatement> {
1263    let raw_text = source[node.byte_range()].to_string();
1264    let byte_range = node.byte_range();
1265
1266    let mut import_path = String::new();
1267    let mut alias = None;
1268
1269    let mut c = node.walk();
1270    if c.goto_first_child() {
1271        loop {
1272            let child = c.node();
1273            match child.kind() {
1274                "interpreted_string_literal" => {
1275                    // Extract the path without quotes
1276                    let text = source[child.byte_range()].to_string();
1277                    import_path = text.trim_matches('"').to_string();
1278                }
1279                "identifier" | "blank_identifier" | "dot" => {
1280                    // This is an alias (e.g., `alias "path"` or `. "path"` or `_ "path"`)
1281                    alias = Some(source[child.byte_range()].to_string());
1282                }
1283                _ => {}
1284            }
1285            if !c.goto_next_sibling() {
1286                break;
1287            }
1288        }
1289    }
1290
1291    if import_path.is_empty() {
1292        return None;
1293    }
1294
1295    let group = classify_group_go(&import_path);
1296
1297    Some(ImportStatement {
1298        module_path: import_path,
1299        names: Vec::new(),
1300        default_import: alias,
1301        namespace_import: None,
1302        kind: ImportKind::Value,
1303        group,
1304        byte_range,
1305        raw_text,
1306    })
1307}
1308
1309/// Public API for Go import line generation (used by add_import handler).
1310pub fn generate_go_import_line_pub(
1311    module_path: &str,
1312    alias: Option<&str>,
1313    in_group: bool,
1314) -> String {
1315    generate_go_import_line(module_path, alias, in_group)
1316}
1317
1318/// Generate a Go import line (public API for command handler).
1319///
1320/// `in_group` controls whether to generate a spec for insertion into an
1321/// existing grouped import (`\t"path"`) or a standalone import (`import "path"`).
1322fn generate_go_import_line(module_path: &str, alias: Option<&str>, in_group: bool) -> String {
1323    if in_group {
1324        // Spec for grouped import block
1325        match alias {
1326            Some(a) => format!("\t{a} \"{module_path}\""),
1327            None => format!("\t\"{module_path}\""),
1328        }
1329    } else {
1330        // Standalone import
1331        match alias {
1332            Some(a) => format!("import {a} \"{module_path}\""),
1333            None => format!("import \"{module_path}\""),
1334        }
1335    }
1336}
1337
1338/// Check if a Go import block has a grouped import declaration.
1339/// Returns the byte range of the import_spec_list if found.
1340pub fn go_has_grouped_import(_source: &str, tree: &Tree) -> Option<Range<usize>> {
1341    let root = tree.root_node();
1342    let mut cursor = root.walk();
1343    if !cursor.goto_first_child() {
1344        return None;
1345    }
1346
1347    loop {
1348        let node = cursor.node();
1349        if node.kind() == "import_declaration" {
1350            let mut c = node.walk();
1351            if c.goto_first_child() {
1352                loop {
1353                    if c.node().kind() == "import_spec_list" {
1354                        return Some(c.node().byte_range());
1355                    }
1356                    if !c.goto_next_sibling() {
1357                        break;
1358                    }
1359                }
1360            }
1361        }
1362        if !cursor.goto_next_sibling() {
1363            break;
1364        }
1365    }
1366    None
1367}
1368
1369/// Skip past a newline character at the given position.
1370fn skip_newline(source: &str, pos: usize) -> usize {
1371    if pos < source.len() {
1372        let bytes = source.as_bytes();
1373        if bytes[pos] == b'\n' {
1374            return pos + 1;
1375        }
1376        if bytes[pos] == b'\r' {
1377            if pos + 1 < source.len() && bytes[pos + 1] == b'\n' {
1378                return pos + 2;
1379            }
1380            return pos + 1;
1381        }
1382    }
1383    pos
1384}
1385
1386// ---------------------------------------------------------------------------
1387// Unit tests
1388// ---------------------------------------------------------------------------
1389
1390#[cfg(test)]
1391mod tests {
1392    use super::*;
1393
1394    fn parse_ts(source: &str) -> (Tree, ImportBlock) {
1395        let grammar = grammar_for(LangId::TypeScript);
1396        let mut parser = Parser::new();
1397        parser.set_language(&grammar).unwrap();
1398        let tree = parser.parse(source, None).unwrap();
1399        let block = parse_imports(source, &tree, LangId::TypeScript);
1400        (tree, block)
1401    }
1402
1403    fn parse_js(source: &str) -> (Tree, ImportBlock) {
1404        let grammar = grammar_for(LangId::JavaScript);
1405        let mut parser = Parser::new();
1406        parser.set_language(&grammar).unwrap();
1407        let tree = parser.parse(source, None).unwrap();
1408        let block = parse_imports(source, &tree, LangId::JavaScript);
1409        (tree, block)
1410    }
1411
1412    // --- Basic parsing ---
1413
1414    #[test]
1415    fn parse_ts_named_imports() {
1416        let source = "import { useState, useEffect } from 'react';\n";
1417        let (_, block) = parse_ts(source);
1418        assert_eq!(block.imports.len(), 1);
1419        let imp = &block.imports[0];
1420        assert_eq!(imp.module_path, "react");
1421        assert!(imp.names.contains(&"useState".to_string()));
1422        assert!(imp.names.contains(&"useEffect".to_string()));
1423        assert_eq!(imp.kind, ImportKind::Value);
1424        assert_eq!(imp.group, ImportGroup::External);
1425    }
1426
1427    #[test]
1428    fn parse_ts_default_import() {
1429        let source = "import React from 'react';\n";
1430        let (_, block) = parse_ts(source);
1431        assert_eq!(block.imports.len(), 1);
1432        let imp = &block.imports[0];
1433        assert_eq!(imp.default_import.as_deref(), Some("React"));
1434        assert_eq!(imp.kind, ImportKind::Value);
1435    }
1436
1437    #[test]
1438    fn parse_ts_side_effect_import() {
1439        let source = "import './styles.css';\n";
1440        let (_, block) = parse_ts(source);
1441        assert_eq!(block.imports.len(), 1);
1442        assert_eq!(block.imports[0].kind, ImportKind::SideEffect);
1443        assert_eq!(block.imports[0].module_path, "./styles.css");
1444    }
1445
1446    #[test]
1447    fn parse_ts_relative_import() {
1448        let source = "import { helper } from './utils';\n";
1449        let (_, block) = parse_ts(source);
1450        assert_eq!(block.imports.len(), 1);
1451        assert_eq!(block.imports[0].group, ImportGroup::Internal);
1452    }
1453
1454    #[test]
1455    fn parse_ts_multiple_groups() {
1456        let source = "\
1457import React from 'react';
1458import { useState } from 'react';
1459import { helper } from './utils';
1460import { Config } from '../config';
1461";
1462        let (_, block) = parse_ts(source);
1463        assert_eq!(block.imports.len(), 4);
1464
1465        let external: Vec<_> = block
1466            .imports
1467            .iter()
1468            .filter(|i| i.group == ImportGroup::External)
1469            .collect();
1470        let relative: Vec<_> = block
1471            .imports
1472            .iter()
1473            .filter(|i| i.group == ImportGroup::Internal)
1474            .collect();
1475        assert_eq!(external.len(), 2);
1476        assert_eq!(relative.len(), 2);
1477    }
1478
1479    #[test]
1480    fn parse_ts_namespace_import() {
1481        let source = "import * as path from 'path';\n";
1482        let (_, block) = parse_ts(source);
1483        assert_eq!(block.imports.len(), 1);
1484        let imp = &block.imports[0];
1485        assert_eq!(imp.namespace_import.as_deref(), Some("path"));
1486        assert_eq!(imp.kind, ImportKind::Value);
1487    }
1488
1489    #[test]
1490    fn parse_js_imports() {
1491        let source = "import { readFile } from 'fs';\nimport { helper } from './helper';\n";
1492        let (_, block) = parse_js(source);
1493        assert_eq!(block.imports.len(), 2);
1494        assert_eq!(block.imports[0].group, ImportGroup::External);
1495        assert_eq!(block.imports[1].group, ImportGroup::Internal);
1496    }
1497
1498    // --- Group classification ---
1499
1500    #[test]
1501    fn classify_external() {
1502        assert_eq!(classify_group_ts("react"), ImportGroup::External);
1503        assert_eq!(classify_group_ts("@scope/pkg"), ImportGroup::External);
1504        assert_eq!(classify_group_ts("lodash/map"), ImportGroup::External);
1505    }
1506
1507    #[test]
1508    fn classify_relative() {
1509        assert_eq!(classify_group_ts("./utils"), ImportGroup::Internal);
1510        assert_eq!(classify_group_ts("../config"), ImportGroup::Internal);
1511        assert_eq!(classify_group_ts("./"), ImportGroup::Internal);
1512    }
1513
1514    // --- Dedup ---
1515
1516    #[test]
1517    fn dedup_detects_same_named_import() {
1518        let source = "import { useState } from 'react';\n";
1519        let (_, block) = parse_ts(source);
1520        assert!(is_duplicate(
1521            &block,
1522            "react",
1523            &["useState".to_string()],
1524            None,
1525            false
1526        ));
1527    }
1528
1529    #[test]
1530    fn dedup_misses_different_name() {
1531        let source = "import { useState } from 'react';\n";
1532        let (_, block) = parse_ts(source);
1533        assert!(!is_duplicate(
1534            &block,
1535            "react",
1536            &["useEffect".to_string()],
1537            None,
1538            false
1539        ));
1540    }
1541
1542    #[test]
1543    fn dedup_detects_default_import() {
1544        let source = "import React from 'react';\n";
1545        let (_, block) = parse_ts(source);
1546        assert!(is_duplicate(&block, "react", &[], Some("React"), false));
1547    }
1548
1549    #[test]
1550    fn dedup_side_effect() {
1551        let source = "import './styles.css';\n";
1552        let (_, block) = parse_ts(source);
1553        assert!(is_duplicate(&block, "./styles.css", &[], None, false));
1554    }
1555
1556    #[test]
1557    fn dedup_type_vs_value() {
1558        let source = "import { FC } from 'react';\n";
1559        let (_, block) = parse_ts(source);
1560        // Type import should NOT match a value import of the same name
1561        assert!(!is_duplicate(
1562            &block,
1563            "react",
1564            &["FC".to_string()],
1565            None,
1566            true
1567        ));
1568    }
1569
1570    // --- Generation ---
1571
1572    #[test]
1573    fn generate_named_import() {
1574        let line = generate_import_line(
1575            LangId::TypeScript,
1576            "react",
1577            &["useState".to_string(), "useEffect".to_string()],
1578            None,
1579            false,
1580        );
1581        assert_eq!(line, "import { useEffect, useState } from 'react';");
1582    }
1583
1584    #[test]
1585    fn generate_default_import() {
1586        let line = generate_import_line(LangId::TypeScript, "react", &[], Some("React"), false);
1587        assert_eq!(line, "import React from 'react';");
1588    }
1589
1590    #[test]
1591    fn generate_type_import() {
1592        let line =
1593            generate_import_line(LangId::TypeScript, "react", &["FC".to_string()], None, true);
1594        assert_eq!(line, "import type { FC } from 'react';");
1595    }
1596
1597    #[test]
1598    fn generate_side_effect_import() {
1599        let line = generate_import_line(LangId::TypeScript, "./styles.css", &[], None, false);
1600        assert_eq!(line, "import './styles.css';");
1601    }
1602
1603    #[test]
1604    fn generate_default_and_named() {
1605        let line = generate_import_line(
1606            LangId::TypeScript,
1607            "react",
1608            &["useState".to_string()],
1609            Some("React"),
1610            false,
1611        );
1612        assert_eq!(line, "import React, { useState } from 'react';");
1613    }
1614
1615    #[test]
1616    fn parse_ts_type_import() {
1617        let source = "import type { FC } from 'react';\n";
1618        let (_, block) = parse_ts(source);
1619        assert_eq!(block.imports.len(), 1);
1620        let imp = &block.imports[0];
1621        assert_eq!(imp.kind, ImportKind::Type);
1622        assert!(imp.names.contains(&"FC".to_string()));
1623        assert_eq!(imp.group, ImportGroup::External);
1624    }
1625
1626    // --- Insertion point ---
1627
1628    #[test]
1629    fn insertion_empty_file() {
1630        let source = "";
1631        let (_, block) = parse_ts(source);
1632        let (offset, _, _) =
1633            find_insertion_point(source, &block, ImportGroup::External, "react", false);
1634        assert_eq!(offset, 0);
1635    }
1636
1637    #[test]
1638    fn insertion_alphabetical_within_group() {
1639        let source = "\
1640import { a } from 'alpha';
1641import { c } from 'charlie';
1642";
1643        let (_, block) = parse_ts(source);
1644        let (offset, _, _) =
1645            find_insertion_point(source, &block, ImportGroup::External, "bravo", false);
1646        // Should insert before 'charlie' (which starts at line 2)
1647        let before_charlie = source.find("import { c }").unwrap();
1648        assert_eq!(offset, before_charlie);
1649    }
1650
1651    // --- Python parsing ---
1652
1653    fn parse_py(source: &str) -> (Tree, ImportBlock) {
1654        let grammar = grammar_for(LangId::Python);
1655        let mut parser = Parser::new();
1656        parser.set_language(&grammar).unwrap();
1657        let tree = parser.parse(source, None).unwrap();
1658        let block = parse_imports(source, &tree, LangId::Python);
1659        (tree, block)
1660    }
1661
1662    #[test]
1663    fn parse_py_import_statement() {
1664        let source = "import os\nimport sys\n";
1665        let (_, block) = parse_py(source);
1666        assert_eq!(block.imports.len(), 2);
1667        assert_eq!(block.imports[0].module_path, "os");
1668        assert_eq!(block.imports[1].module_path, "sys");
1669        assert_eq!(block.imports[0].group, ImportGroup::Stdlib);
1670    }
1671
1672    #[test]
1673    fn parse_py_from_import() {
1674        let source = "from collections import OrderedDict\nfrom typing import List, Optional\n";
1675        let (_, block) = parse_py(source);
1676        assert_eq!(block.imports.len(), 2);
1677        assert_eq!(block.imports[0].module_path, "collections");
1678        assert!(block.imports[0].names.contains(&"OrderedDict".to_string()));
1679        assert_eq!(block.imports[0].group, ImportGroup::Stdlib);
1680        assert_eq!(block.imports[1].module_path, "typing");
1681        assert!(block.imports[1].names.contains(&"List".to_string()));
1682        assert!(block.imports[1].names.contains(&"Optional".to_string()));
1683    }
1684
1685    #[test]
1686    fn parse_py_relative_import() {
1687        let source = "from . import utils\nfrom ..config import Settings\n";
1688        let (_, block) = parse_py(source);
1689        assert_eq!(block.imports.len(), 2);
1690        assert_eq!(block.imports[0].module_path, ".");
1691        assert!(block.imports[0].names.contains(&"utils".to_string()));
1692        assert_eq!(block.imports[0].group, ImportGroup::Internal);
1693        assert_eq!(block.imports[1].module_path, "..config");
1694        assert_eq!(block.imports[1].group, ImportGroup::Internal);
1695    }
1696
1697    #[test]
1698    fn classify_py_groups() {
1699        assert_eq!(classify_group_py("os"), ImportGroup::Stdlib);
1700        assert_eq!(classify_group_py("sys"), ImportGroup::Stdlib);
1701        assert_eq!(classify_group_py("json"), ImportGroup::Stdlib);
1702        assert_eq!(classify_group_py("collections"), ImportGroup::Stdlib);
1703        assert_eq!(classify_group_py("os.path"), ImportGroup::Stdlib);
1704        assert_eq!(classify_group_py("requests"), ImportGroup::External);
1705        assert_eq!(classify_group_py("flask"), ImportGroup::External);
1706        assert_eq!(classify_group_py("."), ImportGroup::Internal);
1707        assert_eq!(classify_group_py("..config"), ImportGroup::Internal);
1708        assert_eq!(classify_group_py(".utils"), ImportGroup::Internal);
1709    }
1710
1711    #[test]
1712    fn parse_py_three_groups() {
1713        let source = "import os\nimport sys\n\nimport requests\n\nfrom . import utils\n";
1714        let (_, block) = parse_py(source);
1715        let stdlib: Vec<_> = block
1716            .imports
1717            .iter()
1718            .filter(|i| i.group == ImportGroup::Stdlib)
1719            .collect();
1720        let external: Vec<_> = block
1721            .imports
1722            .iter()
1723            .filter(|i| i.group == ImportGroup::External)
1724            .collect();
1725        let internal: Vec<_> = block
1726            .imports
1727            .iter()
1728            .filter(|i| i.group == ImportGroup::Internal)
1729            .collect();
1730        assert_eq!(stdlib.len(), 2);
1731        assert_eq!(external.len(), 1);
1732        assert_eq!(internal.len(), 1);
1733    }
1734
1735    #[test]
1736    fn generate_py_import() {
1737        let line = generate_import_line(LangId::Python, "os", &[], None, false);
1738        assert_eq!(line, "import os");
1739    }
1740
1741    #[test]
1742    fn generate_py_from_import() {
1743        let line = generate_import_line(
1744            LangId::Python,
1745            "collections",
1746            &["OrderedDict".to_string()],
1747            None,
1748            false,
1749        );
1750        assert_eq!(line, "from collections import OrderedDict");
1751    }
1752
1753    #[test]
1754    fn generate_py_from_import_multiple() {
1755        let line = generate_import_line(
1756            LangId::Python,
1757            "typing",
1758            &["Optional".to_string(), "List".to_string()],
1759            None,
1760            false,
1761        );
1762        assert_eq!(line, "from typing import List, Optional");
1763    }
1764
1765    // --- Rust parsing ---
1766
1767    fn parse_rust(source: &str) -> (Tree, ImportBlock) {
1768        let grammar = grammar_for(LangId::Rust);
1769        let mut parser = Parser::new();
1770        parser.set_language(&grammar).unwrap();
1771        let tree = parser.parse(source, None).unwrap();
1772        let block = parse_imports(source, &tree, LangId::Rust);
1773        (tree, block)
1774    }
1775
1776    #[test]
1777    fn parse_rs_use_std() {
1778        let source = "use std::collections::HashMap;\nuse std::io::Read;\n";
1779        let (_, block) = parse_rust(source);
1780        assert_eq!(block.imports.len(), 2);
1781        assert_eq!(block.imports[0].module_path, "std::collections::HashMap");
1782        assert_eq!(block.imports[0].group, ImportGroup::Stdlib);
1783        assert_eq!(block.imports[1].group, ImportGroup::Stdlib);
1784    }
1785
1786    #[test]
1787    fn parse_rs_use_external() {
1788        let source = "use serde::{Deserialize, Serialize};\n";
1789        let (_, block) = parse_rust(source);
1790        assert_eq!(block.imports.len(), 1);
1791        assert_eq!(block.imports[0].group, ImportGroup::External);
1792        assert!(block.imports[0].names.contains(&"Deserialize".to_string()));
1793        assert!(block.imports[0].names.contains(&"Serialize".to_string()));
1794    }
1795
1796    #[test]
1797    fn parse_rs_use_crate() {
1798        let source = "use crate::config::Settings;\nuse super::parent::Thing;\n";
1799        let (_, block) = parse_rust(source);
1800        assert_eq!(block.imports.len(), 2);
1801        assert_eq!(block.imports[0].group, ImportGroup::Internal);
1802        assert_eq!(block.imports[1].group, ImportGroup::Internal);
1803    }
1804
1805    #[test]
1806    fn parse_rs_pub_use() {
1807        let source = "pub use super::parent::Thing;\n";
1808        let (_, block) = parse_rust(source);
1809        assert_eq!(block.imports.len(), 1);
1810        // `pub` is stored in default_import as a marker
1811        assert_eq!(block.imports[0].default_import.as_deref(), Some("pub"));
1812    }
1813
1814    #[test]
1815    fn classify_rs_groups() {
1816        assert_eq!(
1817            classify_group_rs("std::collections::HashMap"),
1818            ImportGroup::Stdlib
1819        );
1820        assert_eq!(classify_group_rs("core::mem"), ImportGroup::Stdlib);
1821        assert_eq!(classify_group_rs("alloc::vec"), ImportGroup::Stdlib);
1822        assert_eq!(
1823            classify_group_rs("serde::Deserialize"),
1824            ImportGroup::External
1825        );
1826        assert_eq!(classify_group_rs("tokio::runtime"), ImportGroup::External);
1827        assert_eq!(classify_group_rs("crate::config"), ImportGroup::Internal);
1828        assert_eq!(classify_group_rs("self::utils"), ImportGroup::Internal);
1829        assert_eq!(classify_group_rs("super::parent"), ImportGroup::Internal);
1830    }
1831
1832    #[test]
1833    fn generate_rs_use() {
1834        let line = generate_import_line(LangId::Rust, "std::fmt::Display", &[], None, false);
1835        assert_eq!(line, "use std::fmt::Display;");
1836    }
1837
1838    // --- Go parsing ---
1839
1840    fn parse_go(source: &str) -> (Tree, ImportBlock) {
1841        let grammar = grammar_for(LangId::Go);
1842        let mut parser = Parser::new();
1843        parser.set_language(&grammar).unwrap();
1844        let tree = parser.parse(source, None).unwrap();
1845        let block = parse_imports(source, &tree, LangId::Go);
1846        (tree, block)
1847    }
1848
1849    #[test]
1850    fn parse_go_single_import() {
1851        let source = "package main\n\nimport \"fmt\"\n";
1852        let (_, block) = parse_go(source);
1853        assert_eq!(block.imports.len(), 1);
1854        assert_eq!(block.imports[0].module_path, "fmt");
1855        assert_eq!(block.imports[0].group, ImportGroup::Stdlib);
1856    }
1857
1858    #[test]
1859    fn parse_go_grouped_import() {
1860        let source =
1861            "package main\n\nimport (\n\t\"fmt\"\n\t\"os\"\n\n\t\"github.com/pkg/errors\"\n)\n";
1862        let (_, block) = parse_go(source);
1863        assert_eq!(block.imports.len(), 3);
1864        assert_eq!(block.imports[0].module_path, "fmt");
1865        assert_eq!(block.imports[0].group, ImportGroup::Stdlib);
1866        assert_eq!(block.imports[1].module_path, "os");
1867        assert_eq!(block.imports[1].group, ImportGroup::Stdlib);
1868        assert_eq!(block.imports[2].module_path, "github.com/pkg/errors");
1869        assert_eq!(block.imports[2].group, ImportGroup::External);
1870    }
1871
1872    #[test]
1873    fn parse_go_mixed_imports() {
1874        // Single + grouped
1875        let source = "package main\n\nimport \"fmt\"\n\nimport (\n\t\"os\"\n\t\"github.com/pkg/errors\"\n)\n";
1876        let (_, block) = parse_go(source);
1877        assert_eq!(block.imports.len(), 3);
1878    }
1879
1880    #[test]
1881    fn classify_go_groups() {
1882        assert_eq!(classify_group_go("fmt"), ImportGroup::Stdlib);
1883        assert_eq!(classify_group_go("os"), ImportGroup::Stdlib);
1884        assert_eq!(classify_group_go("net/http"), ImportGroup::Stdlib);
1885        assert_eq!(classify_group_go("encoding/json"), ImportGroup::Stdlib);
1886        assert_eq!(
1887            classify_group_go("github.com/pkg/errors"),
1888            ImportGroup::External
1889        );
1890        assert_eq!(
1891            classify_group_go("golang.org/x/tools"),
1892            ImportGroup::External
1893        );
1894    }
1895
1896    #[test]
1897    fn generate_go_standalone() {
1898        let line = generate_go_import_line("fmt", None, false);
1899        assert_eq!(line, "import \"fmt\"");
1900    }
1901
1902    #[test]
1903    fn generate_go_grouped_spec() {
1904        let line = generate_go_import_line("fmt", None, true);
1905        assert_eq!(line, "\t\"fmt\"");
1906    }
1907
1908    #[test]
1909    fn generate_go_with_alias() {
1910        let line = generate_go_import_line("github.com/pkg/errors", Some("errs"), false);
1911        assert_eq!(line, "import errs \"github.com/pkg/errors\"");
1912    }
1913}