from __future__ import annotations
from typing import List, Tuple, TYPE_CHECKING
from dataclasses import dataclass

import triton.experimental.gluon.language._core as ttgl
from triton.experimental.gluon.language._layouts import PaddedSharedLayout, SwizzledSharedLayout
from triton.experimental.gluon.language.amd.cdna5 import PartitionedSharedLayout
from triton.experimental.gluon.language._core import builtin, _unwrap_if_constexpr

if TYPE_CHECKING:
    from triton._C import ir
    from triton.experimental.gluon.language._core import shared_memory_descriptor

__all__ = [
    "update_tensor_descriptor", "async_load", "async_load_fused", "async_wait", "make_tensor_descriptor",
    "tensor_descriptor", "tensor_descriptor_type", "prefetch", "async_store", "async_scatter", "async_gather"
]


@dataclass(eq=True)
class tensor_descriptor_type(ttgl.base_type):
    """The type for a tensor descriptor."""

    block_type: ttgl.block_type
    shape_type: ttgl.tuple_type
    strides_type: ttgl.tuple_type
    layout: PaddedSharedLayout | SwizzledSharedLayout | PartitionedSharedLayout

    def __str__(self) -> str:
        return f"tensor_descriptor<{self.block_type}, {self.layout}>"

    def _unflatten_ir(self, handles: List[ir.value], cursor: int) -> Tuple[tensor_descriptor, int]:
        handle = handles[cursor]
        cursor += 1
        shape, cursor = self.shape_type._unflatten_ir(handles, cursor)
        strides, cursor = self.strides_type._unflatten_ir(handles, cursor)
        value = tensor_descriptor(handle, shape, strides, self)
        return value, cursor

    def _to_ir(self, builder: ir.builder) -> ir.type:
        is_signed = self.block_type.element_ty.is_int_signed()
        return builder.get_tensor_descriptor_layout_type(
            self.block_type.to_ir(builder),
            is_signed,
            self.layout._to_ir(builder),
        )

    def _flatten_ir_types(self, builder: ir.builder, out: List[ir.type]) -> None:
        out.append(self._to_ir(builder))
        self.shape_type._flatten_ir_types(builder, out)
        self.strides_type._flatten_ir_types(builder, out)

    def mangle(self) -> str:
        return f"TD{self.block_type.mangle()}_{self.shape_type.mangle()}_{self.strides_type.mangle()}_{self.layout.mangle()}TD"


@dataclass
class tensor_descriptor(ttgl.base_value):
    """A descriptor representing a tensor in global memory."""

    handle: ir.value
    shape: ttgl.tuple
    strides: ttgl.tuple
    type: tensor_descriptor_type

    def _set_name(self, builder: ir.builder, name: str) -> None:
        self.handle.set_loc(builder.create_name_loc(name, self.handle.get_loc()))
        self.shape._set_name(builder, name + ".shape")
        self.strides._set_name(builder, name + ".stride")

    def _flatten_ir(self, handles: List[ir.value]) -> None:
        handles.append(self.handle)
        self.shape._flatten_ir(handles)
        self.strides._flatten_ir(handles)

    @property
    def block_type(self):
        return self.type.block_type

    @property
    def block_shape(self):
        return self.type.block_type.shape

    @property
    def dtype(self):
        return self.type.block_type.element_ty

    @property
    def layout(self):
        return self.type.layout


