#!/usr/bin/env python
"""
Static CI gate: audit top-level imports in the main-thread import graph
reachable from sloppy.py. Fails (exit 1) if any heavy module is imported
at the top of a main-thread-reachable file.

The Main Thread Purity Invariant (see conductor/tracks/startup_speedup_20260606/
spec.md:2.1) requires that the main thread's import chain contains only:
  - Python stdlib modules
  - The lean gui_2 skeleton: imgui_bundle, defer, src.imgui_scopes,
    src.theme_2 (default theme only), src.theme_models, src.paths,
    src.models, src.events
  - Modules that have been refactored to be lean (e.g., src.ai_client
    after Phase 3)

Function-level imports inside method bodies are NOT audited (they run
on whichever thread calls the function, and the warmup mechanism in
spec.md:2.2 Layer 3 makes that safe).

Usage:
  uv run python scripts/audit_main_thread_imports.py [--root <path>] [--entry <file>]

Defaults: --root=. --entry=sloppy.py
"""

import argparse
import ast
import sys
from dataclasses import dataclass
from pathlib import Path


STDLIB = set(getattr(sys, "stdlib_module_names", set()) or set())
LEAN_ALLOWLIST: set[str] = {
 "imgui_bundle",
 "defer",
 "defer.sugar",
 "src.imgui_scopes",
 "src.theme_2",
 "src.theme_models",
 "src.paths",
 "src.models",
 "src.events",
 "src.config",
}


@dataclass(frozen=True)
class Violation:
 file: Path
 lineno: int
 module: str
 statement: str

 def render(self) -> str:
  return f"  {self.file}:L{self.lineno}  {self.module:<40}  {self.statement[:80]}"


def _top_module(import_name: str) -> str:
 return import_name.split(".")[0]


def _collect_top_level_imports(path: Path) -> list[tuple[int, str, str]]:
 try:
  source = path.read_text(encoding="utf-8", errors="replace")
 except OSError:
  return []
 try:
  tree = ast.parse(source, filename=str(path))
 except SyntaxError:
  return []
 results: list[tuple[int, str, str]] = []
 for node in tree.body:
  results.extend(_walk_imports(node))
 return results


def _walk_imports(node: ast.AST) -> list[tuple[int, str, str]]:
 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
  return []
 if isinstance(node, ast.Import):
  stmt = ast.unparse(node).strip()
  return [(node.lineno, alias.name, stmt) for alias in node.names]
 if isinstance(node, ast.ImportFrom):
  if node.level and node.level > 0:
   return []
  if not node.module:
   return []
  stmt = ast.unparse(node).strip()
  return [(node.lineno, node.module, stmt)]
 results: list[tuple[int, str, str]] = []
 for child in ast.iter_child_nodes(node):
  results.extend(_walk_imports(child))
 return results


def _resolve_local(import_name: str, root: Path) -> Path | None:
 parts = import_name.split(".")
 base = root.joinpath(*parts[:-1]) if len(parts) > 1 else root
 candidate_py = base / f"{parts[-1]}.py"
 if candidate_py.is_file():
  return candidate_py
 candidate_pkg = base / parts[-1] / "__init__.py"
 if candidate_pkg.is_file():
  return candidate_pkg
 return None


def _walk_import_graph(entry: Path, root: Path) -> list[Path]:
 visited: set[Path] = set()
 queue: list[Path] = [entry.resolve()]
 while queue:
  current = queue.pop(0)
  if current in visited:
   continue
  visited.add(current)
  for _lineno, name, _stmt in _collect_top_level_imports(current):
   resolved = _resolve_local(name, root)
   if resolved is not None:
    queue.append(resolved)
 return sorted(visited)


def _is_allowed(module: str) -> bool:
 if module in STDLIB:
  return True
 if module in LEAN_ALLOWLIST:
  return True
 top = _top_module(module)
 if top in STDLIB or top in LEAN_ALLOWLIST:
  return True
 return False


def audit(root: Path, entry: Path) -> list[Violation]:
 entry = entry.resolve()
 root = root.resolve()
 if not entry.is_file():
  raise FileNotFoundError(f"entry not found: {entry}")
 graph = _walk_import_graph(entry, root)
 violations: list[Violation] = []
 for path in graph:
  for lineno, name, stmt in _collect_top_level_imports(path):
   if _is_allowed(name):
    continue
   violations.append(Violation(
    file=path.relative_to(root),
    lineno=lineno,
    module=name,
    statement=stmt,
   ))
 return violations


def main(argv: list[str]) -> int:
 ap = argparse.ArgumentParser(description="Audit main-thread import graph for heavy modules")
 ap.add_argument("--root", default=".", help="project root (default: cwd)")
 ap.add_argument("--entry", default="sloppy.py", help="entry point file (default: sloppy.py)")
 ap.add_argument("--verbose", action="store_true", help="print the import graph + each file's imports")
 args = ap.parse_args(argv[1:])

 root = Path(args.root).resolve()
 entry = (root / args.entry).resolve()
 try:
  graph = _walk_import_graph(entry, root)
 except FileNotFoundError as e:
  print(f"error: {e}", file=sys.stderr)
  return 2

 if args.verbose:
  print(f"# import graph from {entry.relative_to(root)} ({len(graph)} files reachable)")
  for path in graph:
   rel = path.relative_to(root)
   imports = _collect_top_level_imports(path)
   if not imports:
    continue
   print(f"\n## {rel}")
   for lineno, name, stmt in imports:
    mark = "OK " if _is_allowed(name) else "BAD"
    print(f"  [{mark}] L{lineno:>4}  {name:<40}  {stmt[:60]}")

 try:
  violations = audit(root, entry)
 except FileNotFoundError as e:
  print(f"error: {e}", file=sys.stderr)
  return 2

 if not violations:
  print(f"OK: {len(graph)} files in main-thread import graph; no heavy top-level imports.")
  return 0

 print(f"FAIL: {len(violations)} heavy top-level import(s) in main-thread import graph:")
 for v in violations:
  print(v.render())
 return 1


if __name__ == "__main__":
 raise SystemExit(main(sys.argv))