"""Check that the tile_kernels patch is behaviour-preserving.

The lazy `__init__` and the optional TileLang import must change nothing for a host that has
TileLang. What callers can do with the package -- attribute access, `from x import y`, wildcard
imports inside the kernel modules, the config knobs -- has to keep working, or the patch has
broken the normal path in order to enable an unusual one.

Run with TileLang available.
"""

import sys
import traceback
from collections.abc import Callable

FAILED = []


def check(name: str, fn: Callable[[], None]) -> None:
    try:
        fn()
        print(f"  ok   {name}")
    except Exception as exc:
        print(f"  FAIL {name}: {type(exc).__name__}: {str(exc).splitlines()[0][:110]}")
        traceback.print_exc()
        FAILED.append(name)


def main() -> int:
    print("=== same surface as the eager __init__")

    def submodules_are_attributes() -> None:
        import tile_kernels

        for name in ("config", "utils", "engram", "mhc", "modeling", "moe", "quant", "rand", "transform", "testing"):
            module = getattr(tile_kernels, name)
            assert module is not None, name
            assert module.__name__ == f"tile_kernels.{name}", module.__name__

    def submodule_imports_work() -> None:
        # `from pkg import sub` goes through __getattr__ as well.
        from tile_kernels import quant  # noqa: F401

        assert quant.per_token_cast is not None

    def config_reexports() -> None:
        import tile_kernels

        for name in ("get_num_sms", "set_num_sms", "get_token_alignment", "set_token_alignment"):
            assert callable(getattr(tile_kernels, name)), name
        # `is_ascend` was never a top-level export, and deepinfer reaches it as
        # `tile_kernels.config.is_ascend`; check the path that is actually used.
        assert callable(tile_kernels.config.is_ascend)
        # The knobs the engine sets must still take effect.
        tile_kernels.set_token_alignment(64)
        assert tile_kernels.get_token_alignment() == 64
        tile_kernels.set_token_alignment(128)
        assert tile_kernels.get_token_alignment() == 128

    def star_import_of_common() -> None:
        # Kernel modules do `from tile_kernels.quant.common import *`; the guarded macro block
        # still has to publish the macros at module level.
        namespace: dict = {}
        exec("from tile_kernels.quant.common import *", namespace)  # noqa: S102
        for name in ("get_sf_and_inv", "get_sf_index", "load_sf", "transform_sf", "store_sf", "BaseCastConfig"):
            assert name in namespace, name

    def dir_is_populated() -> None:
        import tile_kernels

        names = dir(tile_kernels)
        assert "quant" in names and "get_num_sms" in names, "dir() lost the eager surface"

    def unknown_attribute_raises() -> None:
        import tile_kernels

        try:
            _ = tile_kernels.definitely_not_exported
        except AttributeError:
            return
        raise AssertionError("missing attribute did not raise")

    for name, fn in [
        ("submodules are attributes", submodules_are_attributes),
        ("from pkg import sub", submodule_imports_work),
        ("config re-exports and knobs", config_reexports),
        ("from quant.common import *", star_import_of_common),
        ("dir() surface", dir_is_populated),
        ("unknown attribute raises", unknown_attribute_raises),
    ]:
        check(name, fn)

    print()
    if FAILED:
        print(f"BROKEN {len(FAILED)}: {FAILED}")
        return 1
    print("PATCH_SAFE")
    return 0


if __name__ == "__main__":
    sys.exit(main())