#!/usr/bin/env python3
"""
PyTorch Block Extractor — Production Drop-in (index-once, batch-by-default)

What this script does:
- Parallel parsing (LibCST→AST), persistent SQLite index (per-file hash)
- Module import graph + recursive dependency closure
- Real definition fetching via DefinitionResolver (no fabricated code)
- Import-aware, scope-accurate free-name analysis (handles lambdas, comps, walrus)
- Indexes & emits module-level constants/aliases as first-class symbols
- Emits original import lines from all contributing modules (so stdlib/3p deps resolve naturally)
- Systematic package→repo mapping computed from the index (no repo hardcoding)
- **Cold start minimized**: index is warmed once per process; default indexing policy = "missing" (only repos not yet in the SQLite index)
- **Batch mode by default**: if --block / --blocks not provided, reads names from nn_block_names.json and extracts one by one
- **Resume & controls**: --redo-existing, --limit, --start-from, --stop-on-fail, --progress-every
"""

from __future__ import annotations

import argparse
import ast
import concurrent.futures as cf
import hashlib
import importlib
import json
import logging
import os
import re
import sqlite3
import threading
import time
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple

# Import the block validator
from .block_validator import BlockValidator

# ----------------------------------------------------------------------------- #
# Internal, deterministic sanitizer (no stub injection)
# ----------------------------------------------------------------------------- #
def sanitize_generated_block(text: str) -> str:
    """
    Enhanced sanitizer that also adds missing imports:
    - normalizes newlines
    - collapses >2 consecutive blank lines to 1
    - de-dupes identical import lines (keeps first occurrence)
    - strips trailing spaces
    - detects undefined symbols and adds necessary imports
    """
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    text = re.sub(r"\n{3,}", "\n\n", text)
    lines = text.split("\n")
    seen_imports: Set[str] = set()
    out: List[str] = []
    
    # First pass: collect existing imports and find undefined symbols
    defined_symbols = set()
    used_symbols = set()
    
    for ln in lines:
        s = ln.rstrip()
        if s.startswith("import ") or s.startswith("from "):
            if s in seen_imports:
                continue
            seen_imports.add(s)
            # Extract imported symbols from import statements
            if s.startswith("from ") and " import " in s:
                # from module import symbol1, symbol2
                parts = s.split(" import ")
                if len(parts) == 2:
                    symbols = [sym.strip() for sym in parts[1].split(",")]
                    defined_symbols.update(symbols)
            elif s.startswith("import ") and " as " in s:
                # import module as alias
                parts = s.split(" as ")
                if len(parts) == 2:
                    defined_symbols.add(parts[1].strip())
            elif s.startswith("import ") and " as " not in s:
                # import module
                module = s.split("import ")[1].strip()
                defined_symbols.add(module)
        
        # Detect class and function definitions
        if s.startswith("class ") or s.startswith("def "):
            # Extract name from "class ClassName" or "def function_name"
            parts = s.split()
            if len(parts) >= 2:
                name = parts[1].split("(")[0].split(":")[0]
                defined_symbols.add(name)
        
        # Detect variable assignments
        if " = " in s and not s.startswith(" ") and not s.startswith("#"):
            # Simple variable assignment (not indented, not comment)
            name = s.split(" = ")[0].strip()
            if name and name.isidentifier():
                defined_symbols.add(name)
        
        out.append(s)
    
    # Second pass: detect undefined symbols
    for ln in lines:
        s = ln.rstrip()
        if not s.startswith("import ") and not s.startswith("from ") and not s.startswith("#") and not s.startswith("class ") and not s.startswith("def "):
            # Look for symbol usage in this line
            words = s.split()
            for word in words:
                # Clean the word (remove punctuation, etc.)
                clean_word = re.sub(r'[^\w]', '', word)
                if clean_word and clean_word.isidentifier() and len(clean_word) > 1:
                    # Check if it's a built-in or already defined
                    if clean_word not in defined_symbols and clean_word not in {
                        "self", "cls", "x", "y", "z", "i", "j", "k", "n", "m", "p", "q", "r", "s", "t", "u", "v", "w",
                        "torch", "nn", "F", "os", "math", "copy", "defaultdict", "logging", "osp", "MODELS", "OptConfigType", "MultiConfig", "ConfigDict"
                    }:
                        used_symbols.add(clean_word)
    
    # Add missing imports for undefined symbols
    missing_imports = []
    
    # Handle common patterns
    if "Swish" in used_symbols and "Swish" not in defined_symbols:
        missing_imports.append("from mmcv.cnn.bricks import Swish")
    if "build_norm_layer" in used_symbols and "build_norm_layer" not in defined_symbols:
        missing_imports.append("from mmcv.cnn.bricks import build_norm_layer")
    if "Conv2dSamePadding" in used_symbols and "Conv2dSamePadding" not in defined_symbols:
        missing_imports.append("from projects.EfficientDet.efficientdet.utils import Conv2dSamePadding")
    if "MaxPool2dSamePadding" in used_symbols and "MaxPool2dSamePadding" not in defined_symbols:
        missing_imports.append("from projects.EfficientDet.efficientdet.utils import MaxPool2dSamePadding")
    if "BiFPNStage" in used_symbols and "BiFPNStage" not in defined_symbols:
        missing_imports.append("from projects.EfficientDet.efficientdet.bifpn import BiFPNStage")
    if "DepthWiseConvBlock" in used_symbols and "DepthWiseConvBlock" not in defined_symbols:
        missing_imports.append("from projects.EfficientDet.efficientdet.bifpn import DepthWiseConvBlock")
    if "DownChannelBlock" in used_symbols and "DownChannelBlock" not in defined_symbols:
        missing_imports.append("from projects.EfficientDet.efficientdet.bifpn import DownChannelBlock")
    if "Sequence" in used_symbols and "Sequence" not in defined_symbols:
        missing_imports.append("from collections.abc import Sequence")
    
    # Insert missing imports after the header comment
    if missing_imports:
        for i, line in enumerate(out):
            if line.startswith("# Auto-generated single-file for"):
                # Insert imports after the header
                for j, imp in enumerate(missing_imports):
                    out.insert(i + 2 + j, imp)
                break
    
    return "\n".join(out).rstrip() + "\n"


# Real definition resolver (returns real source slices; never stubs)
from .utils.definition_resolver import DefinitionResolver, ResolvedSymbol

# Optional LibCST for fast/precise module parsing
try:
    import libcst as cst
    from libcst.metadata import MetadataWrapper, PositionProvider
    LIBCST_AVAILABLE = True
except Exception:
    LIBCST_AVAILABLE = False

# Project-local helpers
from .utils.package_generator import PackageGenerator
from .utils.repo_cache import RepoCache


# ----------------------------------------------------------------------------- #
# Logging
# ----------------------------------------------------------------------------- #
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-7s | %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("extractor")


# ----------------------------------------------------------------------------- #
# Data models
# ----------------------------------------------------------------------------- #
@dataclass
class SymbolInfo:
    name: str
    qualified_name: str         # e.g. "timm.layers.attention.Attention" or "timm.layers.config._USE_REENTRANT_CKPT"
    kind: str                   # 'class' | 'function' | 'const'
    location: str               # repo-relative file path (with .py)
    line_number: int
    source_code: str


@dataclass
class ModuleInfo:
    repo: str
    qual: str                   # dotted module path e.g. "timm.layers.attention"
    path: str                   # repo-relative path without suffix e.g. "timm/layers/attention"
    file: str                   # filename (module leaf)
    source_code: str = ""       # full module text (subset retained in SQLite too)
    all_exports: List[str] = field(default_factory=list)
    imports: List[str] = field(default_factory=list)
    symbols: Dict[str, SymbolInfo] = field(default_factory=dict)   # name -> SymbolInfo


@dataclass
class ImportGraph:
    modules: Dict[str, ModuleInfo] = field(default_factory=dict)         # key: "repo:qual"
    symbol_table: Dict[str, SymbolInfo] = field(default_factory=dict)    # qual.name -> SymbolInfo
    dependents: Dict[str, List[str]] = field(default_factory=dict)       # imported_module -> [module_keys]


@dataclass
class ResolvedDependency:
    name: str
    qualified_name: str
    source_code: str
    resolution_method: str
    confidence: float
    location: Optional[str] = None


@dataclass
class DependencyResolutionResult:
    target_symbol: str
    resolved_dependencies: Dict[str, ResolvedDependency]
    unresolved_dependencies: List[str]
    import_graph: ImportGraph
    resolution_stats: Dict[str, Any]
    topological_order: Optional[List[str]] = None


# ----------------------------------------------------------------------------- #
# Persistent index (SQLite)
# ----------------------------------------------------------------------------- #
class FileIndexStore:
    def __init__(self, db_path: Path = Path(".cache/index.sqlite3")):
        self.db_path = db_path
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = threading.Lock()
        self._init_db()

    def _init_db(self) -> None:
        with sqlite3.connect(self.db_path) as con:
            con.execute("""
                CREATE TABLE IF NOT EXISTS files (
                    repo TEXT NOT NULL,
                    relpath TEXT NOT NULL,
                    mod_qual TEXT NOT NULL,
                    sha1 TEXT NOT NULL,
                    imports_json TEXT NOT NULL,
                    source TEXT,
                    PRIMARY KEY (repo, relpath, sha1)
                )
            """)
            con.execute("""
                CREATE TABLE IF NOT EXISTS symbols (
                    repo TEXT NOT NULL,
                    relpath TEXT NOT NULL,
                    sha1 TEXT NOT NULL,
                    name TEXT NOT NULL,
                    qual TEXT NOT NULL,
                    kind TEXT NOT NULL,
                    line INTEGER NOT NULL,
                    code TEXT NOT NULL,
                    PRIMARY KEY (repo, relpath, sha1, qual)
                )
            """)
            con.execute("""
                CREATE TABLE IF NOT EXISTS repo_index_state (
                    repo TEXT PRIMARY KEY,
                    indexed_at TEXT NOT NULL,
                    file_count INTEGER NOT NULL
                )
            """)
            con.commit()

    @staticmethod
    def _sha1(data: str) -> str:
        return hashlib.sha1(data.encode("utf-8", errors="ignore")).hexdigest()

    def get(self, repo: str, relpath: str, content: str) -> Optional[Tuple[str, List[str], List[SymbolInfo], str]]:
        sha1 = self._sha1(content)
        with self._lock, sqlite3.connect(self.db_path) as con:
            cur = con.execute(
                "SELECT mod_qual, imports_json, source FROM files WHERE repo=? AND relpath=? AND sha1=?",
                (repo, relpath, sha1)
            )
            row = cur.fetchone()
            if not row:
                return None
            mod_qual, imports_json, source = row
            imports = json.loads(imports_json)
            sym_rows = con.execute(
                "SELECT name, qual, kind, line, code FROM symbols WHERE repo=? AND relpath=? AND sha1=?",
                (repo, relpath, sha1)
            ).fetchall()
            syms = [
                SymbolInfo(name=r[0], qualified_name=r[1], kind=r[2], line_number=r[3],
                           source_code=r[4], location=relpath)
                for r in sym_rows
            ]
            return mod_qual, imports, syms, source or ""

    def put(self, repo: str, relpath: str, mod_qual: str, imports: List[str],
            symbols: Iterable[SymbolInfo], source: str) -> None:
        sha1 = self._sha1(source)
        imports_json = json.dumps(sorted(set(imports)))
        with self._lock, sqlite3.connect(self.db_path) as con:
            con.execute(
                "INSERT OR REPLACE INTO files(repo, relpath, mod_qual, sha1, imports_json, source) VALUES(?,?,?,?,?,?)",
                (repo, relpath, mod_qual, sha1, imports_json, source)
            )
            con.execute("DELETE FROM symbols WHERE repo=? AND relpath=? AND sha1=?", (repo, relpath, sha1))
            con.executemany(
                "INSERT INTO symbols(repo, relpath, sha1, name, qual, kind, line, code) VALUES(?,?,?,?,?,?,?,?)",
                [(repo, relpath, sha1, s.name, s.qualified_name, s.kind, s.line_number, s.source_code) for s in symbols]
            )
            con.commit()

    # ------------- indexing state helpers (to avoid re-indexing) -------------- #
    def repo_has_index(self, repo: str) -> bool:
        with sqlite3.connect(self.db_path) as con:
            row = con.execute(
                "SELECT file_count FROM repo_index_state WHERE repo=?",
                (repo,)
            ).fetchone()
            if row and int(row[0]) > 0:
                return True
            # Fallback check (older DBs): look for any file rows for this repo
            row2 = con.execute(
                "SELECT 1 FROM files WHERE repo=? LIMIT 1",
                (repo,)
            ).fetchone()
            return bool(row2)

    def mark_repo_indexed(self, repo: str) -> None:
        with sqlite3.connect(self.db_path) as con:
            cnt = con.execute("SELECT COUNT(1) FROM files WHERE repo=?", (repo,)).fetchone()[0]
            con.execute(
                "INSERT OR REPLACE INTO repo_index_state(repo, indexed_at, file_count) VALUES(?,?,?)",
                (repo, datetime.utcnow().isoformat(), int(cnt))
            )
            con.commit()