@builtin
def make_tensor_descriptor(base: ttgl.tensor, shape: List[ttgl.constexpr | ttgl.tensor],
                           strides: List[ttgl.constexpr | ttgl.tensor], block_shape: List[ttgl.constexpr],
                           layout: PaddedSharedLayout | SwizzledSharedLayout | PartitionedSharedLayout,
                           _semantic=None) -> tensor_descriptor:
    """Make a tensor descriptor object.

    AMD CDNA5 TDM tensor descriptors only support zero padding.

    Args:
        base (tensor): base pointer of the tensor in global memory.
        shape (List[int]): shape of the tensor.
        strides (List[int]): strides of the tensor.
        block_shape (List[int]): block shape of the tensor.
        layout (PaddedSharedLayout | SwizzledSharedLayout | PartitionedSharedLayout): the layout of the tensor in shared memory.

    Returns:
        tensor_descriptor: the created tensor descriptor object
    """
    ndim = len(shape)
    assert 1 <= ndim <= 5, f"Expected 1 <= ndim <= 5 but got {ndim} dimensions"
    assert len(strides) == ndim, f"Expected {ndim} strides but got {len(strides)}"
    block_shape = _unwrap_if_constexpr(block_shape)
    assert len(block_shape) == ndim, f"Expected block_shape to have {ndim} dimensions but got {len(block_shape)}"
    assert isinstance(base.dtype, ttgl.pointer_type), "Expected base to be a pointer"

    layout = _unwrap_if_constexpr(layout)
    assert isinstance(layout, (PaddedSharedLayout, SwizzledSharedLayout, PartitionedSharedLayout)), \
        "Expected layout to be a PaddedSharedLayout, SwizzledSharedLayout, or PartitionedSharedLayout"
    if isinstance(layout, SwizzledSharedLayout):
        assert layout.max_phase == 1, "Expected max_phase to be 1 for SwizzledSharedLayout"
    elif isinstance(layout, PartitionedSharedLayout):
        assert isinstance(layout.partition_layout, (PaddedSharedLayout, SwizzledSharedLayout)), \
            "PartitionedSharedLayout partition_layout must be PaddedSharedLayout or SwizzledSharedLayout"
        if isinstance(layout.partition_layout, SwizzledSharedLayout):
            assert layout.partition_layout.max_phase == 1, \
                "Expected max_phase to be 1 for SwizzledSharedLayout in PartitionedSharedLayout"

    base_handle = base.handle
    shape_handles = _semantic._convert_to_ir_values(shape, require_i64=False)  # i32 shape
    stride_handles = _semantic._convert_to_ir_values(strides, require_i64=True)  # i64 stride

    shape = ttgl.tuple(shape)
    strides = ttgl.tuple(strides)
    block_type = ttgl.block_type(base.type.element_ty, block_shape)
    type = tensor_descriptor_type(block_type, shape.type, strides.type, layout)

    padding = _semantic._str_to_padding_option("zero")
    handle = _semantic.builder.create_make_tensor_descriptor(type._to_ir(_semantic.builder), base_handle, shape_handles,
                                                             stride_handles, padding)

    return tensor_descriptor(handle, shape, strides, type)


def _handle_i32_pred(pred, _semantic):
    pred = _unwrap_if_constexpr(pred)
    if isinstance(pred, bool):
        pred = int(pred)
    pred = _semantic.to_tensor(pred)
    if pred.type.is_int1():
        pred = _semantic.cast(pred, ttgl.int32)
    assert pred.type.is_int32(), f"Expected pred to be an int32 or int1 value, but got {pred.type}"
    return pred


