from __future__ import annotations

from dataclasses import dataclass
from tilelang import tvm as tvm
from tvm.target import Target
from tvm.ir import Range
from tvm import tirx
from tilelang import language as T
from tilelang.utils.language import is_shared, is_fragment, is_tensor_memory
from tilelang.tileop.base import GemmWarpPolicy
from tilelang.language.dtypes import validate_gemm_ab_dtypes
from tvm.ir.base import Node
from tvm.ir import PrimExpr


@dataclass
class GemmBase:
    """Base class for GEMM tile operators.

    Classifies the GEMM variant by the memory scopes of operands A and B
    (SS, SR, RS, TS, RR) and provides common property accessors for the
    underlying ``gemm_node`` IR node.
    """

    gemm_node: Node

    def __post_init__(self) -> None:
        validate_gemm_ab_dtypes(
            self.A.dtype,
            self.B.dtype,
            a_in_tmem=is_tensor_memory(self.A),
            allow_f8f6f4_mixed=self.allow_f8f6f4_mixed_dtypes,
        )

    @property
    def allow_f8f6f4_mixed_dtypes(self) -> bool:
        # TODO(wt): Consider enabling mixed f8f6f4 operands for MMA paths too.
        return False

    def infer_layout(self, target: Target, thread_nums: int):
        raise NotImplementedError("infer_layout is not implemented")

    def lower(
        self,
        layout_map: dict,
        target: Target,
        thread_bounds: Range,
        thread_var: tirx.Var,
        mbar_phase_expr: tirx.PrimExpr | None = None,
    ):
        raise NotImplementedError("lower is not implemented")

    def is_gemm_ss(self) -> bool:
        """Return True if both A and B are in shared memory (SS variant)."""
        return is_shared(self.A) and is_shared(self.B)

    def is_gemm_sr(self) -> bool:
        """Return True if A is in shared memory and B is in registers (SR variant)."""
        return is_shared(self.A) and is_fragment(self.B)

    def is_gemm_rs(self) -> bool:
        """Return True if A is in registers and B is in shared memory (RS variant)."""
        return is_fragment(self.A) and is_shared(self.B)

    def is_gemm_ts(self) -> bool:
        """Return True if A is in tensor memory and B is in shared memory (TS variant)."""
        return is_tensor_memory(self.A) and is_shared(self.B)

    def is_gemm_rr(self) -> bool:
        """Return True if both A and B are in registers (RR variant)."""
        return is_fragment(self.A) and is_fragment(self.B)

    @property
    def M(self) -> int:
        return getattr(self.gemm_node, "m", None)

    @property
    def N(self) -> int:
        return getattr(self.gemm_node, "n", None)

    @property
    def K(self) -> int:
        return getattr(self.gemm_node, "k", None)

    @property
    def trans_A(self) -> bool:
        return getattr(self.gemm_node, "transA", None)

    @property
    def trans_B(self) -> bool:
        return getattr(self.gemm_node, "transB", None)

    @property
    def a_dtype(self):
        """A operand dtype."""
        return self.A.dtype

    @property
    def b_dtype(self):
        """B operand dtype."""
        return self.B.dtype

    @property
    def accum_dtype(self) -> str:
        return self.C.dtype

    @property
    def chunk(self) -> int:
        return self.K

    @property
    def A(self) -> tirx.Buffer:
        return getattr(self.gemm_node, "a", None)

    @property
    def B(self) -> tirx.Buffer:
        return getattr(self.gemm_node, "b", None)

    @property
    def C(self) -> tirx.Buffer:
        return getattr(self.gemm_node, "c", None)

    @property
    def ARegion(self):
        return getattr(self.gemm_node, "aRegion", None)

    @property
    def BRegion(self):
        return getattr(self.gemm_node, "bRegion", None)

    @property
    def CRegion(self):
        return getattr(self.gemm_node, "cRegion", None)

    @property
    def stride_A(self) -> int:
        return getattr(self.gemm_node, "strideA", None)

    @property
    def stride_B(self) -> int:
        return getattr(self.gemm_node, "strideB", None)

    @property
    def offset_A(self) -> int:
        return getattr(self.gemm_node, "offsetA", None)

    @property
    def offset_B(self) -> int:
        return getattr(self.gemm_node, "offsetB", None)

    @property
    def clear_accum(self) -> PrimExpr:
        return getattr(self.gemm_node, "clearAccum", None)

    @property
    def k_pack(self) -> int:
        return getattr(self.gemm_node, "kPack", None)

    @property
    def wg_wait(self) -> int:
        return getattr(self.gemm_node, "wgWait", 0)

    @property
    def is_tcgen05(self) -> bool:
        return getattr(self.gemm_node, "isTcgen05", False)

    @property
    def policy(self) -> GemmWarpPolicy:
        return getattr(self.gemm_node, "policy", None)

    @property
    def mbarptr(self) -> PrimExpr:
        return getattr(self.gemm_node, "mbarPtr", tvm.tirx.const(0, T.uint32))

    @property
    def mbar(self) -> tirx.BufferLoad | None:
        return getattr(self.gemm_node, "mbar", None)

    @property
    def C_coords(self):
        coords = getattr(self.gemm_node, "cCoords", None)
        if coords is None or len(coords) == 0:
            zero = tvm.tirx.const(0, T.int32)
            return [zero, zero]
        return [coords[i] for i in range(len(coords))]

    @property
    def SFARegion(self):
        return getattr(self.gemm_node, "sfaRegion", None)

    @property
    def SFBRegion(self):
        return getattr(self.gemm_node, "sfbRegion", None)

    @property
    def sf_k_start(self) -> PrimExpr:
        return getattr(self.gemm_node, "sfKStart", tvm.tirx.const(0, T.int32))

    @property
    def is_blockscaled(self) -> bool:
        return self.SFARegion is not None and self.SFBRegion is not None

    def get_region_base_offsets(self, region):
        """
        Get the base offset (start index) for each dimension from a BufferRegion.

        For example, if region is A_shared[ko % 2, 0:128, 0:64],
        this returns [ko % 2, 0, 0]

        Args:
            region: BufferRegion object

        Returns:
            List of PrimExpr representing the base offset for each dimension
        """
        if region is None:
            return []
        return [r.min for r in region.region]

    @property
    def A_base_offsets(self):
        """Get base offsets for each dimension of A region"""
        return self.get_region_base_offsets(self.ARegion)

    @property
    def B_base_offsets(self):
        """Get base offsets for each dimension of B region"""
        return self.get_region_base_offsets(self.BRegion)

    @property
    def C_base_offsets(self):
        """Get base offsets for each dimension of C region"""
        return self.get_region_base_offsets(self.CRegion)