# ----------------------------------------------------------------------------- #
# Extractor
# ----------------------------------------------------------------------------- #
class BlockExtractor:
    def __init__(self, max_workers: Optional[int] = None, max_retries: int = 2, index_mode: str = "missing"):
        """
        index_mode:
          - "missing" (default): index only repos that are not present in SQLite index
          - "force": re-index all repos every run
          - "skip": never index (assume index is prebuilt)
        """
        self.repo_cache = RepoCache()
        self.package_generator = PackageGenerator()
        self.max_workers = max_workers or max(os.cpu_count() or 8, 8)
        self.max_retries = max_retries

        self.import_graph = ImportGraph()
        self.parsed_files: Set[str] = set()
        self.failed_files: Set[str] = set()

        self.results_file = "extraction_results.json"
        self.extracted_blocks: List[Dict[str, Any]] = []
        self.failed_blocks: List[str] = []
        self.skipped_blocks: List[str] = []

        self.index = FileIndexStore()
        self._load_results()

        # package root → {repo: count}; and best repo per package root (systematic, from index)
        self._pkg_repo_counts: Dict[str, Dict[str, int]] = {}
        self._pkg_repo_best: Dict[str, str] = {}

        # re-export tracking
        self._reexports: Dict[str, str] = {}                     # exported_qname -> real_qname
        self._pending_star_exports: Dict[str, List[Tuple[str, str]]] = {}

        # indexing policy & warm flag
        self.index_mode = index_mode
        self._index_warmed = False

        # Real-definition resolver (no stubs)
        self.def_resolver = DefinitionResolver(
            repo_cache=self.repo_cache,
            import_graph=self.import_graph,
            index=self.index,
            log=log,
        )

        log.info("Initialized BlockExtractor | workers=%s | LibCST=%s | index_mode=%s",
                 self.max_workers, LIBCST_AVAILABLE, self.index_mode)

    # -------------------------- Utilities & persistence ------------------------ #

    def _load_results(self) -> None:
        fp = Path(self.results_file)
        if fp.exists():
            try:
                data = json.loads(fp.read_text())
                self.extracted_blocks = data.get("extracted_blocks", [])
                self.failed_blocks = data.get("failed_blocks", [])
                self.skipped_blocks = data.get("skipped_blocks", [])
                log.info("Loaded results: %d extracted, %d failed, %d skipped",
                         len(self.extracted_blocks), len(self.failed_blocks), len(self.skipped_blocks))
            except Exception as e:
                log.warning("Could not read results file: %s", e)

    def _save_results(self) -> None:
        data = {
            "timestamp": datetime.now().isoformat(),
            "extracted_blocks": self.extracted_blocks,
            "failed_blocks": self.failed_blocks,
            "skipped_blocks": self.skipped_blocks,
        }
        Path(self.results_file).write_text(json.dumps(data, indent=2))
        log.info("Saved results file.")

    # ------------------------------ Repo indexing ----------------------------- #

    def _should_skip(self, file_path: Path, size_limit_bytes: int = 1_200_000, repo: str = None) -> bool:
        if file_path.suffix != ".py":
            return True
        try:
            if file_path.stat().st_size > size_limit_bytes:
                return True
        except OSError:
            return True
        # Only check directory names that are relative to the repository root
        # The file_path here is the full absolute path, but we only want to check
        # the parts that are relative to the repo root
        if repo and hasattr(self, 'repo_cache') and hasattr(self.repo_cache, 'repos'):
            repo_root = self.repo_cache.get_cached_repo(repo)
            if repo_root:
                try:
                    rel_path = file_path.relative_to(repo_root)
                    # Check only the relative path parts for skipped directories
                    skip = {".git", ".venv", "__pycache__", "build", "dist", "site-packages", "third_party", "tests", "docs", "examples"}
                    if any(part in skip or part.startswith(".") for part in rel_path.parts):
                        return True
                except ValueError:
                    # If we can't determine relative path, use the original logic
                    skip = {".git", ".venv", "__pycache__", "build", "dist", "site-packages", "third_party", "tests", "docs", "examples"}
                    if any(part in skip or part.startswith(".") for part in file_path.parts):
                        return True
        else:
            # Fallback to original logic if no repo context
            skip = {".git", ".venv", "__pycache__", "build", "dist", "site-packages", "third_party", "tests", "docs", "examples"}
            if any(part in skip or part.startswith(".") for part in file_path.parts):
                return True
            
        # Check if file is within allowed paths for this repository
        if repo and hasattr(self, 'repo_cache') and hasattr(self.repo_cache, 'repos'):
            repo_config = self.repo_cache.repos.get(repo, {})
            allowed_paths = repo_config.get("paths", [])
            if allowed_paths:
                # Convert file path to relative path from repo root
                try:
                    repo_root = self.repo_cache.get_cached_repo(repo)
                    if repo_root:
                        rel_path = file_path.relative_to(repo_root)
                        rel_path_str = str(rel_path).replace("\\", "/")  # Normalize path separators
                        # Check if the file is within any of the allowed paths
                        for allowed_path in allowed_paths:
                            if rel_path_str.startswith(allowed_path) or rel_path_str.startswith(allowed_path + "/"):
                                return False  # Don't skip this file
                        # File is not within allowed paths
                        return True  # Skip this file
                except (ValueError, AttributeError):
                    # If we can't determine relative path, allow the file
                    pass
            else:
                # No allowed paths specified, allow all files
                return False
        
        return False

    def _mod_parts(self, repo_root: Path, file_path: Path) -> Tuple[str, str, str]:
        rel = file_path.relative_to(repo_root)
        rel_no_ext = rel.with_suffix("")   # path/to/mod
        parts = list(rel_no_ext.parts)
        mod_qual = ".".join(parts)
        return str(rel), ".".join(parts[:-1]) if len(parts) > 1 else "", parts[-1]

    # --- Hydrate the import graph ---
    def _hydrate_repo_from_index(self, repo: str) -> None:
        """Populate self.import_graph from rows already stored in SQLite for this repo."""
        with sqlite3.connect(self.index.db_path) as con:
            # pick the latest row per (repo, relpath) via ROWID
            rows = con.execute(
                """
                SELECT repo, relpath, mod_qual, sha1, imports_json, source
                FROM files
                WHERE repo=?
                AND ROWID IN (
                    SELECT MAX(ROWID) FROM files WHERE repo=? GROUP BY relpath
                )
                """,
                (repo, repo),
            ).fetchall()

            for _repo, relpath, mod_qual, sha1, imports_json, source in rows:
                imports = json.loads(imports_json or "[]")
                key = f"{repo}:{mod_qual}"
                mi = ModuleInfo(
                    repo=repo,
                    qual=mod_qual,
                    path=str(Path(relpath).with_suffix("")),
                    file=Path(relpath).stem,
                    source_code=source or "",
                    imports=imports,
                )
                # attach symbols for this exact file version
                syms = con.execute(
                    "SELECT name, qual, kind, line, code FROM symbols WHERE repo=? AND relpath=? AND sha1=?",
                    (repo, relpath, sha1),
                ).fetchall()
                for name, qual, kind, line, code in syms:
                    si = SymbolInfo(
                        name=name,
                        qualified_name=qual,
                        kind=kind,
                        location=relpath,
                        line_number=line,
                        source_code=code,
                    )
                    mi.symbols[name] = si
                    self.import_graph.symbol_table[qual] = si

                self.import_graph.modules[key] = mi
                for imp in mi.imports:
                    self.import_graph.dependents.setdefault(imp, []).append(key)
        
        # After modules + symbols are in memory, parse __init__ re-exports:
        for key, mi in list(self.import_graph.modules.items()):
            r, _ = key.split(":", 1)
            if r == repo and self._is_pkg_init(mi):
                self._extract_reexports_for_module(repo, mi)
        self._finalize_star_reexports(repo)

    # ---------- module-level constants harvesting ----------
    def _harvest_module_constants(self, mod_info: ModuleInfo) -> None:
        """
        Capture module-level constants/aliases as SymbolInfo(kind='const'):
        - Assign / AnnAssign with Name targets at module scope
        - 'try/except' blocks that define names (preserve entire try block)
        - simple 'if' guarded assignments (TYPE_CHECKING / feature flags)
        """
        src = mod_info.source_code or ""
        if not src:
            return

        try:
            tree = ast.parse(src)
        except Exception:
            return

        emitted: Set[str] = set()

        def add_const(name: str, node: ast.AST):
            if not name or name in emitted or name in mod_info.symbols:
                return
            if name.startswith("__") and name.endswith("__"):
                return
            code = self._slice_block(src, node)
            # Strip leading indentation for constants to make them valid Python code
            if code:
                lines = code.splitlines()
                if lines:
                    # Find the minimum indentation
                    min_indent = float('inf')
                    for line in lines:
                        if line.strip():  # Skip empty lines
                            indent = len(line) - len(line.lstrip())
                            min_indent = min(min_indent, indent)
                    # Strip the minimum indentation from all lines
                    if min_indent != float('inf'):
                        code = '\n'.join(line[min_indent:] if line.strip() else line for line in lines)
            
            qual = f"{mod_info.qual}.{name}"
            sym = SymbolInfo(
                name=name,
                qualified_name=qual,
                kind="const",
                location=f"{mod_info.path}.py",
                line_number=getattr(node, "lineno", 0),
                source_code=code,
            )
            mod_info.symbols.setdefault(name, sym)
            emitted.add(name)

        def targets_from_assign(node: ast.AST) -> List[str]:
            out: List[str] = []
            if isinstance(node, ast.Assign):
                for t in node.targets:
                    if isinstance(t, ast.Name):
                        out.append(t.id)
                    elif isinstance(t, (ast.Tuple, ast.List)):
                        for e in t.elts:
                            if isinstance(e, ast.Name):
                                out.append(e.id)
            elif isinstance(node, ast.AnnAssign):
                if isinstance(node.target, ast.Name):
                    out.append(node.target.id)
            return out

        for n in tree.body:
            if isinstance(n, (ast.Assign, ast.AnnAssign)):
                for nm in targets_from_assign(n):
                    add_const(nm, n)
            elif isinstance(n, ast.Try):
                handler_assigned: Set[str] = set()
                for h in n.handlers:
                    for s in h.body:
                        if isinstance(s, (ast.Assign, ast.AnnAssign)):
                            for nm in targets_from_assign(s):
                                handler_assigned.add(nm)
                for nm in handler_assigned:
                    add_const(nm, n)
            elif isinstance(n, ast.If):
                for s in list(n.body) + list(n.orelse or []):
                    if isinstance(s, (ast.Assign, ast.AnnAssign)):
                        for nm in targets_from_assign(s):
                            add_const(nm, s)
            elif isinstance(n, ast.FunctionDef):
                # Also capture function definitions as constants
                add_const(n.name, n)

    def _parse_one(self, repo: str, repo_root: Path, file_path: Path) -> Optional[Tuple[str, ModuleInfo]]:
        try:
            if self._should_skip(file_path, repo=repo):
                return None
            content = file_path.read_text(encoding="utf-8", errors="ignore")
            relpath = str(file_path.relative_to(repo_root))
            cached = self.index.get(repo, relpath, content)
            if cached:
                mod_qual, imports, syms, stored_source = cached
                module_key = f"{repo}:{mod_qual}"
                mi = ModuleInfo(
                    repo=repo, qual=mod_qual,
                    path=str(file_path.with_suffix("").relative_to(repo_root)),
                    file=file_path.stem,
                    source_code=stored_source or content,
                    imports=imports
                )
                for s in syms:
                    mi.symbols[s.name] = s
                # enrich cache with constants if needed
                before = set(mi.symbols.keys())
                self._harvest_module_constants(mi)
                after = set(mi.symbols.keys())
                if after - before:
                    self.index.put(repo, relpath, mi.qual, mi.imports, list(mi.symbols.values()), mi.source_code)
                return module_key, mi

            rel, mod_path, mod_file = self._mod_parts(repo_root, file_path)
            mod_qual = ".".join(p for p in (mod_path, mod_file) if p)

            if LIBCST_AVAILABLE and len(content) <= 120_000 and "{{" not in content and "}}" not in content:
                try:
                    tree = cst.parse_module(content)
                    wrapper = MetadataWrapper(tree)
                    positions = wrapper.resolve(PositionProvider)
                    mod_info = ModuleInfo(
                        repo=repo, qual=mod_qual,
                        path=str(file_path.with_suffix("").relative_to(repo_root)),
                        file=mod_file, source_code=content
                    )
                    imports = set()

                    class ImportCollector(cst.CSTVisitor):
                        def visit_Import(self, node: cst.Import) -> None:
                            for n in node.names:
                                imports.add(n.name.code)
                        def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
                            if node.module:
                                imports.add(node.module.code)

                    tree.visit(ImportCollector())
                    for node in tree.body:
                        if isinstance(node, cst.ClassDef):
                            pos = positions[node].start
                            code = tree.code_for_node(node)
                            s = SymbolInfo(
                                name=node.name.value,
                                qualified_name=f"{mod_qual}.{node.name.value}",
                                kind="class",
                                location=rel,
                                line_number=pos.line,
                                source_code=code,
                            )
                            mod_info.symbols[s.name] = s
                        elif isinstance(node, cst.FunctionDef):
                            pos = positions[node].start
                            code = tree.code_for_node(node)
                            s = SymbolInfo(
                                name=node.name.value,
                                qualified_name=f"{mod_qual}.{node.name.value}",
                                kind="function",
                                location=rel,
                                line_number=pos.line,
                                source_code=code,
                            )
                            mod_info.symbols[s.name] = s
                    mod_info.imports = sorted(imports)
                except Exception as e:
                    log.debug("LibCST parsing failed for %s: %s", file_path, e)
                    mod_info = self._parse_with_ast(repo, repo_root, file_path, content, rel, mod_qual)
            else:
                mod_info = self._parse_with_ast(repo, repo_root, file_path, content, rel, mod_qual)

            self._harvest_module_constants(mod_info)
            self.index.put(repo, relpath, mod_info.qual, mod_info.imports, list(mod_info.symbols.values()), content)
            return f"{repo}:{mod_info.qual}", mod_info
        except Exception as e:
            self.failed_files.add(str(file_path))
            log.error("Parse failed for %s: %s", file_path, e)
            return None

    def _parse_with_ast(self, repo: str, repo_root: Path, file_path: Path,
                        content: str, rel: str, mod_qual: str) -> ModuleInfo:
        try:
            tree = ast.parse(content)
        except Exception as e:
            log.error("AST parsing failed for %s: %s", file_path, e)
            raise
        
        mod_info = ModuleInfo(
            repo=repo, qual=mod_qual,
            path=str(file_path.with_suffix("").relative_to(repo_root)),
            file=file_path.stem, source_code=content,
        )
        imports: Set[str] = set()
        for n in ast.walk(tree):
            if isinstance(n, ast.Import):
                imports.update(a.name for a in n.names if a.name)
            elif isinstance(n, ast.ImportFrom) and n.module:
                imports.add(n.module)
        mod_info.imports = sorted(imports)

        for node in tree.body:
            if isinstance(node, ast.ClassDef):
                code = self._slice_block(content, node)
                s = SymbolInfo(
                    name=node.name,
                    qualified_name=f"{mod_qual}.{node.name}",
                    kind="class",
                    location=rel,
                    line_number=getattr(node, "lineno", 0),
                    source_code=code,
                )
                mod_info.symbols.setdefault(s.name, s)
            elif isinstance(node, ast.FunctionDef):
                code = self._slice_block(content, node)
                s = SymbolInfo(
                    name=node.name,
                    qualified_name=f"{mod_qual}.{node.name}",
                    kind="function",
                    location=rel,
                    line_number=getattr(node, "lineno", 0),
                    source_code=code,
                )
                mod_info.symbols.setdefault(s.name, s)

        return mod_info

    @staticmethod
    def _slice_block(full_source: str, node: ast.AST) -> str:
        if hasattr(node, "lineno") and hasattr(node, "end_lineno") and node.lineno and node.end_lineno:
            lines = full_source.splitlines()
            start = node.lineno - 1
            while start > 0 and lines[start - 1].lstrip().startswith("@"):
                start -= 1
            return "\n".join(lines[start: node.end_lineno])
        if hasattr(ast, "unparse"):
            return ast.unparse(node)
        return ""

    def index_repository(self, repo: str) -> None:
        root = self.repo_cache.get_cached_repo(repo)
        if not root:
            log.warning("Repo %s not cached; skipping.", repo)
            return
        py_files = [p for p in root.rglob("*.py") if not self._should_skip(p, repo=repo)]
        log.info("Found %d Python files in %s after path filtering", len(py_files), repo)
        if not py_files:
            log.warning("No Python files found in %s after path filtering", repo)
            return

        log.info("Indexing %s (%d files)...", repo, len(py_files))
        t0 = time.time()
        with cf.ThreadPoolExecutor(max_workers=self.max_workers) as ex:
            futures = [ex.submit(self._parse_one, repo, root, fp) for fp in py_files]
            for fut in cf.as_completed(futures):
                res = fut.result()
                if not res:
                    continue
                key, mod_info = res
                self.import_graph.modules[key] = mod_info
                for s in mod_info.symbols.values():
                    self.import_graph.symbol_table[s.qualified_name] = s
                for imp in mod_info.imports:
                    self.import_graph.dependents.setdefault(imp, []).append(key)

        dt = time.time() - t0
        # Mark repo as indexed in SQLite (for next processes to skip)
        self.index.mark_repo_indexed(repo)
        log.info(
            "Indexed %s: %d modules | %d symbols in %.2fs",
            repo,
            sum(1 for k in self.import_graph.modules if k.startswith(f"{repo}:")),
            len(self.import_graph.symbol_table),
            dt,
        )
        
        # Build re-exports for this repo
        for key, mod in list(self.import_graph.modules.items()):
            r, _ = key.split(":", 1)
            if r != repo:
                continue
            if self._is_pkg_init(mod):
                self._extract_reexports_for_module(repo, mod)
        self._finalize_star_reexports(repo)

    # ---------- package→repo map (SYSTEMATIC, from index; no hardcoding) ----------
    def _refresh_package_repo_map(self) -> None:
        counts: Dict[str, Dict[str, int]] = {}
        for key, mi in self.import_graph.modules.items():
            repo, qual = key.split(":", 1)
            if not qual:
                continue
            root = qual.split(".", 1)[0]
            counts.setdefault(root, {}).setdefault(repo, 0)
            counts[root][repo] += 1
        best: Dict[str, str] = {}
        for root, by_repo in counts.items():
            best_repo = max(by_repo.items(), key=lambda kv: kv[1])[0]
            best[root] = best_repo
        self._pkg_repo_counts = counts
        self._pkg_repo_best = best

    # ------------------------------ Discovery --------------------------------- #

    def _ensure_all_repos_cached(self) -> bool:
        try:
            cfg = json.loads((Path(__file__).parent / "config" / "repo_config.json").read_text())
        except FileNotFoundError:
            log.error("repo_config.json not found.")
            return False
        for name in cfg.keys():
            try:
                if not self.repo_cache.is_repo_cached(name):
                    # If RepoCache attempts pull & fails, it should raise; we deliberately continue.
                    self.repo_cache.ensure_repo_cached(name)
            except Exception as e:
                log.warning("Cache failed for %s (continuing with whatever is present): %s", name, e)
        return True

    def _scan_repo_for_block(self, repo: str, block: str) -> List[Dict[str, Any]]:
        root = self.repo_cache.get_cached_repo(repo)
        if not root:
            return []
        out: List[Dict[str, Any]] = []
        pat = re.compile(rf"^\s*(class|def)\s+{re.escape(block)}\b", re.MULTILINE)
        for fp in root.rglob("*.py"):
            if self._should_skip(fp, repo=repo):
                continue
            try:
                text = fp.read_text(encoding="utf-8", errors="ignore")
            except Exception:
                continue
            if pat.search(text):
                out.append({
                    "repository": repo,
                    "file_path": str(fp.relative_to(root)),
                    "identifier": block,
                    "confidence": 1.0,
                    "type": "class" if re.search(rf"^\s*class\s+{re.escape(block)}\b", text, re.MULTILINE) else "function"
                })
        return out

    def _discover_blocks(self, targets: List[str]) -> Dict[str, List[Dict[str, Any]]]:
        results: Dict[str, List[Dict[str, Any]]] = {t: [] for t in targets}
        for repo in list(self.repo_cache.repos.keys()):
            if not self.repo_cache.is_repo_cached(repo):
                continue
            for block in targets:
                if results[block]:
                    continue
                matches = self._scan_repo_for_block(repo, block)
                if matches:
                    results[block].extend(matches)
        return results

    # ----------------------- Dependency resolution ---------------------------- #

    def _get_dynamic_ignore_externals(self) -> Set[str]:
        """Dynamically build ignore list based on available modules and imports."""
        ignore_set = set()
        try:
            ignore_set.update(dir(__builtins__))
        except Exception:
            pass
        ignore_set.update({
            "torch", "nn", "F", "self", "super", "Config", "ConfigDict", "MODULE2PACKAGE",
            "Parameter", "Self", "Tensor", "_copy_to_script_wrapper", "container_abcs",
            "end_id", "init", "islice", "namedtuple", "start_id", "Sequence"
        })
        # Note: Typing symbols should be resolved to imports, not ignored
        # They will be handled by the dependency resolution system
        ignore_set.update(self._get_stdlib_modules())
        
        # Add common standard library symbols that should be ignored
        # Note: Some of these should be resolved to imports, not completely ignored
        stdlib_symbols = {
            "Enum", "chain", "map", "filter", "reduce", "zip", "range", "len", "str", "int", "float", "bool", "list", "dict", "set", "tuple",
            "ArgumentParser", "BytesIO", "Console", "Dumper", "FileHandler", "Generator", "Loader", "LogRecord", "Logger",
            "OrderedDict", "Path", "SameFileError", "StringIO", "Table", "Text", "colored", "contextmanager", "defaultdict", "find_spec", "gethostname", "getuser",
            "handlers", "import_module", "ismodule", "parse", "urlopen", "yaml", "osp", "np", "mmengine", "master_only", "k", "v", "MODELS", "_format_dict"
        }
        ignore_set.update(stdlib_symbols)
        
        # Special handling for symbols that should be resolved to imports, not ignored
        # These will be handled by the dependency resolution system
        import_resolvable_symbols = {
            "ABCMeta", "abstractmethod", "ast"  # These should resolve to: from abc import ABCMeta, abstractmethod; import ast
        }
        # Don't add these to ignore_set - let them be resolved
        
        # Add common type annotation names that should be ignored
        type_annotations = {"T", "U", "V", "K", "V", "Any", "Optional", "Union", "List", "Dict", "Tuple", "Set", "Callable", "TypeVar", "Generic"}
        ignore_set.update(type_annotations)
        
        # Add more built-in functions and exceptions that should be ignored
        builtin_functions = {
            "Exception", "ImportError", "ValueError", "RuntimeError", "TypeError", "AttributeError",
            "classmethod", "dir", "getattr", "hasattr", "issubclass", "property", "staticmethod",
            "iter", "next", "reversed", "type", "isinstance", "len", "str", "int", "float", "bool",
            "list", "dict", "set", "tuple", "range", "zip", "map", "filter", "reduce", "enumerate",
            "sorted", "min", "max", "sum", "abs", "round", "divmod", "pow", "bin", "hex", "oct",
            "chr", "ord", "repr", "eval", "exec", "compile", "globals", "locals", "vars"
        }
        ignore_set.update(builtin_functions)
        
        ignore_set.add("_")
        return ignore_set

    def _get_stdlib_modules(self) -> Set[str]:
        """Dynamically detect available standard library modules + common DS aliases."""
        stdlib_modules = set()
        common_stdlib = [
            "os", "sys", "re", "math", "random", "time", "datetime", "json", "pickle", "collections",
            "itertools", "functools", "pathlib", "shutil", "tempfile", "logging", "warnings", "traceback",
            "subprocess", "multiprocessing", "concurrent", "asyncio", "socket",
            "importlib", "difflib", "argparse", "io", "types", "platform", "uuid",
            "urllib", "importlib.util", "enum", "os", "yaml", "rich", "numpy",
            "abc", "inspect", "threading", "contextlib", "copy", "weakref", "gc",
            "operator", "heapq", "bisect", "array", "struct", "hashlib", "hmac",
            "base64", "binascii", "zlib", "gzip", "bz2", "lzma", "zipfile", "tarfile",
            "itertools", "collections", "collections.abc"
        ]
        for module_name in common_stdlib:
            try:
                importlib.import_module(module_name)
                stdlib_modules.add(module_name)
            except ImportError:
                pass
        for alias in ("np", "pd", "plt", "cv2", "osp", "yaml", "rich"):
            try:
                if alias == "np":
                    import numpy  # noqa
                    stdlib_modules.add(alias)
                elif alias == "pd":
                    import pandas  # noqa
                    stdlib_modules.add(alias)
                elif alias == "plt":
                    import matplotlib.pyplot  # noqa
                    stdlib_modules.add(alias)
                elif alias == "cv2":
                    import cv2  # noqa
                    stdlib_modules.add(alias)
                elif alias == "osp":
                    import os.path  # noqa
                    stdlib_modules.add(alias)
                elif alias == "yaml":
                    import yaml  # noqa
                    stdlib_modules.add(alias)
                elif alias == "rich":
                    import rich  # noqa
                    stdlib_modules.add(alias)
            except ImportError:
                pass
        return stdlib_modules

    def _free_names(self, code: str) -> Set[str]:
        """Scope-aware free-name collector (module/class/func/lambda/comps/walrus)."""
        try:
            tree = ast.parse(code)
        except Exception:
            return set()

        free: Set[str] = set()
        scope_stack: List[Set[str]] = [set()]
        in_annotation = False

        def define(name: str) -> None:
            if scope_stack:
                scope_stack[-1].add(name)

        def is_defined(name: str) -> bool:
            return any(name in s for s in reversed(scope_stack))

        class V(ast.NodeVisitor):
            def _push(self): scope_stack.append(set())
            def _pop(self): scope_stack.pop()

            def _args(self, args: ast.arguments):
                for a in getattr(args, "posonlyargs", []): define(a.arg)
                for a in getattr(args, "args", []): define(a.arg)
                for a in getattr(args, "kwonlyargs", []): define(a.arg)
                va = getattr(args, "vararg", None)
                if va is not None: define(va.arg)
                ka = getattr(args, "kwarg", None)
                if ka is not None: define(ka.arg)

            def _targets(self, t):
                if isinstance(t, ast.Name): define(t.id)
                elif isinstance(t, (ast.Tuple, ast.List)):
                    for e in t.elts:
                        if isinstance(e, ast.Name): define(e.id)

            def visit_FunctionDef(self, node: ast.FunctionDef):
                define(node.name)
                self._push()
                for dec in node.decorator_list: self.visit(dec)
                if node.returns:
                    nonlocal in_annotation
                    in_annotation = True; self.visit(node.returns); in_annotation = False
                self._args(node.args)
                for arg in list(getattr(node.args, "posonlyargs", [])) + list(getattr(node.args, "args", [])) + list(getattr(node.args, "kwonlyargs", [])):
                    if getattr(arg, "annotation", None) is not None:
                        in_annotation = True; self.visit(arg.annotation); in_annotation = False
                if getattr(node.args, "vararg", None) and node.args.vararg.annotation:
                    in_annotation = True; self.visit(node.args.vararg.annotation); in_annotation = False
                if getattr(node.args, "kwarg", None) and node.args.kwarg.annotation:
                    in_annotation = True; self.visit(node.args.kwarg.annotation); in_annotation = False
                for s in node.body: self.visit(s)
                self._pop()

            def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
                self.visit_FunctionDef(node)

            def visit_Lambda(self, node: ast.Lambda):
                self._push()
                self._args(node.args)
                self.visit(node.body)
                self._pop()

            def visit_ClassDef(self, node: ast.ClassDef):
                define(node.name)
                for b in node.bases: self.visit(b)
                for kw in node.keywords: self.visit(kw)
                for dec in node.decorator_list: self.visit(dec)
                self._push()
                for s in node.body: self.visit(s)
                self._pop()

            def visit_NamedExpr(self, node: ast.NamedExpr):
                if isinstance(node.target, ast.Name):
                    define(node.target.id)
                self.visit(node.value)

            def visit_AnnAssign(self, node: ast.AnnAssign):
                if node.target: self._targets(node.target)
                if node.annotation:
                    nonlocal in_annotation
                    in_annotation = True; self.visit(node.annotation); in_annotation = False
                if node.value: self.visit(node.value)

            def visit_Assign(self, node: ast.Assign):
                for t in node.targets: self._targets(t)
                if node.value: self.visit(node.value)

            def visit_AugAssign(self, node: ast.AugAssign):
                self._targets(node.target); self.visit(node.value)

            def _visit_comp(self, node):
                self._push()
                for gen in node.generators:
                    self._targets(gen.target)
                    self.visit(gen.iter)
                    for if_ in gen.ifs:
                        self.visit(if_)
                if hasattr(node, "elt") and node.elt is not None:
                    self.visit(node.elt)
                if hasattr(node, "key") and node.key is not None:
                    self.visit(node.key)
                if hasattr(node, "value") and node.value is not None:
                    self.visit(node.value)
                self._pop()

            def visit_ListComp(self, node: ast.ListComp): self._visit_comp(node)
            def visit_SetComp(self, node: ast.SetComp): self._visit_comp(node)
            def visit_DictComp(self, node: ast.DictComp): self._visit_comp(node)
            def visit_GeneratorExp(self, node: ast.GeneratorExp): self._visit_comp(node)

            def visit_For(self, node: ast.For):
                self._targets(node.target); self.visit(node.iter)
                for s in node.body: self.visit(s)
                for s in node.orelse: self.visit(s)

            def visit_AsyncFor(self, node: ast.AsyncFor): self.visit_For(node)
            def visit_With(self, node: ast.With):
                for item in node.items:
                    if item.optional_vars: self._targets(item.optional_vars)
                    self.visit(item.context_expr)
                for s in node.body: self.visit(s)
            def visit_AsyncWith(self, node: ast.AsyncWith): self.visit_With(node)

            def visit_ExceptHandler(self, node: ast.ExceptHandler):
                if node.name: define(node.name)
                if node.type: self.visit(node.type)
                for s in node.body: self.visit(s)

            def visit_Import(self, node: ast.Import):
                for a in node.names:
                    alias = a.asname or a.name.split(".", 1)[0]
                    define(alias)

            def visit_ImportFrom(self, node: ast.ImportFrom):
                for a in node.names:
                    if a.name == "*": continue
                    alias = a.asname or a.name
                    define(alias)

            def visit_Name(self, node: ast.Name):
                if isinstance(node.ctx, ast.Load):
                    if not is_defined(node.id):
                        free.add(node.id)

            def visit_Global(self, node: ast.Global):
                for n in node.names: define(n)

            def visit_Nonlocal(self, node: ast.Nonlocal):
                for n in node.names: define(n)

        V().visit(tree)
        return free

    # ---------- package helpers ----------
    def _collect_import_roots(self, repo: str, mod_qual: str) -> Set[str]:
        roots: Set[str] = set()
        mi = self.import_graph.modules.get(f"{repo}:{mod_qual}")
        if not mi:
            return roots
        for imp in mi.imports:
            if not imp:
                continue
            roots.add(imp.split(".", 1)[0])
        return roots

    def _module_ctx(self, sym: SymbolInfo) -> Tuple[str, str]:
        """Return (repo, module_qual) for this symbol, or infer via package→repo map (systematic)."""
        for key, mi in self.import_graph.modules.items():
            repo, mod = key.split(":", 1)
            if any(s.qualified_name == sym.qualified_name for s in mi.symbols.values()):
                return repo, mod
        q = sym.qualified_name
        mod_qual = q.rsplit(".", 1)[0]
        root = mod_qual.split(".", 1)[0] if mod_qual else q.split(".", 1)[0]
        repo = self._pkg_repo_best.get(root, "")
        return repo, mod_qual

    def _alias_map(self, repo: str, mod_qual: str) -> Dict[str, str]:
        key = f"{repo}:{mod_qual}"
        mi = self.import_graph.modules.get(key)
        out: Dict[str, str] = {}
        if not mi or not mi.source_code:
            return out
        try:
            tree = ast.parse(mi.source_code)
            for n in tree.body:
                if isinstance(n, ast.Import):
                    for a in n.names:
                        base = a.name
                        alias = a.asname or base.split(".", 1)[0]
                        out[alias] = base
                elif isinstance(n, ast.ImportFrom):
                    base = n.module or ""
                    if n.level:
                        parts = mod_qual.split(".")
                        base = ".".join(parts[:len(parts)-n.level] + ([base] if base else []))
                    for a in n.names:
                        if a.name == "*":
                            continue
                        alias = a.asname or a.name
                        out[alias] = f"{base}.{a.name}".lstrip(".")
        except Exception:
            pass
        return out

    def _disambiguate_candidates(
        self,
        candidates: List[str],
        current_repo: str,
        current_mod: str,
        import_roots: Set[str],
    ) -> Optional[str]:
        """Pick best qname using import affinity + package mapping (no repo hardcoding)."""
        if not candidates:
            return None
        if len(candidates) == 1:
            return candidates[0]

        current_root = current_mod.split(".", 1)[0] if current_mod else ""
        scored: List[Tuple[Tuple[int, int, int, int], str]] = []
        for q in candidates:
            q_root = q.split(".", 1)[0]
            import_affinity = 1 if q_root in import_roots else 0
            same_root = 1 if q_root == current_root and current_root else 0
            pkg_strength = self._pkg_repo_counts.get(q_root, {}).get(self._pkg_repo_best.get(q_root, ""), 0)
            module_len = -len(q.rsplit(".", 1)[0])  # prefer deeper module paths slightly
            scored.append(((import_affinity, same_root, pkg_strength, module_len), q))
        scored.sort(reverse=True)
        return scored[0][1]

    def _map_refs(self, refs: Iterable[str], sym: SymbolInfo) -> Tuple[List[str], Set[str]]:
        """
        Map bare names to fully-qualified candidates via:
          1) same-module top-level defs (incl. harvested consts)
          2) imported aliases from the owning module
          3) global exact matches
          4) import-aware unique name search (prefer packages the module imports)

        Returns:
          (qualified_symbol_refs, externals_satisfied_by_imports)
        """
        repo, mod = self._module_ctx(sym)
        alias_map = self._alias_map(repo, mod)
        import_roots = self._collect_import_roots(repo, mod)
        out_syms: Set[str] = set()
        externals_ok: Set[str] = set()

        ignore_externals = self._get_dynamic_ignore_externals()

        for r in refs:
            if r in ignore_externals:
                continue

            # Handle import-resolvable standard library symbols systematically
            if r in ["ABCMeta", "abstractmethod"]:
                # These should resolve to existing symbols in the import graph
                # Look for existing ABCMeta/abstractmethod symbols
                existing_symbols = [q for q, si in self.import_graph.symbol_table.items() 
                                  if si.name == r and si.kind == "const" and "import" in si.source_code]
                if existing_symbols:
                    # Use the first one found (they all have the same import source)
                    out_syms.add(existing_symbols[0])
                    continue
                else:
                    # Fallback: mark as satisfied by external imports
                    externals_ok.add(r)
                    continue
            elif r in ["ast", "defaultdict", "logging", "copy", "os", "math", "collections", "itertools", "functools", "pathlib", "shutil", "tempfile", "warnings", "traceback", "subprocess", "multiprocessing", "concurrent", "asyncio", "socket", "importlib", "difflib", "argparse", "io", "types", "platform", "uuid", "urllib", "enum", "yaml", "rich", "numpy", "pandas", "matplotlib", "cv2"]:
                # Try to resolve these to their actual definitions instead of importing them
                existing_symbols = [q for q, si in self.import_graph.symbol_table.items() 
                                  if si.name == r]
                if existing_symbols:
                    # Prioritize class/function definitions over import statements
                    definition_symbols = [q for q in existing_symbols 
                                        if self.import_graph.symbol_table[q].kind in ["class", "function"]]
                    if definition_symbols:
                        out_syms.add(definition_symbols[0])
                    else:
                        # Use the first one found if no class/function definitions
                        out_syms.add(existing_symbols[0])
                    continue
                else:
                    # Only if we can't find the definition, mark as external
                    externals_ok.add(r)
                    continue
            elif r in ["FileHandler", "Logger", "LogRecord", "ManagerMixin", "osp"]:
                # These are common aliases or specific symbols that should be resolved
                # Look for existing symbols with these names, prioritizing class definitions
                existing_symbols = [q for q, si in self.import_graph.symbol_table.items() 
                                  if si.name == r]
                if existing_symbols:
                    # Prioritize class definitions over import statements
                    class_symbols = [q for q in existing_symbols 
                                   if self.import_graph.symbol_table[q].kind == "class"]
                    if class_symbols:
                        out_syms.add(class_symbols[0])
                    else:
                        # Use the first one found if no class definitions
                        out_syms.add(existing_symbols[0])
                    continue
                else:
                    # Mark as satisfied by external imports
                    externals_ok.add(r)
                    continue

            # 1) same-module - but prefer class definitions over constants/imports
            cand_local = f"{mod}.{r}" if mod else r
            if cand_local in self.import_graph.symbol_table:
                local_sym = self.import_graph.symbol_table[cand_local]
                # If it's just a constant/import, check if there's a better class/function definition
                if local_sym.kind == "const" and local_sym.source_code and "import" in local_sym.source_code:
                    # Look for class/function definitions with the same name
                    class_candidates = [q for q, si in self.import_graph.symbol_table.items() 
                                      if si.name == r and si.kind == "class"]
                    function_candidates = [q for q, si in self.import_graph.symbol_table.items() 
                                         if si.name == r and si.kind == "function"]
                    
                    if class_candidates:
                        # Prefer the class definition
                        chosen = self._disambiguate_candidates(class_candidates, repo, mod, import_roots)
                        if chosen:
                            out_syms.add(chosen)
                            continue
                    elif function_candidates:
                        # Use function definition if no class definitions
                        chosen = self._disambiguate_candidates(function_candidates, repo, mod, import_roots)
                        if chosen:
                            out_syms.add(chosen)
                            continue
                # Otherwise use the local symbol
                out_syms.add(cand_local)
                continue

            # 2) imported alias? (can point at symbol OR package re-export)
            if r in alias_map:
                q = alias_map[r]  # may be a leaf symbol, or a package-level export like "timm.layers.X"
                # follow re-exports if needed
                q_real = self._reexports.get(q, q)
                if q_real in self.import_graph.symbol_table:
                    out_syms.add(q_real)
                else:
                    # If re-export points to a symbol that doesn't exist, we need to resolve it
                    # Don't treat as satisfied import - let dependency resolution handle it
                    pass
                continue
            
            # 3) Global search with class definition prioritization
            # Look for all symbols with this name globally
            global_candidates = [q for q, si in self.import_graph.symbol_table.items() 
                               if si.name == r]
            if global_candidates:
                # Always prioritize class definitions over import statements
                class_candidates = [q for q in global_candidates 
                                  if self.import_graph.symbol_table[q].kind == "class"]
                function_candidates = [q for q in global_candidates 
                                    if self.import_graph.symbol_table[q].kind == "function"]
                
                if class_candidates:
                    # Use the best class definition
                    chosen = self._disambiguate_candidates(class_candidates, repo, mod, import_roots)
                    if chosen:
                        out_syms.add(chosen)
                        continue
                elif function_candidates:
                    # Use the best function definition if no class definitions
                    chosen = self._disambiguate_candidates(function_candidates, repo, mod, import_roots)
                    if chosen:
                        out_syms.add(chosen)
                        continue
                else:
                    # Only use const/import symbols if no class/function definitions exist
                    # This prevents import statements from being used when actual definitions exist
                    pass
            
            # 2.5) Handle relative imports (e.g., from .weight_init import ...)
            # Look for symbols that might be from relative imports
            if mod:
                # Try to resolve relative imports by looking for symbols in sibling modules
                base_mod = mod.rsplit(".", 1)[0] if "." in mod else mod
                relative_candidates = [q for q, si in self.import_graph.symbol_table.items() 
                                    if si.name == r and q.startswith(base_mod + ".") and q != cand_local]
                if relative_candidates:
                    # Use the first relative candidate found
                    out_syms.add(relative_candidates[0])
                continue

            # 3) already-qualified
            q3 = self._reexports.get(r, r)
            if q3 in self.import_graph.symbol_table:
                out_syms.add(q3)
                continue

            # 4) name-only search with import-aware disambiguation
            name_matches = [q for q, si in self.import_graph.symbol_table.items() if si.name == r]
            if name_matches:
                chosen = self._disambiguate_candidates(name_matches, repo, mod, import_roots)
                if chosen:
                    out_syms.add(chosen)
                    continue

            # 5) if matches an imported package root, treat as satisfied by import
            if r in import_roots:
                externals_ok.add(r)

        return sorted(out_syms), externals_ok

    # ---------- re-export helpers ----------
    def _is_pkg_init(self, mi: ModuleInfo) -> bool:
        # We normalize package qual as "....__init__" in our index; treat it as a package
        return mi.file == "__init__" or mi.qual.endswith(".__init__")

    def _pkg_qual(self, mi: ModuleInfo) -> str:
        # Convert "timm.layers.__init__" -> "timm.layers"
        return mi.qual[:-9] if mi.qual.endswith(".__init__") else mi.qual

    def _resolve_relative(self, base_pkg: str, level: int, module: Optional[str]) -> str:
        """
        Resolve 'from ....foo import bar' relative module to absolute qual under base_pkg.
        base_pkg: e.g., 'timm.layers'
        level: number of dots (e.g., 1 for '.', 2 for '..')
        module: module name (e.g., 'adaptive_avgmax_pool')
        """
        parts = base_pkg.split(".")
        # For relative imports, we want to keep the current package level
        # from .sub import X -> base_pkg.sub
        # from ..sub import X -> base_pkg.parent.sub
        if level == 1:
            # Single dot: from .sub import X -> base_pkg.sub
            if module:
                return f"{base_pkg}.{module}"
            else:
                return base_pkg
        elif level > 1:
            # Multiple dots: from ..sub import X -> base_pkg.parent.sub
            up = level - 1  # Number of levels to go up
            root = parts[: max(0, len(parts) - up)]
            if module:
                root.append(module)
            return ".".join([p for p in root if p])
        else:
            # level == 0: absolute import
            if module:
                return module
            else:
                return ""

    def _extract_reexports_for_module(self, repo: str, mi: ModuleInfo) -> None:
        """Parse a module (usually a package __init__) and record explicit/STAR re-exports."""
        if not mi.source_code:
            return
        try:
            tree = ast.parse(mi.source_code)
        except Exception:
            return

        pkg = self._pkg_qual(mi)

        for node in tree.body:
            if isinstance(node, ast.ImportFrom):
                level = getattr(node, "level", 0) or 0
                mod = node.module
                
                # Handle both relative and absolute imports for re-exports
                if level > 0:
                    # Relative import: from .sub import X
                    target_mod = self._resolve_relative(pkg, level, mod)
                elif mod and mod.startswith(pkg.split(".")[0]) and len(mod.split(".")) > len(pkg.split(".")):
                    # Absolute import within the same package hierarchy: from timm.layers.sub import X
                    # But NOT from timm.config import X (which is external to timm.layers)
                    target_mod = mod
                elif mod and mod.startswith("mmengine"):
                    # Special case: allow MMEngine imports for re-exports
                    # This handles cases like "from mmengine.registry import MODELS as MMENGINE_MODELS"
                    target_mod = mod
                else:
                    # External import, skip
                    continue
                
                for alias in node.names:
                    if alias.name == "*":
                        # Defer star expansion until all modules in this repo are loaded/hydrated
                        self._pending_star_exports.setdefault(repo, []).append((pkg, target_mod))
                        continue
                    exported = alias.asname or alias.name
                    real = f"{target_mod}.{alias.name}"
                    self._reexports[f"{pkg}.{exported}"] = real
                    
                    # Also handle the case where the alias is used directly (e.g., MMENGINE_MODELS)
                    if alias.asname:
                        # This is an import alias: from mod import X as Y
                        # Map Y -> mod.X
                        self._reexports[alias.asname] = real

    def _finalize_star_reexports(self, repo: str) -> None:
        """Expand 'from .sub import *' using the already-populated symbol table."""
        pending = self._pending_star_exports.get(repo) or []
        if not pending:
            return
        for pkg, submod in pending:
            prefix = f"{submod}."
            # Collect all symbol basenames from submodule
            for q, si in self.import_graph.symbol_table.items():
                if q.startswith(prefix):
                    basename = q.split(".")[-1]
                    self._reexports.setdefault(f"{pkg}.{basename}", q)
        # Clear after processing to avoid duplication across runs
        self._pending_star_exports[repo] = []

    @staticmethod
    def _head_kind(snippet: str) -> str:
        try:
            t = ast.parse(snippet)
            if t.body and isinstance(t.body[0], ast.ClassDef): return "class"
            if t.body and isinstance(t.body[0], ast.FunctionDef): return "function"
        except Exception:
            pass
        return "const"

    def _dfs_all(self, root: SymbolInfo) -> Tuple[Dict[str, ResolvedDependency], Set[str], List[str]]:
        # Fundamental PyTorch classes that should not be included as dependencies
        FUNDAMENTAL_TORCH_CLASSES = {
            'torch.nn.modules.module.Module',
            'torch.nn.parameter.Parameter',
            'torch.Tensor'
        }
        
        resolved: Dict[str, ResolvedDependency] = {}
        unresolved: Set[str] = set()
        seen: Set[str] = set()
        order: List[str] = []

        def dfs(qname: str):
            if qname in seen:
                return
            seen.add(qname)
            
            # Skip fundamental PyTorch classes
            if qname in FUNDAMENTAL_TORCH_CLASSES:
                return
                
            sym = self.import_graph.symbol_table.get(qname)
            if not sym:
                nm = qname.split(".")[-1]
                candidates = [s for s in self.import_graph.symbol_table.values() if s.name == nm]
                if candidates:
                    # Prioritize class/function definitions over import statements
                    definition_candidates = [s for s in candidates if s.kind in ["class", "function"]]
                    if definition_candidates:
                        sym = definition_candidates[0]
                    else:
                        sym = candidates[0]
                    qname = sym.qualified_name
                else:
                    unresolved.add(qname)
                    return

            # First, add the current symbol to resolved dependencies
            resolved[qname] = ResolvedDependency(
                name=sym.name,
                qualified_name=qname,
                source_code=sym.source_code,
                resolution_method="index",
                confidence=0.95,
                location=sym.location
            )
            order.append(qname)

            # Then process its dependencies
            refs = self._free_names(sym.source_code)
            mapped_syms, externals_ok = self._map_refs(refs, sym)

            mapped_basenames = {m.split(".")[-1] for m in mapped_syms}
            ignore_externals = self._get_dynamic_ignore_externals()
            for nm in refs - mapped_basenames:
                if nm in externals_ok:
                    continue
                if nm not in ignore_externals:
                    unresolved.add(nm)

            # Recursively process all mapped dependencies
            for child in mapped_syms:
                dfs(child)

        dfs(root.qualified_name)
        return resolved, unresolved, order

    def _inject_resolved_symbols(self, fetched: Dict[str, ResolvedSymbol]) -> None:
        for qname, rsym in fetched.items():
            repo = rsym.repo or "__resolved__"
            mod_qual = rsym.module_qual or qname.rsplit(".", 1)[0]
            key = f"{repo}:{mod_qual}"
            if key not in self.import_graph.modules:
                self.import_graph.modules[key] = ModuleInfo(
                    repo=repo, qual=mod_qual, path=rsym.module_path or mod_qual.replace(".", "/"),
                    file=mod_qual.rsplit(".", 1)[-1], source_code=rsym.module_source or ""
                )
            mi = self.import_graph.modules[key]
            name = qname.split(".")[-1]
            kind = rsym.kind or self._head_kind(rsym.source_code)
            sym = SymbolInfo(
                name=name, qualified_name=qname, kind=kind,
                location=rsym.file_path or mi.path + ".py",
                line_number=rsym.line or 0,
                source_code=rsym.source_code,
            )
            mi.symbols[name] = sym
            self.import_graph.symbol_table[qname] = sym
            if rsym.module_imports:
                mi.imports = sorted(set(mi.imports) | set(rsym.module_imports))

    def _find_unique_symbol_in_index(self, name: str, current_repo: str, current_mod: str) -> Optional[ResolvedSymbol]:
        """
        Import-aware, index-backed symbol lookup by bare name.
        Returns real, sliced code only.
        """
        try:
            with sqlite3.connect(self.index.db_path) as con:
                rows = con.execute(
                    "SELECT repo, relpath, qual, kind, line, code FROM symbols WHERE name=?",
                    (name,)
                ).fetchall()
        except Exception as e:
            log.debug("Index lookup failed for %s: %s", name, e)
            return None

        if not rows:
            return None

        quals = list({r[2] for r in rows})
        # trivial unique-case
        if len(quals) == 1:
            repo, relpath, qual, kind, line, code = rows[0]
            mod_qual = qual.rsplit(".", 1)[0]
            return ResolvedSymbol(
                qname=qual,
                name=name,
                kind=kind if kind in ("class", "function", "const") else self._head_kind(code),
                repo=repo,
                file_path=relpath,
                module_qual=mod_qual,
                module_path=mod_qual.replace(".", "/"),
                module_source=None,
                source_code=code,
                line=line,
                module_imports=None,
            )

        import_roots = self._collect_import_roots(current_repo, current_mod)
        candidates = [q for _, _, q, _, _, _ in rows]
        chosen = self._disambiguate_candidates(candidates, current_repo, current_mod, import_roots)
        if not chosen:
            return None

        for repo, relpath, qual, kind, line, code in rows:
            if qual == chosen:
                mod_qual = qual.rsplit(".", 1)[0]
                return ResolvedSymbol(
                    qname=qual,
                    name=name,
                    kind=kind if kind in ("class", "function", "const") else self._head_kind(code),
                    repo=repo,
                    file_path=relpath,
                    module_qual=mod_qual,
                    module_path=mod_qual.replace(".", "/"),
                    module_source=None,
                    source_code=code,
                    line=line,
                    module_imports=None,
                )
        return None

    def resolve_block_dependencies(self, target_name: str, candidate_file: Optional[str]):
        # Prioritize class definitions over constants/imports
        target = None
        # First try to find a class with this name
        target = next((s for s in self.import_graph.symbol_table.values() if s.name == target_name and s.kind == "class"), None)
        # If no class found, fall back to any symbol with this name
        if not target:
            target = next((s for s in self.import_graph.symbol_table.values() if s.name == target_name), None)
        if not target:
            return DependencyResolutionResult(
                target_symbol=target_name,
                resolved_dependencies={},
                unresolved_dependencies=[target_name],
                import_graph=self.import_graph,
                resolution_stats={"method": "not_found"},
                topological_order=[],
            )

        # 1) Initial closure
        resolved, unresolved, order = self._dfs_all(target)

        # 2) Bring in missing defs via DefinitionResolver (iterative closure)
        remaining = set(unresolved)
        for _ in range(2):
            if not remaining:
                break
            fetched = self.def_resolver.resolve_symbols(sorted(remaining))
            if not fetched:
                break
            self._inject_resolved_symbols(fetched)
            new_resolved, new_unresolved, new_order = self._dfs_all(target)
            # Merge the resolved dependencies instead of overwriting
            resolved.update(new_resolved)
            unresolved = new_unresolved
            order = new_order
            remaining = set(unresolved)

        # 3) Import-aware index lookup for any remaining bare names (real slices only)
        if remaining:
            repo, mod = self._module_ctx(target)
            unique_map: Dict[str, ResolvedSymbol] = {}
            for nm in list(remaining):
                rs = self._find_unique_symbol_in_index(nm, repo, mod)
                if rs:
                    unique_map[rs.qname] = rs
            if unique_map:
                self._inject_resolved_symbols(unique_map)
                new_resolved, new_unresolved, new_order = self._dfs_all(target)
                # Merge the resolved dependencies instead of overwriting
                resolved.update(new_resolved)
                unresolved = new_unresolved
                order = new_order
                remaining = set(unresolved)

        stats = {"resolved": len(resolved), "unresolved": len(remaining), "method": "index+resolver"}
        return DependencyResolutionResult(
            target_symbol=target.name,
            resolved_dependencies=resolved,
            unresolved_dependencies=sorted(remaining),
            import_graph=self.import_graph,
            resolution_stats=stats,
            topological_order=order
        )

    # ------------------------- Generation & validation ------------------------- #

    @staticmethod
    def _collect_import_lines_from_source(source: str) -> List[str]:
        """Normalize import lines from a module's source."""
        try:
            tree = ast.parse(source)
        except Exception:
            return []
        lines: List[str] = []
        for n in tree.body:
            if isinstance(n, ast.Import):
                for a in n.names:
                    if a.asname:
                        lines.append(f"import {a.name} as {a.asname}")
                    else:
                        lines.append(f"import {a.name}")
            elif isinstance(n, ast.ImportFrom):
                mod = ('.' * n.level) + (n.module or "")
                names = []
                for a in n.names:
                    if a.name == "*":
                        names.append("*")
                    elif a.asname:
                        names.append(f"{a.name} as {a.asname}")
                    else:
                        names.append(a.name)
                lines.append(f"from {mod} import {', '.join(names)}")
        return lines

    def _extract_imported_symbols(self, import_line: str) -> List[str]:
        """Extract symbol names from an import line."""
        try:
            if import_line.startswith("from ") and " import " in import_line:
                parts = import_line.split(" import ")
                if len(parts) == 2:
                    symbols_part = parts[1].strip()
                    symbols = []
                    for symbol in symbols_part.split(","):
                        symbol = symbol.strip()
                        if " as " in symbol:
                            symbol = symbol.split(" as ")[0].strip()
                        symbols.append(symbol)
                    return symbols
            return []
        except Exception:
            return []

    def _get_header_imports(self, qnames: List[str]) -> Set[str]:
        """Get the set of modules that are already imported in headers by contributing modules."""
        header_imports = set()
        for qname in qnames:
            sym = self.import_graph.symbol_table.get(qname)
            if sym:
                for key, mi in self.import_graph.modules.items():
                    if any(s.qualified_name == qname for s in mi.symbols.values()):
                        header_imports.update(mi.imports)
                        break
        return header_imports

    def _collect_import_lines_for_symbols(self, qnames: List[str], target_qname: str) -> List[str]:
        """Collect and de-dupe import lines from modules owning the given symbols (and the target)."""
        # Fundamental PyTorch classes that should always be imported directly from torch.nn
        FUNDAMENTAL_TORCH_CLASSES = {
            'Module', 'Parameter', 'Tensor', 'Buffer', 'Sequential', 'ModuleList', 'ModuleDict'
        }
        
        module_keys: Set[str] = set()
        for q in list(qnames) + [target_qname]:
            if not q:
                continue
            sym = self.import_graph.symbol_table.get(q)
            if not sym:
                continue
            for key, mi in self.import_graph.modules.items():
                if any(s.qualified_name == q for s in mi.symbols.values()):
                    module_keys.add(key)
                    break

        import_lines: Set[str] = set()
        for key in module_keys:
            mi = self.import_graph.modules.get(key)
            if not mi or not mi.source_code:
                continue
            for line in self._collect_import_lines_from_source(mi.source_code):
                import_lines.add(line)

        # Filter imports to only include symbols that are actually used in the emitted code
        emitted_symbols = set()
        for qname in qnames:
            sym = self.import_graph.symbol_table.get(qname)
            if sym:
                emitted_symbols.add(sym.name)

        filtered = []
        header_imports = self._get_header_imports(qnames)
        for line in sorted(import_lines):
            if line.startswith("from .") or line.startswith("from .."):
                continue
            import_symbols = self._extract_imported_symbols(line)
            
            # Skip imports of fundamental PyTorch classes from internal modules
            if any(sym in FUNDAMENTAL_TORCH_CLASSES for sym in import_symbols):
                if line.startswith("from torch.nn.modules.") or line.startswith("from torch.nn.parameter"):
                    continue
            
            # Only include this import if at least one of its symbols is actually used
            if import_symbols:
                used_symbols = [sym for sym in import_symbols if sym in emitted_symbols]
                if not used_symbols:
                    # None of the imported symbols are used, skip this import
                    continue
                # If only some symbols are used, create a filtered import line
                if len(used_symbols) < len(import_symbols):
                    # Reconstruct the import line with only used symbols
                    if line.startswith("from "):
                        module_part = line.split(" import ")[0]
                        line = f"{module_part} import {', '.join(used_symbols)}"
                    else:
                        # Handle 'import module' case
                        continue  # Skip module imports if we can't determine usage
            
            # Already covered by header imports?
            redundant = False
            for h in header_imports:
                if line.startswith(f"import {h}") or line.startswith(f"from {h}"):
                    redundant = True
                    break
            if not redundant:
                filtered.append(line)
        return filtered

    def _analyze_import_usage(self, source_code: str) -> Set[str]:
        """Analyze source code to determine what imports are actually used."""
        try:
            tree = ast.parse(source_code)
        except Exception:
            # Try to fix common issues like indentation
            try:
                # Strip leading indentation if it's causing parsing issues
                lines = source_code.splitlines()
                if lines:
                    # Find the minimum indentation
                    min_indent = float('inf')
                    for line in lines:
                        if line.strip():  # Skip empty lines
                            indent = len(line) - len(line.lstrip())
                            min_indent = min(min_indent, indent)
                    # Strip the minimum indentation from all lines
                    if min_indent != float('inf'):
                        fixed_code = '\n'.join(line[min_indent:] if line.strip() else line for line in lines)
                        tree = ast.parse(fixed_code)
                    else:
                        return set()
                else:
                    return set()
            except Exception:
                return set()

        used_names = set()

        class UsageVisitor(ast.NodeVisitor):
            def visit_Name(self, node):
                if isinstance(node.ctx, ast.Load):
                    used_names.add(node.id)
                self.generic_visit(node)

            def visit_Attribute(self, node):
                if isinstance(node.ctx, ast.Load):
                    if isinstance(node.value, ast.Name):
                        used_names.add(node.value.id)
                    elif isinstance(node.value, ast.Attribute):
                        if isinstance(node.value.value, ast.Name):
                            used_names.add(node.value.value.id)
                self.generic_visit(node)
            
            def visit_FunctionDef(self, node):
                # Check docstring for type annotations
                if node.body and isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant):
                    docstring = node.body[0].value.value
                    if isinstance(docstring, str):
                        # Look for :obj:`TypeName` patterns in docstrings
                        import re
                        type_matches = re.findall(r':obj:`([A-Za-z_][A-Za-z0-9_]*)`', docstring)
                        for type_name in type_matches:
                            used_names.add(type_name)
                
                # Check return type annotation
                if node.returns:
                    self.visit(node.returns)
                
                # Check parameter type annotations
                for arg in node.args.args:
                    if arg.annotation:
                        self.visit(arg.annotation)
                
                self.generic_visit(node)
            
            def visit_ClassDef(self, node):
                # Check docstring for type annotations
                if node.body and isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant):
                    docstring = node.body[0].value.value
                    if isinstance(docstring, str):
                        # Look for :obj:`TypeName` patterns in docstrings
                        import re
                        type_matches = re.findall(r':obj:`([A-Za-z_][A-Za-z0-9_]*)`', docstring)
                        for type_name in type_matches:
                            used_names.add(type_name)
                
                # Check base class type annotations
                for base in node.bases:
                    self.visit(base)
                
                # Check keyword arguments in base classes
                for keyword in node.keywords:
                    self.visit(keyword)
                
                self.generic_visit(node)

        UsageVisitor().visit(tree)
        return used_names

    def _collect_required_imports(self, symbols: List[str], target_source: str) -> List[str]:
        """Dynamically collect required imports based on actual symbol usage."""
        required_imports = []
        target_usage = self._analyze_import_usage(target_source)
        all_imports = set()
        
        # Analyze all dependency source code for import usage
        for qname in symbols:
            sym = self.import_graph.symbol_table.get(qname)
            if sym:
                sym_imports = self._analyze_import_usage(sym.source_code)
                all_imports.update(sym_imports)
        
        # Also analyze the target source
        all_imports.update(target_usage)

        if "torch" in all_imports:
            required_imports.append("import torch")
        if "nn" in all_imports:
            required_imports.append("import torch.nn as nn")
        if "F" in all_imports:
            required_imports.append("import torch.nn.functional as F")
        # PyTorch-specific imports
        if "Module" in all_imports:
            required_imports.append("from torch.nn import Module")
        if "Parameter" in all_imports:
            required_imports.append("from torch.nn.parameter import Parameter")
        if "Tensor" in all_imports:
            required_imports.append("from torch import Tensor")

        # Standard typing imports (excluding PyTorch-specific types)
        typing_imports = []
        for item in ["Any", "Dict", "List", "Optional", "Tuple", "Union", "Callable", "TypeVar", "Generic", "Literal", "TypedDict", "Protocol", "runtime_checkable", "overload", "final", "ParamSpec", "TypeAlias", "NotRequired", "Required"]:
            if item in all_imports:
                typing_imports.append(item)
        if typing_imports:
            required_imports.append(f"from typing import {', '.join(typing_imports)}")

        if "Enum" in all_imports:
            required_imports.append("from enum import Enum")
        if "chain" in all_imports:
            required_imports.append("from itertools import chain")
        
        # Add other standard library imports
        if "os" in all_imports:
            required_imports.append("import os")
        if "osp" in all_imports:
            required_imports.append("import os.path as osp")
        if "re" in all_imports:
            required_imports.append("import re")
        if "math" in all_imports:
            required_imports.append("import math")
        if "random" in all_imports:
            required_imports.append("import random")
        if "sys" in all_imports:
            required_imports.append("import sys")
        if "warnings" in all_imports:
            required_imports.append("import warnings")
        if "getuser" in all_imports or "gethostname" in all_imports:
            required_imports.append("from getpass import getuser")
            required_imports.append("from socket import gethostname")
        if "handlers" in all_imports:
            required_imports.append("from logging import handlers")
        
        # Add collections imports
        if "defaultdict" in all_imports or "OrderedDict" in all_imports or "Counter" in all_imports:
            collections_items = []
            if "defaultdict" in all_imports:
                collections_items.append("defaultdict")
            if "OrderedDict" in all_imports:
                collections_items.append("OrderedDict")
            if "Counter" in all_imports:
                collections_items.append("Counter")
            if collections_items:
                required_imports.append(f"from collections import {', '.join(collections_items)}")
        
        # Add collections.abc imports
        if "Sequence" in all_imports:
            required_imports.append("from collections.abc import Sequence")
        
        # Add logging imports
        if "Logger" in all_imports or "FileHandler" in all_imports or "LogRecord" in all_imports:
            logging_items = []
            if "Logger" in all_imports:
                logging_items.append("Logger")
            if "FileHandler" in all_imports:
                logging_items.append("FileHandler")
            if "LogRecord" in all_imports:
                logging_items.append("LogRecord")
            if logging_items:
                required_imports.append(f"from logging import {', '.join(logging_items)}")
        
        # Add logging module import separately
        if "logging" in all_imports:
            required_imports.append("import logging")
        
        # Add copy import
        if "copy" in all_imports:
            required_imports.append("import copy")
        
        # Add pathlib imports
        if "Path" in all_imports:
            required_imports.append("from pathlib import Path")
        
        # Add numpy import
        if "np" in all_imports or "numpy" in all_imports:
            required_imports.append("import numpy as np")
        
        # Add other common imports
        if "yaml" in all_imports:
            required_imports.append("import yaml")
        if "rich" in all_imports:
            required_imports.append("import rich")
        if "cv2" in all_imports:
            required_imports.append("import cv2")
        if "pandas" in all_imports:
            required_imports.append("import pandas as pd")
        if "matplotlib" in all_imports:
            required_imports.append("import matplotlib.pyplot as plt")
        
        # Note: We don't add hardcoded imports here anymore
        # The system will resolve dependencies to their actual definitions
        # and include the full source code in the generated file
        
        # Note: We don't add hardcoded imports here anymore
        # The system will resolve dependencies to their actual definitions
        # and include the full source code in the generated file



        return required_imports

    def _emit_single_file(self, block_name: str, source_info: Dict[str, Any], deps: DependencyResolutionResult) -> Dict[str, Any]:
        out_dir = Path("generated_packages")
        out_dir.mkdir(parents=True, exist_ok=True)
        outfile = out_dir / f"{block_name}.py"

        out_lines: List[str] = []
        out_lines += [
            f"# Auto-generated single-file for {block_name}",
            "# Dependencies are emitted in topological order (utilities first).",
        ]

        order = deps.topological_order or []
        seen_defs: Set[str] = set()
        emitted_q: Set[str] = set()
        
        # First pass: build class name to qualified name mapping for ALL classes
        class_dependencies = {}  # class_name -> qualified_name mapping
        
        for q in order:
            if q in emitted_q:
                continue
            s = deps.resolved_dependencies.get(q)
            if not s:
                continue
            if s.name == block_name or q.endswith(f".{block_name}"):
                continue
            
            # Build class name to qualified name mapping
            class_dependencies[s.name] = q
        
        # Second pass: analyze inheritance relationships between all classes
        inheritance_map = {}  # child -> parent mapping
        
        for q in order:
            if q in emitted_q:
                continue
            s = deps.resolved_dependencies.get(q)
            if not s:
                continue
            if s.name == block_name or q.endswith(f".{block_name}"):
                continue
            
            # Analyze inheritance relationships for classes
            if s.source_code and "class " in s.source_code:
                cls_lines = s.source_code.splitlines()
                for line in cls_lines:
                    if line.strip().startswith("class ") and "(" in line and ")" in line:
                        # Extract parent class name from class ClassName(ParentClass):
                        parent_match = re.search(r'class\s+(\w+)\s*\(([^)]+)\)', line)
                        if parent_match:
                            child_name = parent_match.group(1).strip()
                            parent_name = parent_match.group(2).strip()
                            # Store inheritance relationship
                            if parent_name in class_dependencies:
                                inheritance_map[q] = class_dependencies[parent_name]
                    break

        # Second pass: build dependency order respecting inheritance
        processed_classes = set()
        ordered_classes = []
        
        def add_class_with_parents(class_qname):
            """Recursively add class and its parents in correct order"""
            if class_qname in processed_classes:
                return
            if class_qname in inheritance_map:
                # Add parent first
                parent_qname = inheritance_map[class_qname]
                add_class_with_parents(parent_qname)
            # Add this class
            ordered_classes.append(class_qname)
            processed_classes.add(class_qname)
        
        # Process all classes respecting inheritance order
        for q in order:
            if q in emitted_q:
                continue
            s = deps.resolved_dependencies.get(q)
            if not s:
                continue
            if s.name == block_name or q.endswith(f".{block_name}"):
                continue
            
            if s.source_code and "class " in s.source_code:
                add_class_with_parents(q)
        
        # Third pass: categorize remaining dependencies by priority
        type_aliases_and_constants = []
        utility_functions = []
        classes_and_main_functions = []
        
        for q in order:
            if q in emitted_q:
                continue
            s = deps.resolved_dependencies.get(q)
            if not s:
                continue
            if s.name == block_name or q.endswith(f".{block_name}"):
                continue
            
            # Skip classes that are already ordered
            if q in ordered_classes:
                continue
            
            # Categorize remaining dependencies by priority
            original_sym = self.import_graph.symbol_table.get(q)
            if original_sym:
                if original_sym.kind == "const" or "Union" in s.source_code or "TypeVar" in s.source_code:
                    # Type aliases, constants, and type definitions come first
                    type_aliases_and_constants.append(q)
                elif original_sym.kind == "class" and ("Enum" in s.source_code or "class" in s.source_code and "Enum" in s.source_code):
                    # Enum classes and classes used in type aliases come first
                    type_aliases_and_constants.append(q)
                elif original_sym.kind == "function" and not any(keyword in s.source_code for keyword in ["class ", "def __init__", "def forward"]):
                    # Utility functions come second
                    utility_functions.append(q)
                else:
                    # Classes and main functions come last
                    classes_and_main_functions.append(q)
            else:
                # Default to utility functions if we can't determine the kind
                utility_functions.append(q)
        
        # Track imported types to prevent duplicates
        imported_types = set()
        
        # Emit dependencies in correct order: inheritance-ordered classes first, then others by priority
        for q in ordered_classes + type_aliases_and_constants + utility_functions + classes_and_main_functions:
            if q in emitted_q:
                continue
            s = deps.resolved_dependencies.get(q)
            if not s:
                continue
            
            head_name = self._top_level_def_name(s.source_code)
            if head_name and head_name in seen_defs:
                continue
            
            # Check if this dependency contains types that are already imported
            if self._contains_duplicate_types(s.source_code, imported_types):
                continue
                
            out_lines.append(f"# ---- {q} ----")
            
            # Clean and validate the source code before emitting
            cleaned_source = self._clean_and_validate_source_code(s.source_code, q)
            if cleaned_source:
                out_lines.append(cleaned_source.rstrip())
                out_lines.append("")
                if head_name:
                    seen_defs.add(head_name)
                emitted_q.add(q)
                
                # Track the types that were imported
                self._track_imported_types(cleaned_source, imported_types)
            else:
                # Skip this dependency if it couldn't be cleaned properly
                continue
        
        # Now collect and add required imports AFTER dependency analysis
        all_dependency_symbols = ordered_classes + type_aliases_and_constants + utility_functions + classes_and_main_functions
        qnames_for_imports = [q for q in all_dependency_symbols if q in deps.resolved_dependencies]
        required_imports = self._collect_required_imports(qnames_for_imports, source_info["content"])
        
        # Also collect dynamic imports from contributing modules
        target_qname = None
        for q, s in deps.resolved_dependencies.items():
            if s.name == block_name or q.endswith(f".{block_name}"):
                target_qname = q
                break
        if target_qname is None:
            for q, s in deps.resolved_dependencies.items():
                if s.name == block_name:
                    target_qname = q
                    break
        
        dyn_imports = self._collect_import_lines_for_symbols(qnames_for_imports, target_qname or "")
        
        # Build the final file with imports at the very top
        final_lines: List[str] = []
        final_lines.append(f"# Auto-generated single-file for {block_name}")
        final_lines.append("# Dependencies are emitted in topological order (utilities first).")
        
        # Add imports at the very top
        if required_imports:
            final_lines.append("# Standard library and external imports")
            final_lines.extend(required_imports)
            final_lines.append("")
        
        if dyn_imports:
            final_lines.append("# ---- original imports from contributing modules ----")
            final_lines.extend(dyn_imports)
            final_lines.append("")
        
        # Add all the dependency content we gathered into out_lines
        # (skip the 2 header lines we already added above)
        final_lines.extend(out_lines[2:])
        
        # Replace the lines list with our properly ordered final_lines
        lines = final_lines
        
        # Ensure the first dependency starts with a class definition, not a method
        # Calculate header length dynamically instead of using magic indices
        header_length = len(final_lines)
        if header_length > 0:
            # Find the first non-empty, non-comment line after the header
            first_content_line = -1
            for i, line in enumerate(lines):
                stripped = line.strip()
                if stripped and not stripped.startswith('#') and not stripped.startswith('"""') and not stripped.startswith("'''"):
                    first_content_line = i
                    break
            
            if first_content_line > 0 and lines[first_content_line].strip().startswith('def '):
                # The first dependency is missing its class declaration
                # Find the first dependency that should be a class
                if ordered_classes:
                    first_class_qname = ordered_classes[0]
                    first_class_sym = deps.resolved_dependencies.get(first_class_qname)
                    if first_class_sym:
                        # Insert the class declaration at the beginning of the dependency content
                        class_header = f"# ---- {first_class_qname} ----"
                        lines.insert(first_content_line, class_header)
                        lines.insert(first_content_line + 1, first_class_sym.source_code.rstrip())
                        lines.insert(first_content_line + 2, "")
        


        lines.append(f"# ---- {block_name} (target) ----")
        lines.append(source_info["content"].rstrip())
        lines.append("")

        result = sanitize_generated_block("\n".join(lines))

        outfile.write_text(result, encoding="utf-8")
        try:
            compile(outfile.read_text(encoding="utf-8"), str(outfile), "exec")
            ok = True
            err = []
        except Exception as e:
            ok = False
            err = [str(e)]

        return {"success": ok, "file_path": str(outfile), "errors": err, "mode": "single_file"}

    @staticmethod
    def _top_level_def_name(code: str) -> Optional[str]:
        m = re.search(r"^\s*(class|def)\s+([A-Za-z_][A-Za-z0-9_]*)\b", code, re.MULTILINE)
        return m.group(2) if m else None
    
    @staticmethod
    def _dedent_code_block(code: str) -> str:
        """
        Safely dedent a code block while preserving relative indentation.
        This is especially important for multi-line blocks like try/except, if/else, etc.
        
        Args:
            code: The source code to dedent
            
        Returns:
            Dedented code with relative indentation preserved
        """
        if not code or not code.strip():
            return code
            
        lines = code.splitlines()
        if not lines:
            return code
            
        # Find the minimum indentation from non-empty lines
        min_indent = float('inf')
        for line in lines:
            if line.strip():  # Skip empty lines
                indent = len(line) - len(line.lstrip())
                min_indent = min(min_indent, indent)
        
        # If no indentation to strip or already at top level, return as-is
        if min_indent == float('inf') or min_indent == 0:
            return code
            
        # Strip the minimum indentation while preserving relative structure
        dedented_lines = []
        for line in lines:
            if line.strip():  # Non-empty line
                # Strip the minimum indentation
                dedented_lines.append(line[min_indent:])
            else:
                # Keep empty lines as-is
                dedented_lines.append(line)
        
        return '\n'.join(dedented_lines)

    def _clean_and_validate_source_code(self, source_code: str, qname: str) -> Optional[str]:
        """Clean and validate source code before emitting to prevent syntax errors."""
        if not source_code or not source_code.strip():
            return None
        
        try:
            # First, try to parse the code to check for syntax errors
            ast.parse(source_code)
        except SyntaxError:
            # If there's a syntax error, try to clean it up
            cleaned = self._clean_malformed_code(source_code, qname)
            if cleaned:
                try:
                    # Verify the cleaned code parses
                    ast.parse(cleaned)
                    return cleaned
                except SyntaxError:
                    # If still can't parse, skip this dependency
                    return None
            else:
                return None
        
        # If the code parses correctly, clean it up for emission
        original_sym = self.import_graph.symbol_table.get(qname)
        if original_sym and original_sym.kind == "const":
            # Use the improved dedenting method that preserves relative indentation
            dedented_code = self._dedent_code_block(source_code)
            # Also replace relative imports in constants
            cleaned_code = self._replace_relative_imports(dedented_code)
            return cleaned_code
        else:
            # Replace relative imports with absolute imports in the source code
            cleaned_source = self._replace_relative_imports(source_code)
            return cleaned_source

    def _clean_malformed_code(self, source_code: str, qname: str) -> Optional[str]:
        """Attempt to clean malformed code that has syntax errors."""
        if not source_code:
            return None
        
        lines = source_code.splitlines()
        cleaned_lines = []
        
        # Look for common patterns that cause issues
        in_import_section = False
        import_lines = []
        in_parentheses = False
        paren_count = 0
        
        for line in lines:
            stripped = line.strip()
            
            # Skip empty lines
            if not stripped:
                continue
            
            # Track parentheses for multi-line imports
            if '(' in stripped:
                paren_count += stripped.count('(')
            if ')' in stripped:
                paren_count += stripped.count(')')
            
            # Check if we're in a multi-line import
            if paren_count > 0:
                in_parentheses = True
            else:
                in_parentheses = False
            
            # Check if this is an import section
            if stripped.startswith('from ') and ' import ' in stripped:
                in_import_section = True
                import_lines.append(line)
                continue
            elif stripped.startswith('import '):
                in_import_section = True
                import_lines.append(line)
                continue
            
            # If we're in an import section and hit non-import code, end the section
            if in_import_section and not (stripped.startswith('from ') or stripped.startswith('import ') or in_parentheses):
                in_import_section = False
                # Add all collected import lines
                cleaned_lines.extend(import_lines)
                import_lines = []
            
            # Skip lines that are just type names without proper context
            if stripped.endswith(',') and not stripped.startswith(('def ', 'class ', 'if ', 'for ', 'while ')):
                # This might be a malformed import line, skip it
                continue
            
            # Skip lines that are just type names (common PyTorch type aliases)
            pytorch_types = [
                '_ratio_2_t,', '_ratio_3_t,', '_size_1_t,', '_size_2_opt_t,', '_size_2_t,',
                '_size_3_opt_t,', '_size_3_t,', '_size_any_opt_t,', '_size_any_t,',
                '_ratio_2_t', '_ratio_3_t', '_size_1_t', '_size_2_opt_t', '_size_2_t',
                '_size_3_opt_t', '_size_3_t', '_size_any_opt_t', '_size_any_t'
            ]
            if stripped in pytorch_types:
                continue
            
            # Skip orphaned closing parentheses
            if stripped == ')' and not in_parentheses:
                continue
            
            # Add valid lines
            cleaned_lines.append(line)
        
        # Add any remaining import lines
        if import_lines:
            cleaned_lines.extend(import_lines)
        
        if not cleaned_lines:
            return None
        
        return '\n'.join(cleaned_lines)

    def _contains_duplicate_types(self, source_code: str, imported_types: set) -> bool:
        """Check if source code contains types that are already imported."""
        if not source_code:
            return False
        
        # Common PyTorch type aliases to check for
        pytorch_types = {
            '_ratio_2_t', '_ratio_3_t', '_size_1_t', '_size_2_opt_t', '_size_2_t',
            '_size_3_opt_t', '_size_3_t', '_size_any_opt_t', '_size_any_t'
        }
        
        for line in source_code.splitlines():
            stripped = line.strip()
            for type_name in pytorch_types:
                if type_name in stripped and type_name in imported_types:
                    return True
        
        return False

    def _track_imported_types(self, source_code: str, imported_types: set) -> None:
        """Track which types have been imported to prevent duplicates."""
        if not source_code:
            return
        
        # Common PyTorch type aliases to track
        pytorch_types = {
            '_ratio_2_t', '_ratio_3_t', '_size_1_t', '_size_2_opt_t', '_size_2_t',
            '_size_3_opt_t', '_size_3_t', '_size_any_opt_t', '_size_any_t'
        }
        
        for line in source_code.splitlines():
            stripped = line.strip()
            for type_name in pytorch_types:
                if type_name in stripped:
                    imported_types.add(type_name)
    
    def _replace_relative_imports(self, source_code: str) -> str:
        """Replace relative imports with absolute imports in source code."""
        lines = source_code.splitlines()
        cleaned_lines = []
        i = 0
        
        while i < len(lines):
            line = lines[i]
            
            # Replace relative imports with absolute imports
            if line.strip().startswith("from .weight_init import"):
                line = line.replace("from .weight_init import", "from mmengine.model.weight_init import")
            elif line.strip().startswith("from .wrappers.utils import"):
                line = line.replace("from .wrappers.utils import", "from mmengine.model.wrappers.utils import")
            elif line.strip().startswith("from .utils import"):
                line = line.replace("from .utils import", "from mmdet.utils import")
            elif line.strip().startswith("from ."):
                # Generic relative import replacement - handle multiple dots correctly
                # Replace "from .something" with "from mmengine.something"
                # Replace "from ..something" with "from mmengine.something"
                # Replace "from ...something" with "from mmengine.something"
                import re
                # Match "from \.+(\w+.*)" and replace with "from mmengine.\1"
                line = re.sub(r'from \.+(\w+.*)', r'from mmengine.\1', line)
            
            # Handle multi-line imports that should be commented out
            if line.strip().startswith("from mmdet.utils import MultiConfig, OptConfigType"):
                line = "# from mmdet.utils import MultiConfig, OptConfigType  # Already imported at top"
            elif line.strip().startswith("from mmdet.utils import OptConfigType"):
                line = "# from mmdet.utils import OptConfigType  # Already imported at top"
            elif line.strip().startswith("from mmcv.cnn.bricks import Swish, build_norm_layer"):
                line = "# from mmcv.cnn.bricks import Swish, build_norm_layer  # Already imported at top"
            elif line.strip().startswith("from mmengine.runner.checkpoint import"):
                # Handle multi-line checkpoint import
                line = "# from mmengine.runner.checkpoint import (_load_checkpoint_with_prefix, load_checkpoint, load_state_dict)  # Already imported at top"
                # Skip the continuation lines
                while i + 1 < len(lines) and lines[i + 1].strip().startswith(("(", "load_checkpoint", "load_state_dict")):
                    i += 1
                    continue
            
            cleaned_lines.append(line)
            i += 1
        
        return "\n".join(cleaned_lines)

    # ------------------------------ Warm index once --------------------------- #

    def warm_index_once(self) -> bool:
        """
        Prepare caches and build the import graph once per process.
        Indexing policy is controlled by self.index_mode:
          - "missing": only index repos not present in SQLite index
          - "force": index all repos
          - "skip": don't index (assume prebuilt DB)
        """
        if self._index_warmed:
            return True
        if not self._ensure_all_repos_cached():
            return False

        # Index according to policy
        for repo in list(self.repo_cache.repos.keys()):
            if not self.repo_cache.is_repo_cached(repo):
                continue

            if self.index_mode == "skip":
                if self.index.repo_has_index(repo):
                    log.info("Index present for %s — hydrating from SQLite.", repo)
                    self._hydrate_repo_from_index(repo)
                continue

            if self.index_mode == "missing":
                if self.index.repo_has_index(repo):
                    log.info("Index present for %s — skipping indexing and hydrating from SQLite.", repo)
                    self._hydrate_repo_from_index(repo)
                    continue
                # no index yet → index now
                self.index_repository(repo)
                continue

            # index_mode == "force"
            self.index_repository(repo)

        # Build package→repo map (derived from index)
        self._refresh_package_repo_map()

        self._index_warmed = True
        return True

    # ------------------------------ Top-level API ------------------------------ #

    def extract_block(self, block_name: str) -> Dict[str, Any]:
        log.info("Extracting block: %s", block_name)

        # Warm caches + build index ONCE (policy-controlled)
        if not self.warm_index_once():
            return {"success": False, "reason": "cache failure", "block_name": block_name}

        # Discover (uses the already-built import graph + repo_cache)
        discoveries = self._discover_blocks([block_name])
        candidates = discoveries.get(block_name) or []
        if not candidates:
            self.failed_blocks.append(block_name)
            self._save_results()
            return {"success": False, "reason": "block not found", "block_name": block_name}

        best = candidates[0]
        repo = best["repository"]
        rel = best["file_path"]
        repo_root = self.repo_cache.get_cached_repo(repo)
        content = (repo_root / rel).read_text(encoding="utf-8", errors="ignore")

        # Exact target slice (preserve decorators/comments)
        target_code = self._extract_named_block(content, block_name)
        if not target_code:
            self.failed_blocks.append(block_name)
            self._save_results()
            return {"success": False, "reason": "could not slice block", "block_name": block_name}

        source_info = {"repository": repo, "file_path": rel, "content": target_code, "identifier": block_name}

        # Resolve all dependencies using the warmed import graph
        deps = self.resolve_block_dependencies(block_name, rel)

        if deps.unresolved_dependencies:
            self.failed_blocks.append(block_name)
            self._save_results()
            return {
                "success": False,
                "reason": "unresolved dependencies",
                "block_name": block_name,
                "unresolved": deps.unresolved_dependencies
            }

        gen = self._emit_single_file(block_name, source_info, deps)

        result = {
            "success": gen["success"],
            "block_name": block_name,
            "file_path": gen.get("file_path"),
            "errors": gen.get("errors", []),
            "dependencies": {
                "resolved": deps.resolution_stats.get("resolved", 0),
                "unresolved": deps.resolution_stats.get("unresolved", 0)
            },
            "timestamp": datetime.now().isoformat()
        }
        if gen["success"]:
            self.extracted_blocks.append(result)
        else:
            self.failed_blocks.append(block_name)
        self._save_results()
        return result

    # ------------------------------ Helpers ----------------------------------- #

    @staticmethod
    def _extract_named_block(content: str, name: str) -> Optional[str]:
        """Return the full text of class/def named `name`, including body + decorators."""
        try:
            tree = ast.parse(content)
        except Exception:
            return None
        for node in tree.body:
            if isinstance(node, (ast.ClassDef, ast.FunctionDef)) and node.name == name:
                if hasattr(node, "lineno") and hasattr(node, "end_lineno"):
                    lines = content.splitlines()
                    start = node.lineno - 1
                    while start > 0 and lines[start - 1].lstrip().startswith("@"):
                        start -= 1
                    return "\n".join(lines[start: node.end_lineno])
                if hasattr(ast, "unparse"):
                    return ast.unparse(node)
        return None


