# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from __future__ import annotations

import functools
from collections.abc import Callable
from functools import cache
from typing import TYPE_CHECKING, Any

from vllm.platforms import current_platform
from vllm.utils.import_utils import has_tilelang

if TYPE_CHECKING or current_platform.is_cuda():
    if not has_tilelang():
        raise ImportError(
            "tilelang is required for mhc but is not installed. Install it with "
            "`pip install tilelang`."
        )
    import tilelang
    import tilelang.language as T
else:
    tilelang = None  # type: ignore[assignment]
    T = None  # type: ignore[assignment]


def _ensure_tilelang_imported() -> None:
    """Bind the `tilelang` and `T` module globals, importing them if needed.

    On ROCm, this runs on the first kernel call instead of at import time.

    Raises:
        ImportError: If TileLang is not installed.

    """
    global T, tilelang

    if tilelang is not None:
        return
    if not has_tilelang():
        raise ImportError(
            "tilelang is required for mhc but is not installed. Install it with "
            "`pip install tilelang`."
        )
    import tilelang as tilelang_module
    import tilelang.language as tilelang_language

    tilelang = tilelang_module
    T = tilelang_language


@cache
def _get_pass_configs() -> dict[Any, Any]:
    _ensure_tilelang_imported()
    pass_configs: dict[Any, Any] = {
        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
    }
    if current_platform.is_cuda():
        pass_configs[tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL] = 10
    return pass_configs


class _DeferredTileLangJitKernel:
    """Stand-in for a `tilelang.jit` kernel that decorates on first use.

    Both attribute access and calling apply the decoration. Required for
    compile-only JIT warmup on platforms that defer import of tilelang.
    """

    _kernel_function: Callable[..., Any] | None = None
    _jit_kernel: Any = None

    def __init__(self, kernel_function: Callable[..., Any]) -> None:
        self._kernel_function = kernel_function
        functools.update_wrapper(self, kernel_function)

    def _ensure_jit_kernel(self) -> Any:
        if self._jit_kernel is None:
            _ensure_tilelang_imported()
            kernel_function = self._kernel_function
            assert kernel_function is not None
            kernel_function.__globals__["tilelang"] = tilelang
            kernel_function.__globals__["T"] = T
            self._jit_kernel = tilelang.jit(pass_configs=_get_pass_configs())(
                kernel_function
            )
        return self._jit_kernel

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        return self._ensure_jit_kernel()(*args, **kwargs)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._ensure_jit_kernel(), name)


def tilelang_jit(kernel_function: Callable[..., Any]) -> Callable[..., Any]:
    """Apply `tilelang.jit`, deferring until first use on ROCm.

    ROCm defers JIT decoration so importing the caller's module does not
    require TileLang immediately. CUDA keeps the eager decoration behavior.

    The kernel body parsed by TileLang references `T` as an unqualified
    global, so on the deferred ROCm path this rebinds `T`/`tilelang` in the
    decorated function's own module globals once they become available.
    """
    if not current_platform.is_rocm():
        _ensure_tilelang_imported()
        return tilelang.jit(pass_configs=_get_pass_configs())(kernel_function)

    return _DeferredTileLangJitKernel(kernel_function)