"""Kernel launching language interface in TileLang."""

from __future__ import annotations
from collections import deque
import os
from tvm import tirx
from tvm.tirx import Var
from tvm.tirx.script.builder import evaluate as T_evaluate
from tvm.tirx.script.builder.frame import TIRFrame
from tvm.tirx.script.builder.frame import SBlockFrame
from tvm.ffi import register_object
from tilelang import _ffi_api
from tilelang.jit.exceptions import JITNoBuilderError
import threading

# Ensure single-dimension kernel bindings can be unpacked like iterables.
# especially for issue https://github.com/tile-ai/tilelang/issues/830
if not hasattr(Var, "__iter__"):

    def _var_iter(self):
        yield self

    Var.__iter__ = _var_iter  # type: ignore[attr-defined]

if not hasattr(Var, "__len__"):
    Var.__len__ = lambda self: 1  # type: ignore[attr-defined]


class FrameStack:
    """
    A simple stack-like wrapper around a deque that provides
    push, pop, and top methods for convenience.
    """

    def __init__(self):
        self._stack = deque()

    def push(self, item):
        """Pushes an item onto the top of the stack."""
        self._stack.append(item)

    def pop(self):
        """
        Pops and returns the top of the stack, or returns None
        if the stack is empty.
        """
        if self._stack:
            return self._stack.pop()
        raise IndexError(f"{self.__class__.__name__} is empty")

    def top(self):
        """
        Returns the item on the top of the stack without removing it,
        or None if the stack is empty.
        """
        if self._stack:
            return self._stack[-1]
        raise IndexError(f"{self.__class__.__name__} is empty")

    def size(self):
        """Returns the number of items in the stack."""
        return len(self._stack)

    def __len__(self):
        """Returns the number of items in the stack."""
        return len(self._stack)

    def __bool__(self):
        """
        Allows truthy checks on the stack object itself,
        e.g., 'if stack: ...'
        """
        return bool(self._stack)


# Use thread local to store the stack
# This is to avoid the cross-thread interference
_local = threading.local()


def _get_current_stack() -> FrameStack:
    if not hasattr(_local, "kernel_launch_frame_stack"):
        _local.kernel_launch_frame_stack = FrameStack()
    return _local.kernel_launch_frame_stack


def _normalize_bindings(bindings: list[Var]) -> Var | list[Var]:
    """
    Return a bare Var when we only have a single binding so that users may write either
    `with T.Kernel(...) as pid:` or `with T.Kernel(...) as (pid,)`.
    Otherwise, keep the list semantics for multi-dimensional launches.
    """
    if len(bindings) == 1:
        return bindings[0]
    return bindings


def _normalize_threads(
    threads: int | list[int] | tuple | None,
) -> list[int]:
    """Normalize a thread-block specification into a 3-D extent list.

    Args:
        threads: A thread count, a per-dimension extent list/tuple, or None for the default.

    Returns:
        The extents as ``[x, y, z]``, padding missing dimensions with 1.

    Raises:
        ValueError: If ``threads`` has an unsupported type, or any concrete extent
            is not positive.
    """
    if threads is None:
        threads = 128  # default thread number

    if isinstance(threads, int):
        normalized = [threads, 1, 1]
    elif isinstance(threads, list):
        normalized = threads + [1] * (3 - len(threads))
    elif isinstance(threads, tuple):
        normalized = list(threads) + [1] * (3 - len(threads))
    else:
        raise ValueError("threads must be an integer or a list of integers")

    # A block extent must be positive. A non-positive value would otherwise reach
    # codegen and launch a kernel with an empty thread block, silently writing nothing.
    if any(isinstance(extent, int) and extent <= 0 for extent in normalized):
        raise ValueError(f"threads must be positive, got {threads}")

    return normalized


def _normalize_cluster_dims(
    cluster_dims: int | tuple[int, int, int] | list[int] | None,
) -> list[int] | None:
    if cluster_dims is None:
        return None

    if isinstance(cluster_dims, (list, tuple)):
        cluster_dims = list(cluster_dims) + [1] * (3 - len(cluster_dims))
    elif isinstance(cluster_dims, int):
        cluster_dims = [cluster_dims, 1, 1]
    else:
        raise ValueError("cluster_dims must be a list or tuple of integers")

    return None if cluster_dims == [1, 1, 1] else cluster_dims