@builtin
def update_tensor_descriptor(desc: tensor_descriptor, add_offsets: List[ttgl.constexpr | ttgl.tensor] = None,
                             set_bounds: List[ttgl.constexpr | ttgl.tensor] = None, pred=None,
                             clamp_bounds: bool = False, _semantic=None) -> tensor_descriptor:
    """Update selected fields of a TDM descriptor; return a new descriptor SSA value.

    Each parameter is independently optional; only the fields the caller
    names are written.  Everything else is inherited from the input
    descriptor.

    NOTE: Unlike the standard tensor-descriptor mental model, ``add_offsets``
    here moves the tile position only.  It does NOT update the descriptor's
    bounds.  If the loop crosses an OOB boundary, you must also pass
    ``set_bounds`` explicitly, or pass ``clamp_bounds=True`` to derive the OOB
    extent from ``add_offsets``.

    Args:
        desc (tensor_descriptor): the input descriptor.
        add_offsets (List[int], optional): per-dim deltas in element units
            that move the tile position.  Does not touch the bounds.
        set_bounds (List[int], optional): per-dim absolute rewrite of the
            descriptor's bounds.  Use to install OOB extent at a peel
            epilogue.
        pred (int, optional): set the descriptor's predicate.
        clamp_bounds (bool, optional): if True, also shrink the descriptor's
            bounds by ``add_offsets`` (``tensor_dim[i] = max(0, tensor_dim[i] - add_offsets[i])``)
            to derive the advanced tile's OOB extent. Requires ``add_offsets``;
            mutually exclusive with ``set_bounds``.

    Returns:
        tensor_descriptor: a new descriptor SSA value with the requested
        fields rewritten.

    Raises:
        ValueError: if no parameter is provided (no-op updates are forbidden).

    **Example**

    .. code-block:: python

        # K-loop interior: bump tile position only
        d = tdm.update_tensor_descriptor(d, add_offsets=[0, BLOCK_K])

        # Prologue: position at first tile + set the predicate
        d = tdm.update_tensor_descriptor(d, add_offsets=[pid_m * BLOCK_M, 0],
                                         pred=do_load)

        # Peel epilogue: install real OOB extent for the partial last tile
        d = tdm.update_tensor_descriptor(d, set_bounds=[M - pid_m * BLOCK_M, K - k_main])
    """
    if add_offsets is None and set_bounds is None and pred is None:
        raise ValueError("tdm.update_tensor_descriptor requires at least one of add_offsets, "
                         "set_bounds, pred")

    clamp_bounds = _unwrap_if_constexpr(clamp_bounds)
    if clamp_bounds:
        if add_offsets is None:
            raise ValueError("tdm.update_tensor_descriptor: clamp_bounds requires add_offsets")
        if set_bounds is not None:
            raise ValueError("tdm.update_tensor_descriptor: clamp_bounds and set_bounds are mutually exclusive")

    rank = len(desc.block_shape)

    add_offset_handles = []
    if add_offsets is not None:
        if len(add_offsets) != rank:
            raise ValueError(f"add_offsets must have length {rank} (descriptor rank), got {len(add_offsets)}")
        add_offset_handles = _semantic._convert_to_ir_values(add_offsets, require_i64=False)

    set_bounds_handles = []
    if set_bounds is not None:
        if len(set_bounds) != rank:
            raise ValueError(f"set_bounds must have length {rank} (descriptor rank), got {len(set_bounds)}")
        set_bounds_handles = _semantic._convert_to_ir_values(set_bounds, require_i64=False)

    pred_handle = ttgl.ir.value()
    if pred is not None:
        pred_handle = _handle_i32_pred(pred, _semantic).handle

    new_handle = _semantic.builder.create_update_tensor_descriptor(desc.handle, add_offset_handles, set_bounds_handles,
                                                                   pred_handle, bool(clamp_bounds))
    # Rebuild shape/strides tuples so the returned descriptor's tuple objects
    # don't alias the original's (the frontend's SSA tracking assumes distinct
    # tuples per descriptor value).
    new_shape = ttgl.tuple(list(desc.shape))
    new_strides = ttgl.tuple(list(desc.strides))
    return tensor_descriptor(new_handle, new_shape, new_strides, desc.type)


