from __future__ import annotations

from tilelang.tileop.gemm.gemm_base import GemmBase
from tilelang.utils.language import is_shared, is_full_region, is_fragment, is_metal_simdgroup
from tilelang import tvm as tvm
from tvm.target import Target
from tvm.ir import Range
from tvm import tirx as tir
from tilelang import language as T
from tilelang.transform.simplify import _Simplify


GEMM_INST_METAL = "metal.simdgroup"


class GemmMetal(GemmBase):
    def is_gemm_ss(self) -> bool:
        return is_shared(self.A) and is_shared(self.B)

    def infer_layout(self, target: Target, thread_nums: int):
        return {}

    def lower(
        self, layout_map: dict, target: Target, thread_bounds: Range, thread_var: tir.Var, mbar_phase_expr: tir.PrimExpr | None = None
    ):
        thread_nums = thread_bounds.extent
        m_warp, n_warp = self.policy.compute_warp_partition(self.M, self.N, thread_nums, target, GEMM_INST_METAL)
        warp_row_tiles = int(self.M // m_warp)
        warp_col_tiles = int(self.N // n_warp)

        from tilelang.metal.intrinsics.metal_macro_generator import MPSIntrinEmitter

        mps_emitter = MPSIntrinEmitter(
            a_dtype=self.a_dtype,
            b_dtype=self.b_dtype,
            accum_dtype=self.accum_dtype,
            a_transposed=self.trans_A,
            b_transposed=self.trans_B,
            block_row_warps=m_warp,
            block_col_warps=n_warp,
            warp_row_tiles=warp_row_tiles,
            warp_col_tiles=warp_col_tiles,
            chunk=self.chunk,
            thread_var=thread_var,
        )

        a_dtype = self.a_dtype
        b_dtype = self.b_dtype
        accum_dtype = self.accum_dtype
        warp_rows = mps_emitter.warp_rows
        warp_cols = mps_emitter.warp_cols
        num_simd_c = warp_rows * warp_cols
        block_K = mps_emitter.chunk
        micro_size_k = mps_emitter.micro_size_k

        A_region = self.ARegion
        B_region = self.BRegion
        C_region = self.CRegion

        C_buf = C_region.buffer

        clear_accum = self.clear_accum
        c_in_register = is_fragment(C_buf) or is_metal_simdgroup(C_buf)

        assert block_K >= micro_size_k, f"block_K ({block_K}) must be >= micro_size_k ({micro_size_k})"
        assert is_full_region(C_region), "Fragment output C must be a full region"
        assert c_in_register or is_shared(C_buf), (
            f"Metal GEMM requires C in local.fragment, metal.simdgroup or shared scope, got {C_buf.scope()}"
        )

        if self.is_gemm_ss():
            if c_in_register:

                @T.prim_func
                def _gemm_ss_simdgroup() -> None:
                    A_local = T.alloc_local((warp_rows * 64), a_dtype, scope="metal.simdgroup")
                    B_local = T.alloc_local((warp_cols * 64), b_dtype, scope="metal.simdgroup")
                    if clear_accum:
                        for _i in T.serial(num_simd_c):
                            T.make_filled_simdgroup_matrix(C_buf.data, _i, T.cast(0, accum_dtype))
                    for ki in T.serial(0, (block_K // micro_size_k)):
                        mps_emitter.ldmatrix_a(A_local, A_region, ki)
                        mps_emitter.ldmatrix_b(B_local, B_region, ki)
                        mps_emitter.mma(A_local, B_local, C_buf)

                return _Simplify(_gemm_ss_simdgroup, inline_let=True)
            else:

                @T.prim_func
                def _gemm_ss_shared() -> None:
                    A_local = T.alloc_local((warp_rows * 64), a_dtype, scope="metal.simdgroup")
                    B_local = T.alloc_local((warp_cols * 64), b_dtype, scope="metal.simdgroup")
                    C_simd = T.alloc_local((num_simd_c * 64), accum_dtype, scope="metal.simdgroup")
                    if clear_accum:
                        for _i in T.serial(num_simd_c):
                            T.make_filled_simdgroup_matrix(C_simd.data, _i, T.cast(0, accum_dtype))
                    else:
                        mps_emitter.simd_load(C_simd, C_buf)
                    for ki in T.serial(0, (block_K // micro_size_k)):
                        mps_emitter.ldmatrix_a(A_local, A_region, ki)
                        mps_emitter.ldmatrix_b(B_local, B_region, ki)
                        mps_emitter.mma(A_local, B_local, C_simd)

                    mps_emitter.simd_store(C_simd, C_buf)

                return _Simplify(_gemm_ss_shared, inline_let=True)
        else:
            raise ValueError(f"Unsupported gemm combination, A: {self.A.scope()}, B: {self.B.scope()}")