@register_object("tl.KernelLaunchFrame")
class KernelLaunchFrame(TIRFrame):
    """
    KernelLaunchFrame is a custom TIRFrame that manages block/thread indices
    and handles the entry and exit of the kernel launch scope.
    """

    def __enter__(self) -> Var | list[Var]:
        """
        Enters the KernelLaunchFrame scope and pushes this frame onto the stack.
        Returns one Var if we detect exactly 5 frames (meaning there is a single
        block dimension), or a list of Vars otherwise.
        """
        super().__enter__()
        _get_current_stack().push(self)

        last_block_frame = self.frames[-1]
        assert isinstance(last_block_frame, SBlockFrame), f"Last frame must be a block frame, got {last_block_frame}"

        # Return a list of grid loop vars (excluding the last 4 frames:
        # threadIdx.x, threadIdx.y, threadIdx.z and the block frame with attributes).
        return _normalize_bindings([frame.vars[0] for frame in self.frames[0:-4]])

    def __exit__(self, ptype, value, trace):
        """
        Exits the KernelLaunchFrame scope and pops this frame from the stack,
        but only if it's indeed the topmost frame.
        """
        stack = _get_current_stack()
        if stack.top() is self:
            stack.pop()
        super().__exit__(ptype, value, trace)

    @classmethod
    def Current(cls) -> KernelLaunchFrame | None:
        """
        Returns the topmost (current) KernelLaunchFrame from the stack if it exists,
        or None if the stack is empty.
        """
        stack = _get_current_stack()
        return stack.top() if stack else None

    def get_block_extent(self, dim: int) -> int:
        """
        Returns the block extent for the given dimension.
        dim=0 corresponds to blockIdx.x, dim=1 to blockIdx.y, and dim=2 to blockIdx.z.
        """
        iter_var = self.frames[dim].doms[0]
        return int(iter_var.extent)

    def get_block_extents(self) -> list[int]:
        """
        Returns the block extents for all three dimensions.
        """
        return [self.get_block_extent(dim) for dim in range(3)]

    def get_thread_extent(self, dim: int) -> int:
        """
        Returns the thread extent for the given dimension.
        dim=0 corresponds to threadIdx.x, dim=1 to threadIdx.y, and dim=2 to threadIdx.z.
        """
        iter_var = self.frames[-4 + dim].doms[0]
        return int(iter_var.extent)

    def get_thread_extents(self) -> list[int]:
        """
        Returns the thread extents for all three dimensions.
        """
        return [self.get_thread_extent(dim) for dim in range(3)]

    def get_thread_binding(self, dim: int = 0) -> Var:
        """
        Returns the thread binding for the given dimension.
        dim=0 corresponds to threadIdx.x, dim=1 to threadIdx.y, and dim=2 to threadIdx.z.
        """
        return self.frames[-4 + dim].vars[0]

    def get_thread_bindings(self) -> list[Var]:
        """
        Returns the thread binding for the given dimension.
        dim=0 corresponds to threadIdx.x, dim=1 to threadIdx.y, and dim=2 to threadIdx.z.
        """
        return [frame.vars[0] for frame in self.frames[-4:-1]]

    def get_num_threads(self) -> int:
        """
        Returns the thread indices from the topmost frame.
        """
        num_threads: int = 1
        for thread_dim in range(3):
            num_threads *= self.get_thread_extent(thread_dim)
        return num_threads

    def get_block_binding(self, dim: int = 0) -> Var:
        """
        Returns the block binding for the given dimension.
        dim=0 corresponds to blockIdx.x, dim=1 to blockIdx.y, and dim=2 to blockIdx.z.
        """
        return self.frames[dim].vars[0]

    def get_block_bindings(self) -> list[Var]:
        """
        Returns all three block bindings.
        """
        return [frame.vars[0] for frame in self.frames[0:-4]]

    @property
    def blocks(self) -> list[Var]:
        """
        Returns the block indices from the topmost frame.
        """
        return [frame.vars[0] for frame in self.frames[0:-4]]

    @property
    def threads(self) -> list[Var]:
        """
        Returns the thread indices from the topmost frame.
        """
        return [frame.vars[0] for frame in self.frames[-4:-1]]

    @property
    def num_threads(self) -> int:
        """
        Returns the total number of threads.
        """
        return self.get_num_threads()