@builtin
def async_load(src: tensor_descriptor, offsets: List[ttgl.constexpr | ttgl.tensor] = None,
               dest: shared_memory_descriptor = None, pred=None, mbarrier: shared_memory_descriptor = None,
               warp_used_hint=None, cache_modifier="", _semantic=None) -> None:
    """Load a block of tensor specified in tensor descriptor from global memory to shared memory asynchronously.

    This operation expects a prior :func:`update_tensor_descriptor` to position the
    descriptor for the load offsets and bounds.  Doing this explicitly lets the
    developer control when and where the descriptor update happens, which has
    performance implications for downstream code generation.  For convenience,
    ``offsets`` (and optionally ``pred``) may be passed here, in which case an
    :func:`update_tensor_descriptor` is emitted immediately before the load.

    Args:
        src (tensor_descriptor): the source tensor descriptor.
        offsets (List[int], optional): if given, the offsets from the base
            pointer used to position the descriptor before the load.
        dest (shared_memory_descriptor): the shared memory destination to store the loaded data.
        pred (bool, optional): if given, predicate to enable or disable the load.
        mbarrier (shared_memory_descriptor, optional): The barrier object to signal "arrive" on.
        warp_used_hint (int, optional): Bitmask selecting the active warp
            subset for descriptor layout (bit ``n`` => warp ``n``).  Doesn't
            affect the data in ``dest``, only the work split.
            The number of active warps must be a power of two, and the active
            warps must follow a regular bit pattern for efficient lowering.
            Examples: ``0b00001111`` (warps 0..3), ``0b11110000`` (warps
            4..7), ``0b01010101`` (warps 0,2,4,6).  Omit / ``None`` = all
            warps participate; explicit ``0`` and other invalid hints are
            rejected by the verifier.
        cache_modifier (str, optional): Cache behavior.
    """
    assert dest is not None, "async_load requires a dest shared_memory_descriptor"
    if offsets:
        eff_pred = True if pred is None else pred
        src = update_tensor_descriptor(src, add_offsets=offsets, pred=eff_pred, clamp_bounds=True, _semantic=_semantic)
    elif pred is not None:
        src = update_tensor_descriptor(src, pred=pred, _semantic=_semantic)

    mbarrier = _unwrap_if_constexpr(mbarrier)
    mbarrier_handle = mbarrier.handle if mbarrier is not None else ttgl.ir.value()
    cache_modifier = _semantic._str_to_load_cache_modifier(cache_modifier)

    warp_used_hint = _unwrap_if_constexpr(warp_used_hint)
    if warp_used_hint is not None:
        warp_used_hint = int(warp_used_hint)

    _semantic.builder.create_async_tdm_copy_global_to_local(src.handle, dest.handle, mbarrier_handle, cache_modifier,
                                                            warp_used_hint)


@builtin
def async_load_fused(members: List[Tuple[tensor_descriptor, shared_memory_descriptor, ttgl.constexpr | int]],
                     cache_modifier="", _semantic=None) -> None:
    """Emit one explicit fused TDM load for 2-4 descriptor/destination pairs.

    This can perform better than several consecutive separate TDM loads,
    especially for more than two loads. Under the hood, different warps load
    different descriptor/destination pairs according to each member's
    ``warp_used_hint``; collectively all participating warps load all pairs at
    the block level.

    Each member is ``(desc, dest, warp_used_hint)``. The descriptors must
    already encode their tile offsets, predicates, and bounds; use
    :func:`update_tensor_descriptor` before calling this helper when needed.
    All members share one cache modifier, matching the fused IR operation.

    Args:
        members: 2-4 ``(desc, dest, warp_used_hint)`` tuples. Hints must be
            legal, pairwise-disjoint bitmasks.
        cache_modifier (str, optional): Cache behavior shared by all members.
    """
    members = _unwrap_if_constexpr(members)
    if not 2 <= len(members) <= 4:
        raise ValueError(f"tdm.async_load_fused requires 2 to 4 members, got {len(members)}")

    desc_handles = []
    dest_handles = []
    warp_used_hints = []
    rank = None
    for idx, member in enumerate(members):
        member = _unwrap_if_constexpr(member)
        if len(member) != 3:
            raise ValueError("tdm.async_load_fused members must be (desc, dest, warp_used_hint) tuples")
        desc, dest, warp_used_hint = member
        if not isinstance(desc, tensor_descriptor):
            raise TypeError(f"tdm.async_load_fused member {idx}: expected tensor_descriptor")
        if rank is None:
            rank = len(desc.block_shape)
        if len(desc.block_shape) != rank:
            raise ValueError("tdm.async_load_fused requires all descriptors to have the same rank")

        warp_used_hint = _unwrap_if_constexpr(warp_used_hint)
        if warp_used_hint is None:
            raise ValueError(f"tdm.async_load_fused member {idx}: warp_used_hint is required")

        desc_handles.append(desc.handle)
        dest_handles.append(dest.handle)
        warp_used_hints.append(int(warp_used_hint))

    cache_modifier = _semantic._str_to_load_cache_modifier(cache_modifier)
    _semantic.builder.create_async_tdm_fused_copy_global_to_local(desc_handles, dest_handles, warp_used_hints,
                                                                  cache_modifier)


