"""The dependency policy, enforced.

Every module under align/ and scripts/ is AST-scanned. Any import whose top-level
package is not (a) stdlib, (b) the four allowed runtime deps, or (c) the package
itself, fails this test. Lazy imports inside functions count too -- a lazy
`import transformers` is still a dependency, just a hidden one.

Whitelist is per-FILE and per-PACKAGE, so an optional extra can exist in exactly
one place and nowhere else.
"""
from __future__ import annotations

import ast
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
ALLOWED = {"torch", "triton", "numpy", "safetensors", "align"}
# file (relative to ROOT) -> extra packages it alone may import lazily
FILE_WHITELIST: dict[str, set[str]] = {
    "align/data/ingest_parquet.py": {"pyarrow"},
    "scripts/teachers/download_teacher.py": {"huggingface_hub"},  # offline weight download, never imported at runtime
}
STDLIB = set(sys.stdlib_module_names)


def _imports(path: Path):
    tree = ast.parse(path.read_text(), filename=str(path))
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for a in node.names:
                yield node.lineno, a.name.split(".")[0]
        elif isinstance(node, ast.ImportFrom):
            if node.level:  # relative import -> the package itself
                continue
            if node.module:
                yield node.lineno, node.module.split(".")[0]


def test_no_forbidden_imports():
    bad = []
    for sub in ("align", "scripts"):
        for py in sorted((ROOT / sub).rglob("*.py")):
            rel = str(py.relative_to(ROOT))
            extra = FILE_WHITELIST.get(rel, set())
            for lineno, top in _imports(py):
                if top in STDLIB or top in ALLOWED or top in extra:
                    continue
                bad.append(f"{rel}:{lineno}: {top}")
    assert not bad, "forbidden third-party imports:\n" + "\n".join(bad)


def test_extras_are_not_reachable_from_core():
    """extras/ may import anything, but nothing under align/ may import extras/."""
    for py in (ROOT / "align").rglob("*.py"):
        for _, top in _imports(py):
            assert top != "extras", f"{py} imports extras/"