"""CPU-only structural checks for the unified kernel tree."""

from __future__ import annotations

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

import pytest

import sglang.kernels as kernels
from sglang.test.ci.ci_register import register_cpu_ci

register_cpu_ci(est_time=5, suite="base-a-test-cpu")

REPO_ROOT = Path(__file__).resolve().parents[3]
KERNELS_ROOT = REPO_ROOT / "python" / "sglang" / "kernels"
OPS_ROOT = KERNELS_ROOT / "ops"
JIT_CSRC_ROOT = KERNELS_ROOT / "jit" / "csrc"
AOT_ROOT = KERNELS_ROOT / "aot"


def _directory_names(root: Path) -> set[str]:
    return {
        path.name
        for path in root.iterdir()
        if path.is_dir()
        and not path.name.startswith((".", "__"))
        and any(path.rglob("*.py"))
    }


def _target_names(target: ast.expr) -> set[str]:
    if isinstance(target, ast.Name):
        return {target.id}
    if isinstance(target, (ast.List, ast.Tuple)):
        return {name for element in target.elts for name in _target_names(element)}
    return set()


def _bound_names(statements: list[ast.stmt]) -> set[str]:
    """Collect names a module can bind without importing it."""
    names: set[str] = set()
    for statement in statements:
        if isinstance(statement, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            names.add(statement.name)
        elif isinstance(statement, ast.Assign):
            for target in statement.targets:
                names.update(_target_names(target))
        elif isinstance(statement, (ast.AnnAssign, ast.AugAssign)):
            names.update(_target_names(statement.target))
        elif isinstance(statement, (ast.Import, ast.ImportFrom)):
            for alias in statement.names:
                names.add(alias.asname or alias.name.split(".", 1)[0])
        elif isinstance(statement, (ast.For, ast.AsyncFor)):
            names.update(_target_names(statement.target))
            names.update(_bound_names(statement.body))
            names.update(_bound_names(statement.orelse))
        elif isinstance(statement, ast.If):
            names.update(_bound_names(statement.body))
            names.update(_bound_names(statement.orelse))
        elif isinstance(statement, (ast.With, ast.AsyncWith)):
            names.update(_bound_names(statement.body))
        elif isinstance(statement, ast.Try):
            names.update(_bound_names(statement.body))
            names.update(_bound_names(statement.orelse))
            names.update(_bound_names(statement.finalbody))
            for handler in statement.handlers:
                names.update(_bound_names(handler.body))
        elif isinstance(statement, ast.Match):
            for case in statement.cases:
                names.update(_bound_names(case.body))
    return names


def _module_string_constants(tree: ast.Module) -> dict[str, str]:
    constants: dict[str, str] = {}
    for statement in tree.body:
        if not isinstance(statement, (ast.Assign, ast.AnnAssign)):
            continue
        value = statement.value
        if not isinstance(value, ast.Constant) or not isinstance(value.value, str):
            continue
        targets = (
            statement.targets
            if isinstance(statement, ast.Assign)
            else [statement.target]
        )
        for target in targets:
            for name in _target_names(target):
                constants[name] = value.value
    return constants


def _source_patterns(expression: ast.expr, constants: dict[str, str]) -> list[str]:
    if isinstance(expression, (ast.List, ast.Tuple)):
        return [
            pattern
            for element in expression.elts
            for pattern in _source_patterns(element, constants)
        ]
    if isinstance(expression, ast.Constant) and isinstance(expression.value, str):
        return [expression.value]
    if isinstance(expression, ast.Name) and expression.id in constants:
        return [constants[expression.id]]
    if isinstance(expression, ast.JoinedStr):
        parts = []
        for value in expression.values:
            if isinstance(value, ast.Constant):
                parts.append(str(value.value))
            elif isinstance(value, ast.FormattedValue):
                parts.append("*")
            else:
                raise AssertionError(f"Unsupported f-string segment: {ast.dump(value)}")
        return ["".join(parts)]
    raise AssertionError(
        f"Unsupported JIT source declaration: {ast.unparse(expression)}"
    )


def test_declared_operator_groups_match_packages():
    assert set(kernels.ops.__all__) == _directory_names(OPS_ROOT)


def test_registered_kernel_test_groups_are_known():
    declared_groups = set(kernels.ops.__all__)
    registered_root = REPO_ROOT / "test" / "registered" / "kernels"
    for kind in ("ops", "benchmark"):
        unknown = _directory_names(registered_root / kind) - declared_groups
        assert (
            not unknown
        ), f"Unknown {kind} kernel group directories: {sorted(unknown)}"


def test_internal_registry_target_attributes_are_declared():
    missing = []
    for spec in kernels.registry.all_specs():
        module_name, _, attribute_path = spec.target.partition(":")
        if not module_name.startswith("sglang.kernels."):
            continue
        module_spec = importlib.util.find_spec(module_name)
        if (
            module_spec is None
            or module_spec.origin is None
            or not module_spec.origin.endswith(".py")
        ):
            continue
        tree = ast.parse(Path(module_spec.origin).read_text())
        root_attribute = attribute_path.split(".", 1)[0]
        if root_attribute not in _bound_names(tree.body):
            missing.append(spec.target)
    assert not missing, f"KernelSpec targets missing attributes: {missing}"


def test_jit_source_declarations_exist():
    missing = []
    unsupported = []
    for python_file in OPS_ROOT.rglob("*.py"):
        tree = ast.parse(python_file.read_text())
        constants = _module_string_constants(tree)
        for call in (node for node in ast.walk(tree) if isinstance(node, ast.Call)):
            function_name = (
                call.func.id
                if isinstance(call.func, ast.Name)
                else call.func.attr if isinstance(call.func, ast.Attribute) else None
            )
            if function_name != "load_jit":
                continue
            for keyword in call.keywords:
                if keyword.arg not in {"cpp_files", "cuda_files"}:
                    continue
                try:
                    patterns = _source_patterns(keyword.value, constants)
                except AssertionError as exc:
                    unsupported.append(f"{python_file.relative_to(REPO_ROOT)}: {exc}")
                    continue
                for pattern in patterns:
                    matches = list(JIT_CSRC_ROOT.glob(pattern))
                    if not matches:
                        missing.append(
                            f"{python_file.relative_to(REPO_ROOT)} -> {pattern}"
                        )
    assert not unsupported, "Unsupported JIT source declarations:\n" + "\n".join(
        unsupported
    )
    assert not missing, "Missing JIT sources:\n" + "\n".join(missing)


def test_aot_compilation_units_are_accounted_for():
    manifests = [
        AOT_ROOT / "CMakeLists.txt",
        AOT_ROOT / "setup_metal.py",
        AOT_ROOT / "setup_musa.py",
        AOT_ROOT / "setup_rocm.py",
        AOT_ROOT / "csrc" / "cpu" / "CMakeLists.txt",
        *sorted((AOT_ROOT / "cmake").rglob("*.cmake")),
    ]
    manifest_text = "\n".join(path.read_text() for path in manifests)
    source_text = {
        path: path.read_text(errors="ignore")
        for path in (AOT_ROOT / "csrc").rglob("*")
        if path.is_file()
    }
    compilation_suffixes = {".cc", ".cpp", ".cu", ".hip", ".metal", ".mu"}
    missing = []
    for source in source_text:
        if source.suffix not in compilation_suffixes:
            continue
        if AOT_ROOT / "csrc" / "cpu" in source.parents:
            # The CPU build intentionally uses file(GLOB_RECURSE ... *.cpp).
            continue
        relative_path = source.relative_to(AOT_ROOT).as_posix()
        if relative_path in manifest_text:
            continue
        if any(
            source.name in text
            for other_source, text in source_text.items()
            if other_source != source
        ):
            # Some CUDA translation units are included by another source.
            continue
        missing.append(relative_path)
    assert not missing, f"AOT compilation units missing from build manifests: {missing}"


if __name__ == "__main__":
    sys.exit(pytest.main([__file__, "-v"]))