@builtin
def async_store(dest: tensor_descriptor, offsets: List[ttgl.constexpr | ttgl.tensor] = None,
                src: shared_memory_descriptor = None, mbarrier: shared_memory_descriptor = None, cache_modifier="",
                _semantic=None) -> None:
    """Store a block of tensor specified in tensor descriptor from shared memory to global memory asynchronously.

    See :func:`async_load` for how the descriptor is positioned (a prior
    :func:`update_tensor_descriptor` or the ``offsets`` convenience). Stores take no
    ``pred``.

    Args:
        dest (tensor_descriptor): the destination tensor descriptor.
        offsets (List[int], optional): if given, the offsets from the base
            pointer used to position the descriptor before the store.
        src (shared_memory_descriptor): the shared memory source to load the data.
        mbarrier (shared_memory_descriptor, optional): The barrier object to signal "arrive" on.
    """
    assert src is not None, "async_store requires a src shared_memory_descriptor"
    if offsets:
        dest = update_tensor_descriptor(dest, add_offsets=offsets, clamp_bounds=True, _semantic=_semantic)

    mbarrier = _unwrap_if_constexpr(mbarrier)
    mbarrier_handle = mbarrier.handle if mbarrier is not None else ttgl.ir.value()
    cache_modifier = _semantic._str_to_store_cache_modifier(cache_modifier)
    _semantic.builder.create_async_tdm_copy_local_to_global(dest.handle, src.handle, mbarrier_handle, cache_modifier)


@builtin
def async_wait(num_outstanding=0, _semantic=None) -> None:
    """Wait for the outstanding asynchronous tensor operations to complete.

    Args:
        num_outstanding (int): number of outstanding async tensor operations to wait for.
    """
    num_outstanding = _unwrap_if_constexpr(num_outstanding)
    _semantic.builder.create_async_tdm_wait(num_outstanding)


@builtin
def async_scatter(desc: tensor_descriptor, dst_row_indices: ttgl.tensor, src: shared_memory_descriptor,
                  mbarrier: shared_memory_descriptor = None, _semantic=None) -> None:
    """Scatter data from shared memory to non-contiguous rows in global memory asynchronously.

    This operation uses TDM scatter mode to write data to non-contiguous rows in global memory.
    Unlike :func:`async_store`, which writes to contiguous rows, scatter allows writing to
    arbitrary rows specified by the ``dst_row_indices`` tensor.

    The dtype of ``dst_row_indices`` determines the index size:

    * ``int16``: up to 16 rows can be scattered per TDM instruction.
    * ``int32``: up to 8 rows can be scattered per TDM instruction.

    If more rows are needed, multiple TDM instructions will be automatically issued.

    The column offset is carried by the descriptor: position it beforehand with
    :func:`update_tensor_descriptor` with ``add_offsets=[0, col]``.

    Args:
        desc (tensor_descriptor): The destination tensor descriptor. Must be 2D and positioned
            through :func:`update_tensor_descriptor` for the column.
        dst_row_indices (tensor): 1D tensor of row indices (int16 or int32) in the destination tensor.
        src (shared_memory_descriptor): the shared memory source containing data to scatter. Must be 2D.
        mbarrier (shared_memory_descriptor, optional): The barrier object to signal "arrive" on.
    """
    ndim = len(desc.block_shape)
    assert ndim == 2, f"TDM scatter only supports 2D tensors, got {ndim}D"

    src_ndim = len(src.shape)
    assert src_ndim == 2, f"TDM scatter src must be 2D, got {src_ndim}D"

    mbarrier = _unwrap_if_constexpr(mbarrier)
    mbarrier_handle = mbarrier.handle if mbarrier is not None else ttgl.ir.value()

    _semantic.builder.create_async_tdm_scatter(desc.handle, dst_row_indices.handle, src.handle, mbarrier_handle)