def Kernel(
    *blocks: int | tirx.PrimExpr,
    threads: int | list[int] | tuple | None = None,
    prelude: str | None = None,
):
    """Tools to quickly construct a kernel launch frame.

    The launch nest is emitted in a target-neutral form (thread_binding
    For loops); each backend pipeline materializes it via
    MaterializeKernelLaunch. Backends without SIMT (e.g. CPU) simply
    ignore the thread extents at compile time, so the same kernel can be
    compiled for any target.

    Parameters
    ----------
    blocks : int
        A list of extent, can be 1-3 dimension, representing gridDim.(x|y|z)
    threads : int
        A integer representing blockDim.x
        Or a list of integers representing blockDim.(x|y|z)
        if the value is -1, we skip the threadIdx.x binding.
    prelude : str
        The import c code of the kernel,
        will be injected before the generated kernel code.

    Returns
    -------
    res : Tuple[frame.LaunchThreadFrame]
        The result LaunchThreadFrame.

    Examples
    --------
    Create a 1-D CUDA kernel launch and unpack the single block index:

    .. code-block:: python

        with T.Kernel(T.ceildiv(N, 128), threads=128) as bx:
            # bx is the blockIdx.x binding (also iterable as (bx,))
            ...

    Launch a 2-D grid while requesting two thread dimensions:

    .. code-block:: python

        with T.Kernel(grid_x, grid_y, threads=(64, 2)) as (bx, by):
            tx, ty = T.get_thread_bindings()
            ...
    """
    # In eager mode, we construct AST directly without prim_func,
    # so there must be a Builder available. If not, this function
    # is being called outside of a JIT/prim_func context.
    # lazy import to avoid circular import
    from tilelang.language.eager.builder import Builder

    if Builder.current() is None:
        raise JITNoBuilderError("T.Kernel() can only be used inside @tilelang.jit or @T.prim_func context. No Builder is available.")

    attrs: dict = {}
    threads = _normalize_threads(threads)

    if prelude is not None:
        attrs["pragma_import_c"] = prelude

    return _ffi_api.KernelLaunch(blocks, threads, attrs)


def ClusterKernel(
    *blocks: int | tirx.PrimExpr,
    cluster_dims: int | tuple[int, int, int] | list[int],
    threads: int | list[int] | tuple | None = None,
    prelude: str | None = None,
):
    """Construct a kernel launch frame with a CUDA thread block cluster
    (SM90+ only).

    This is the CUDA-specific variant of :func:`Kernel`: identical launch
    semantics and bindings, plus a ``cluster_dims`` annotation. The kernel
    will be launched with cudaLaunchKernelEx using
    cudaLaunchAttributeClusterDimension.

    Parameters
    ----------
    blocks : int
        A list of extent, can be 1-3 dimension, representing gridDim.(x|y|z)
    cluster_dims : int | tuple[int, int, int] | list[int]
        The cluster dimensions. For example, use 2 or (2, 1, 1) to create
        2-CTA clusters.
    threads : int
        A integer representing blockDim.x
        Or a list of integers representing blockDim.(x|y|z)
    prelude : str
        The import c code of the kernel,
        will be injected before the generated kernel code.

    Examples
    --------
    .. code-block:: python

        with T.ClusterKernel(grid_x, grid_y, cluster_dims=2, threads=128) as (bx, by):
            ...
    """
    from tilelang.language.eager.builder import Builder

    if Builder.current() is None:
        raise JITNoBuilderError("T.ClusterKernel() can only be used inside @tilelang.jit or @T.prim_func context. No Builder is available.")

    attrs: dict = {}
    threads = _normalize_threads(threads)

    if prelude is not None:
        attrs["pragma_import_c"] = prelude

    cluster_dims = _normalize_cluster_dims(cluster_dims)
    if cluster_dims is not None:
        attrs["cluster_dims"] = cluster_dims

    return _ffi_api.KernelLaunch(blocks, threads, attrs)


# For CUDA source kernels, we need to load the source code from a file or string.


def _load_cuda_source(source_code_or_path: str | os.PathLike[str]) -> str:
    source = os.fspath(source_code_or_path)
    if not isinstance(source, str) or not source.strip():
        raise ValueError("source_code_or_path must be a non-empty source string or source path")

    expanded = os.path.expanduser(source)
    if os.path.isfile(expanded):
        with open(expanded, encoding="utf-8") as f:
            return f.read()

    source_markers = ("\n", "__global__", 'extern "C"', "#include")
    if any(marker in source for marker in source_markers):
        return source

    contains_path_sep = os.path.sep in source or (os.path.altsep is not None and os.path.altsep in source)
    if contains_path_sep or source.endswith((".cu", ".cuh", ".cuda", ".cpp", ".cc", ".c")):
        raise FileNotFoundError(f"CUDA source file not found: {source}")

    return source