# ----------------------------------------------------------------------------- #
# CLI
# ----------------------------------------------------------------------------- #
def _load_block_list(json_path: Path) -> List[str]:
    if not json_path.exists():
        raise FileNotFoundError(f"Names file not found: {json_path}")
    data = json.loads(json_path.read_text(encoding='utf-8'))
    if not isinstance(data, list):
        raise ValueError("names JSON must be a JSON array of strings.")
    out: List[str] = []
    for x in data:
        if isinstance(x, str) and x.strip():
            out.append(x.strip())
    seen = set()
    uniq = []
    for n in out:
        if n not in seen:
            seen.add(n)
            uniq.append(n)
    return uniq


def main():
    p = argparse.ArgumentParser(description="Extract reusable PyTorch blocks")
    p.add_argument("--block", type=str, help="Extract a specific block by name")
    p.add_argument("--blocks", nargs="+", help="Extract multiple blocks")
    p.add_argument("--retry-failed", action="store_true", help="Retry failed blocks from previous run")
    p.add_argument("--names-json", type=Path, default=Path(__file__).parent / "config" / "nn_block_names.json",
                   help="If --block/--blocks not provided, read names from this JSON (default: ab/rag/config/nn_block_names.json)")
    p.add_argument("--limit", type=int, default=None, help="Max number of names to process from the list")
    p.add_argument("--start-from", type=str, default=None,
                   help="Skip names until (and including) this one, then start")
    p.add_argument("--skip-existing", action="store_true",
                   help="Skip blocks that already appear in extraction_results.json (default: process all blocks)")
    p.add_argument("--stop-on-fail", action="store_true", help="Stop batch as soon as one extraction fails")
    p.add_argument("--progress-every", type=int, default=10, help="Log progress every N blocks")
    p.add_argument("--index-mode", type=str, choices=("missing", "force", "skip"),
                   default="missing", help="Indexing policy: missing (default), force, or skip")
    p.add_argument("--no-validate", action="store_true", help="Disable automatic validation and movement of valid blocks to 'block' directory")
    p.add_argument("--cleanup-invalid", action="store_true", help="Remove invalid blocks after validation")
    args = p.parse_args()

    extractor = BlockExtractor(index_mode=args.index_mode)

    # Warm caches + build index ONCE (policy-controlled)
    ok = extractor.warm_index_once()
    if not ok:
        print(json.dumps({"success": False, "reason": "cache failure during warm_index"}, indent=2))
        return

    def run_one(name: str) -> Dict[str, Any]:
        try:
            res = extractor.extract_block(name)
            return res
        except Exception as e:
            return {"success": False, "reason": f"exception: {type(e).__name__}: {e}", "block_name": name}
    
    def validate_block_parallel(name: str) -> Tuple[str, Dict[str, Any]]:
        """Validate a block in parallel and return (name, validation_result)."""
        try:
            if not args.no_validate:
                validator = BlockValidator()
                is_valid, error = validator.validate_and_move_block(name)
                
                validation_result = {
                    "name": name,
                    "status": "valid" if is_valid else "invalid",
                    "moved_to_block_dir": is_valid,
                    "error": error if not is_valid else None
                }
                
                # Cleanup invalid blocks if requested
                if args.cleanup_invalid and not is_valid:
                    invalid_file = Path("generated_packages") / f"{name}.py"
                    if invalid_file.exists():
                        invalid_file.unlink()
                        validation_result["cleaned_up"] = True
                
                return name, validation_result
            else:
                return name, {"name": name, "status": "skipped", "reason": "validation disabled"}
        except Exception as e:
            return name, {"name": name, "status": "validation_error", "error": str(e)}

    # Direct single
    if args.block:
        res = run_one(args.block)
        
        # Validate the single block
        if not args.no_validate and res.get("success"):
            name, validation_result = validate_block_parallel(args.block)
            res["validation"] = validation_result
        
        print(json.dumps(res, indent=2))
        return

    # Direct list
    if args.blocks:
        results = {b: run_one(b) for b in args.blocks}
        
        # Validate all blocks in parallel
        if not args.no_validate:
            with cf.ThreadPoolExecutor(max_workers=min(len(args.blocks), 4)) as executor:
                validation_futures = [executor.submit(validate_block_parallel, b) for b in args.blocks]
                validation_results = {}
                
                for future in cf.as_completed(validation_futures):
                    name, validation_result = future.result()
                    validation_results[name] = validation_result
                
                # Add validation results to extraction results
                for name in args.blocks:
                    if name in validation_results:
                        results[name]["validation"] = validation_results[name]
        
        print(json.dumps(results, indent=2))
        return

    # Retry failed
    if args.retry_failed:
        retried: Dict[str, Any] = {}
        failed_blocks_list = list(extractor.failed_blocks)
        
        for item in failed_blocks_list:
            retried[item] = run_one(item)
            if args.stop_on_fail and not retried[item].get("success"):
                break
        
        # Validate retried blocks in parallel
        if not args.no_validate:
            successful_retries = [item for item in failed_blocks_list if retried[item].get("success")]
            if successful_retries:
                with cf.ThreadPoolExecutor(max_workers=min(len(successful_retries), 4)) as executor:
                    validation_futures = [executor.submit(validate_block_parallel, item) for item in successful_retries]
                    validation_results = {}
                    
                    for future in cf.as_completed(validation_futures):
                        name, validation_result = future.result()
                        validation_results[name] = validation_result
                    
                    # Add validation results to retry results
                    for name in successful_retries:
                        if name in validation_results:
                            retried[name]["validation"] = validation_results[name]
        
        print(json.dumps(retried, indent=2))
        return

    # Default mode: read nn_block_names.json and process one by one
    try:
        names = _load_block_list(args.names_json)
    except Exception as e:
        print(json.dumps({"success": False, "reason": f"names file error: {e}", "path": str(args.names_json)}, indent=2))
        return

    # Default behavior: process all blocks every time
    # Only skip already-extracted blocks if --skip-existing is specified (opposite of --redo-existing)
    already = {x.get("block_name") for x in extractor.extracted_blocks} if extractor.extracted_blocks else set()
    if args.start_from:
        if args.start_from in names:
            pos = names.index(args.start_from)
            names = names[pos + 1 :]
        # else ignore

    plan = []
    for n in names:
        if args.limit and len(plan) >= args.limit:
            break
        # Process all blocks by default, only skip if --skip-existing is specified
        # (This is the opposite of the original --redo-existing behavior)
        if hasattr(args, 'skip_existing') and args.skip_existing and n in already:
            continue
        plan.append(n)

    if not plan:
        print(json.dumps({"success": True, "reason": "nothing to do (empty plan)"}, indent=2))
        return

    print(f"Starting batch extraction for {len(plan)} block(s) from {args.names_json} ...")
    batch_results: Dict[str, Any] = {}
    t0 = time.time()
    
    # Process blocks with immediate validation after each extraction
    for i, name in enumerate(plan, 1):
        # Extract single block
        res = run_one(name)
        batch_results[name] = res
        
        # Validate immediately after extraction (unless --no-validate is specified)
        if not args.no_validate and res.get("success"):
            try:
                validator = BlockValidator()
                is_valid, error = validator.validate_and_move_block(name)
                
                validation_result = {
                    "name": name,
                    "status": "valid" if is_valid else "invalid",
                    "moved_to_block_dir": is_valid,
                    "error": error if not is_valid else None
                }
                
                # Cleanup invalid blocks if requested
                if args.cleanup_invalid and not is_valid:
                    invalid_file = Path("generated_packages") / f"{name}.py"
                    if invalid_file.exists():
                        invalid_file.unlink()
                        validation_result["cleaned_up"] = True
                
                # Add validation result to batch results
                batch_results[name]["validation"] = validation_result
                
                # Log validation result
                if is_valid:
                    print(f"✓ {name}: Validated and moved to block directory")
                else:
                    print(f"✗ {name}: Validation failed - {error}")
                    
            except Exception as e:
                validation_result = {"name": name, "status": "validation_error", "error": str(e)}
                batch_results[name]["validation"] = validation_result
                print(f"✗ {name}: Validation error - {e}")
        
        # Save progress incrementally
        extractor._save_results()
        
        # Progress reporting
        if i % max(1, args.progress_every) == 0 or i == len(plan):
            ok_n = sum(1 for r in batch_results.values() if r.get("success"))
            print(f"[{i}/{len(plan)}] ok={ok_n} fail={i - ok_n} elapsed={time.time()-t0:.1f}s")
        
        # Check for stop on fail
        if args.stop_on_fail and not res.get("success"):
            print(f"Stopping on first failure: {name}")
            break

    ok_n = sum(1 for r in batch_results.values() if r.get("success"))
    fail_n = len(batch_results) - ok_n
    summary = {
        "success": fail_n == 0,
        "processed": len(batch_results),
        "ok": ok_n,
        "fail": fail_n,
        "elapsed_sec": round(time.time() - t0, 3),
    }
    print(json.dumps(summary, indent=2))
    
    # Final validation summary (validation already done during extraction)
    if not args.no_validate:
        # Count blocks that were successfully validated and moved
        validated_blocks = sum(1 for r in batch_results.values() 
                              if r.get("success") and r.get("validation", {}).get("moved_to_block_dir"))
        
        print(f"\nValidation Summary:")
        print(f"Total blocks processed: {len(batch_results)}")
        print(f"Successfully extracted: {ok_n}")
        print(f"Successfully validated and moved: {validated_blocks}")
        print(f"Validation success rate: {(validated_blocks/ok_n*100):.1f}%" if ok_n > 0 else "Validation success rate: 0.0%")
        
        if args.cleanup_invalid:
            # Clean up invalid blocks that weren't moved
            cleanup_count = 0
            for name, result in batch_results.items():
                if result.get("success") and not result.get("validation", {}).get("moved_to_block_dir"):
                    invalid_file = Path("generated_packages") / f"{name}.py"
                    if invalid_file.exists():
                        try:
                            invalid_file.unlink()
                            cleanup_count += 1
                        except Exception as e:
                            print(f"Failed to cleanup {name}: {e}")
            
            if cleanup_count > 0:
                print(f"Cleaned up {cleanup_count} invalid blocks from generated_packages directory")


if __name__ == "__main__":
    main()
