#!/usr/bin/env python3
"""Verify that the packaged sources only relocate internal imports."""

from __future__ import annotations

import ast
import importlib.util
import subprocess
import sys
from pathlib import Path


PINNED_COMMIT = "8a26fb0ec9e353125ead798cb2e312d5ce48cded"
UPSTREAM_SUBDIR = Path("techniques/sparse_backends/sol_attn")
IGNORED_PARTS = {"__pycache__"}


def source_files(root: Path) -> dict[Path, Path]:
    return {
        path.relative_to(root): path
        for path in root.rglob("*")
        if path.is_file() and not (set(path.parts) & IGNORED_PARTS)
    }


def module_package(relative_path: Path) -> str:
    parent_parts = relative_path.with_suffix("").parts[:-1]
    return ".".join(("sol_attn", *parent_parts))


def normalized_imports(tree: ast.AST, package: str):
    imports = []
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            imports.extend(
                (alias.name, alias.asname or alias.name.split(".")[0])
                for alias in node.names
            )
        elif isinstance(node, ast.ImportFrom):
            module = node.module or ""
            if node.level:
                relative = "." * node.level + module
                module = importlib.util.resolve_name(relative, package)
            imports.extend(
                (
                    f"{module}.{alias.name}" if module else alias.name,
                    alias.asname or alias.name,
                )
                for alias in node.names
            )
    return sorted(imports)


def without_imports(text: str) -> str:
    tree = ast.parse(text)
    lines = text.splitlines(keepends=True)
    for node in ast.walk(tree):
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            for index in range(node.lineno - 1, node.end_lineno):
                lines[index] = "\n" if lines[index].endswith("\n") else ""
    return "".join(lines)


def main() -> int:
    if len(sys.argv) != 2:
        print("usage: verify_upstream.py /path/to/Sana", file=sys.stderr)
        return 2

    checkout = Path(sys.argv[1]).resolve()
    commit = subprocess.check_output(
        ["git", "-C", str(checkout), "rev-parse", "HEAD"],
        text=True,
    ).strip()
    if commit != PINNED_COMMIT:
        raise SystemExit(f"expected Sana {PINNED_COMMIT}, found {commit}")

    upstream_root = checkout / UPSTREAM_SUBDIR
    packaged_root = Path(__file__).resolve().parents[1] / "torch-ext/sol_attn"
    upstream = source_files(upstream_root)
    packaged = source_files(packaged_root)
    if upstream.keys() != packaged.keys():
        missing = sorted(upstream.keys() - packaged.keys())
        extra = sorted(packaged.keys() - upstream.keys())
        raise SystemExit(f"source inventory mismatch; missing={missing}, extra={extra}")

    relocated = []
    for relative_path in sorted(upstream):
        upstream_bytes = upstream[relative_path].read_bytes()
        packaged_bytes = packaged[relative_path].read_bytes()
        if upstream_bytes == packaged_bytes:
            continue
        if relative_path.suffix != ".py":
            raise SystemExit(f"non-Python source differs: {relative_path}")

        upstream_text = upstream_bytes.decode()
        packaged_text = packaged_bytes.decode()
        if without_imports(upstream_text) != without_imports(packaged_text):
            raise SystemExit(f"non-import source differs: {relative_path}")

        package = module_package(relative_path)
        upstream_imports = normalized_imports(ast.parse(upstream_text), package)
        packaged_imports = normalized_imports(ast.parse(packaged_text), package)
        if upstream_imports != packaged_imports:
            raise SystemExit(f"import semantics differ: {relative_path}")
        relocated.append(relative_path)

    print(
        "Verified pinned NVlabs/Sana source; only semantically equivalent "
        f"import relocation is present in {len(relocated)} files."
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())