def CUDASourceCodeKernel(
    *blocks: int | tirx.PrimExpr,
    threads: int | list[int] | tuple | None = None,
    source_code_or_path: str | os.PathLike[str],
    entry_name: str = "main_kernel",
    cluster_dims: int | tuple[int, int, int] | list[int] | None = None,
    prelude: str | None = None,
) -> None:
    """Launch a kernel from CUDA source code or a CUDA source file.

    The code must follows the following rules:
    1. The kernel source must be a valid CUDA kernel which can be correctly compiled under TileLang's context.
    2. The kernel source must either contains only one `__global__` function as an entry, or have a `__global__` entry function named `main_kernel`.

    Parameters
    ----------
    source_code_or_path : str | os.PathLike[str]
        Inline CUDA source code, or a path to a CUDA source file.
        If the argument resolves to an existing file, the file contents are
        loaded. Otherwise it is treated as inline CUDA source code.
    blocks : int
        A list of extent, can be 1-3 dimension, representing gridDim.(x|y|z)
    entry_name : str | None
        Optional name of the `__global__` CUDA entry function inside the
        provided source. When specified, TileLang launches that external CUDA
        entry directly.
    threads : int
        A integer representing blockDim.x
        Or a list of integers representing blockDim.(x|y|z)
        if the value is -1, we skip the threadIdx.x binding.
    cluster_dims : int | tuple[int, int, int] | list[int] | None
        The cluster dimensions for SM90+ cluster launch.
        For example, use 2 or (2, 1, 1) to create 2-CTA clusters.
        When specified, the kernel will be launched using cudaLaunchKernelEx
        with cudaLaunchAttributeClusterDimension.
    prelude : str
        The import c code of the kernel,
        will be injected before the generated kernel code.
    """
    from tilelang.language.eager.builder import Builder

    if Builder.current() is None:
        raise JITNoBuilderError(
            "T.CUDASourceCodeKernel() can only be used inside @tilelang.jit or @T.prim_func context. No Builder is available."
        )

    source = _load_cuda_source(source_code_or_path)
    if prelude is not None:
        source = prelude + "\n" + source

    attrs: dict = {"code_block_source": source}
    if not isinstance(entry_name, str) or not entry_name.strip():
        raise ValueError("entry_name must be a non-empty string when provided")
    attrs["code_block_entry_name"] = entry_name

    threads = _normalize_threads(threads)

    cluster_dims = _normalize_cluster_dims(cluster_dims)
    if cluster_dims is not None:
        attrs["cluster_dims"] = cluster_dims

    with _ffi_api.KernelLaunch(blocks, threads, attrs):
        # Keep the launch frame alive until SplitHostDevice can lift the
        # external CUDA source pragma onto the device PrimFunc.
        T_evaluate(tirx.call_extern("int32", entry_name))


def get_thread_binding(dim: int = 0) -> Var:
    """Returns the thread binding for the given dimension."""
    assert KernelLaunchFrame.Current() is not None, "KernelLaunchFrame is not initialized"
    return KernelLaunchFrame.Current().get_thread_binding(dim)


def get_thread_bindings() -> list[Var]:
    """Returns all three thread bindings."""
    assert KernelLaunchFrame.Current() is not None, "KernelLaunchFrame is not initialized"
    return KernelLaunchFrame.Current().get_thread_bindings()


def get_block_binding(dim: int = 0) -> Var:
    """Returns the block binding for the given dimension."""
    assert KernelLaunchFrame.Current() is not None, "KernelLaunchFrame is not initialized"
    return KernelLaunchFrame.Current().get_block_binding(dim)


def get_block_bindings() -> list[Var]:
    """Returns all three block bindings."""
    assert KernelLaunchFrame.Current() is not None, "KernelLaunchFrame is not initialized"
    return KernelLaunchFrame.Current().get_block_bindings()


def get_thread_extent(dim: int = 0) -> int:
    """Returns the thread extent for the given dimension."""
    assert KernelLaunchFrame.Current() is not None, "KernelLaunchFrame is not initialized"
    return KernelLaunchFrame.Current().get_thread_extent(dim)


def get_thread_extents() -> list[int]:
    """Returns all three thread extents."""
    assert KernelLaunchFrame.Current() is not None, "KernelLaunchFrame is not initialized"
    return KernelLaunchFrame.Current().get_thread_extents()


def get_block_extent(dim: int = 0) -> int:
    """Returns the block extent for the given dimension."""
    assert KernelLaunchFrame.Current() is not None, "KernelLaunchFrame is not initialized"
    return KernelLaunchFrame.Current().get_block_extent(dim)


def get_block_extents() -> list[int]:
    """Returns all three block extents."""
    assert KernelLaunchFrame.Current() is not None, "KernelLaunchFrame is not initialized"
    return KernelLaunchFrame.Current().get_block_extents()
