"""Let the reference implementations sync, without turning the guard off everywhere.

The engine wraps the forward pass in `forbid_cuda_sync()`, which sets torch's sync debug mode to
"error"; one blocking copy or one `nonzero()` raises. The compiled kernels never sync, so the guard
is a real property of the production path and worth keeping.

The torch references cannot honour it, and not because of an oversight that could be patched away:

* `get_mapping_from_psum` reads `psum_num_tokens_per_expert` to the host to compute per-expert
  alignment ranges in Python.
* `moe_topk_gate_forward` compacts valid rows with `active.nonzero()` and branches on
  `normal_mask.any()`, both of which need the result on the host.

Making those sync-free means rewriting how the reference does its addressing. That is the wrong
trade: the reference exists to make the model runnable and checkable, and it is not the path whose
performance matters.

So each compat module wraps its reference module once:

    _ref = reference_module(tile_kernels.torch)

and every call through it runs inside `allow_cuda_sync()`. The wrapper is at the module boundary
rather than in each branch, so no call site knows about it and a newly used reference entry point
is covered automatically.

Deliberately *not* done: `torch.cuda.set_sync_debug_mode("default")` for the whole process. That
would also stop the guard from catching accidental syncs in the compiled path, which is exactly
what it is for.
"""

from __future__ import annotations

import functools
from collections.abc import Callable
from typing import Any

try:  # The guard lives in deepinfer; a bare install of this package would not have it.
    from deepinfer.utils import allow_cuda_sync

    _HAVE_GUARD = True
except Exception:  # pragma: no cover
    _HAVE_GUARD = False


def _wrap(fn: Callable[..., Any]) -> Callable[..., Any]:
    if not _HAVE_GUARD:
        return fn

    @functools.wraps(fn)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        with allow_cuda_sync():
            return fn(*args, **kwargs)

    return wrapper


class _ReferenceProxy:
    """Attribute access that grants sync permission for every callable it returns."""

    __slots__ = ("_module",)

    def __init__(self, module: Any) -> None:
        object.__setattr__(self, "_module", module)

    def __getattr__(self, name: str) -> Any:
        value = getattr(object.__getattribute__(self, "_module"), name)
        return _wrap(value) if callable(value) else value

    def __repr__(self) -> str:
        return f"reference_module({object.__getattribute__(self, '_module')!r})"


def reference_module(module: Any) -> Any:
    """Wrap a `tile_kernels.torch` submodule so its calls may sync."""
    return _ReferenceProxy(module)


if _HAVE_GUARD:
    allow_sync = allow_cuda_sync
else:  # pragma: no cover - no guard installed, nothing to lift
    import contextlib

    allow_sync = contextlib.nullcontext