@builtin
def async_gather(desc: tensor_descriptor, src_row_indices: ttgl.tensor, dst: shared_memory_descriptor,
                 mbarrier: shared_memory_descriptor = None, _semantic=None) -> None:
    """Gather data from non-contiguous rows in global memory to shared memory asynchronously.

    This operation uses TDM gather mode to read data from non-contiguous rows in global memory.
    Unlike :func:`async_load`, which reads from contiguous rows, gather allows reading from
    arbitrary rows specified by the ``src_row_indices`` tensor.

    The dtype of ``src_row_indices`` determines the index size:

    * ``int16``: up to 16 rows can be gathered per TDM instruction.
    * ``int32``: up to 8 rows can be gathered per TDM instruction.

    If more rows are needed, multiple TDM instructions will be automatically issued.

    The column offset and predicate are carried by the descriptor: position it
    beforehand with :func:`update_tensor_descriptor`, passing
    ``add_offsets=[0, col]`` and ``pred=p``.

    Args:
        desc (tensor_descriptor): The source tensor descriptor. Must be 2D and positioned
            through :func:`update_tensor_descriptor` for the column and predicate.
        src_row_indices (tensor): 1D tensor of row indices (int16 or int32) in the source tensor.
        dst (shared_memory_descriptor): the shared memory destination to store gathered data. Must be 2D.
        mbarrier (shared_memory_descriptor, optional): The barrier object to signal "arrive" on.
    """
    ndim = len(desc.block_shape)
    assert ndim == 2, f"TDM gather only supports 2D tensors, got {ndim}D"

    dst_ndim = len(dst.shape)
    assert dst_ndim == 2, f"TDM gather dst must be 2D, got {dst_ndim}D"

    mbarrier = _unwrap_if_constexpr(mbarrier)
    mbarrier_handle = mbarrier.handle if mbarrier is not None else ttgl.ir.value()

    _semantic.builder.create_async_tdm_gather(desc.handle, src_row_indices.handle, dst.handle, mbarrier_handle)


@builtin
def prefetch(src: tensor_descriptor, offsets: List[ttgl.constexpr | ttgl.tensor], pred: bool = True,
             speculative: bool = False, _semantic=None) -> None:
    """Prefetches a block of tensor specified in tensor descriptor from global memory into L2.

    Speculative prefetches can generate more efficient assembly because they do not require out of bounds checks.
    However, they are dropped by the hardware if their virtual address translation is not cached.
    So speculative should only be set if previous iterations have accessed the same virtual page (e.g. column major).

    Args:
        src (tensor_descriptor): the source tensor descriptor.
        offsets (List[int]): the offsets from the base pointer in the tensor descriptor.
        pred (bool, optional): Predicate to enable or disable the prefetch. Defaults to True.
        speculative (bool, optional): Whether the prefetch is speculative. Defaults to False.
    """
    offset_handles = _semantic._convert_to_ir_values(offsets, require_i64=False)
    pred = _semantic.to_tensor(pred)
    pred_handle = pred.handle
    speculative = _unwrap_if_constexpr(speculative)
    _semantic.builder.create_tdm_prefetch(src.handle, offset_handles, pred_handle, speculative, False)


@builtin
def _test_prefetch_with_offsets(src: tensor_descriptor, offsets: List[ttgl.constexpr | ttgl.tensor], pred: bool = True,
                                speculative: bool = False, _semantic=None) -> ttgl.tensor:
    """Test-only prefetch variant that returns offsets for validation."""
    offset_handles = _semantic._convert_to_ir_values(offsets, require_i64=False)
    pred = _semantic.to_tensor(pred)
    pred_handle = pred.handle
    speculative = _unwrap_if_constexpr(speculative)
    handle = _semantic.builder.create_tdm_prefetch(src.handle, offset_handles, pred_handle, speculative, True)
    shape = _semantic.builder.get_shape_from_tensor(handle)
    layout = _semantic.builder.get_gluon_layout_from_tensor(handle)
    ret_ty = ttgl.distributed_type(ttgl.int64, shape, layout)
    tensor = ttgl.tensor(handle, ret_ty)
    return tensor
