"""List every `tile_kernels.<module>.<name>` deepinfer uses, so the compat layer can be complete.

Worth doing as a scan rather than discovering gaps one traceback at a time: the e2e run is the
only thing that would catch a missing attribute, and it costs a minute per attempt.
"""

import ast
import pathlib
import sys

ROOT = pathlib.Path(__file__).resolve().parents[2] / "deepinfer"
SKIP = ("tilelang_ops/cuda", "tilelang_ops/asc", "tk_compat")

# Names reached through a module-level `from deepinfer import tk_compat as tile_kernels`.
used: dict[str, dict[str, set[str]]] = {}
# Names imported directly from a submodule: `from deepinfer.tk_compat.quant import per_token_cast`.
imported: dict[str, set[str]] = {}

for path in sorted(ROOT.rglob("*.py")):
    rel = path.as_posix()
    if any(s in rel for s in SKIP):
        continue
    try:
        tree = ast.parse(path.read_text(encoding="utf-8"))
    except (SyntaxError, UnicodeDecodeError):
        continue
    for node in ast.walk(tree):
        if isinstance(node, ast.ImportFrom) and node.module and node.module.startswith("deepinfer.tk_compat"):
            parts = node.module.split(".")
            if len(parts) >= 3:
                sub = parts[2]
                for a in node.names:
                    imported.setdefault(sub, set()).add(a.name)
        if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Attribute):
            inner = node.value
            if isinstance(inner.value, ast.Name) and inner.value.id == "tile_kernels":
                used.setdefault(inner.attr, {}).setdefault(node.attr, set()).add(
                    rel.replace(str(ROOT.parent) + "/", "")
                )

print("=== 通过 `tile_kernels.<mod>.<name>` 访问的")
for sub in sorted(used):
    print(f"\n{ sub }:")
    for name in sorted(used[sub]):
        files = ", ".join(sorted(used[sub][name]))
        print(f"  {name:<38} {files}")

print("\n=== `from deepinfer.tk_compat.<mod> import ...` 的")
for sub in sorted(imported):
    print(f"\n{sub}: {', '.join(sorted(imported[sub]))}")

# Cross-check against what the compat modules actually export.
sys.path.insert(0, str(ROOT.parent))
import importlib

print("\n=== 对照 compat 实际导出")
missing = []
for sub, names in sorted(used.items()):
    try:
        module = importlib.import_module(f"deepinfer.tk_compat.{sub}")
    except Exception as exc:  # noqa: BLE001
        print(f"  {sub}: 无法导入 ({exc})")
        continue
    absent = [n for n in names if not hasattr(module, n)]
    if absent:
        missing.append((sub, absent))
        print(f"  {sub}: MISSING {absent}")
    else:
        print(f"  {sub}: ok")

if missing:
    print(f"\n缺口 {sum(len(a) for _, a in missing)} 个")
    sys.exit(1)
print("\n无缺口")