# ruff: noqa
from tilelang import tvm as tvm
import tilelang as tl
from tilelang.backend.target import determine_target
import tilelang.language as T
import tilelang.testing
from tilelang.cuda.pipeline import CUDAPassPipelineBodyPrologue
from tvm import tirx

auto_target = tvm.target.Target(determine_target("auto"))


def _apply(func):
    mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
    mod = tvm.tirx.transform.BindTarget(auto_target)(mod)
    mod = tl.transform.MaterializeKernelLaunch()(mod)
    mod = tl.cuda.transform.LowerSharedBarrier()(mod)
    return mod


def _collect_calls(stmt, op_name: str | set[str]):
    op_names = {op_name} if isinstance(op_name, str) else op_name
    calls = []

    def visitor(node):
        if isinstance(node, tvm.tirx.Call) and hasattr(node, "op") and hasattr(node.op, "name") and str(node.op.name) in op_names:
            calls.append(node)

    tvm.tirx.stmt_functor.post_order_visit(stmt, visitor)
    return calls


def _collect_storage_syncs(stmt):
    return _collect_calls(stmt, "tirx.tvm_storage_sync")


def _collect_init_barrier_calls(stmt):
    return _collect_calls(stmt, "tirx.ptx_init_barrier_thread_count")


def _collect_fence_barrier_init(stmt):
    return _collect_calls(stmt, "tl.ptx_fence_barrier_init")


def _collect_shuffle_elect(stmt):
    return _collect_calls(stmt, "tl.tl_shuffle_elect")


def _collect_barrier_blocks(stmt):
    blocks = []

    def visitor(node):
        if isinstance(node, tvm.tirx.SBlock):
            barrier_bufs = [buf for buf in node.alloc_buffers if buf.scope() in ("shared.barrier", "shared.cluster_barrier")]
            if barrier_bufs:
                blocks.append(node)

    tvm.tirx.stmt_functor.post_order_visit(stmt, visitor)
    return blocks


def test_single_barrier():
    """Single barrier with one arrive count."""

    @T.prim_func
    def func():
        with T.Kernel(1, threads=128):
            mbar = T.alloc_barrier(128)  # noqa: F841

    mod = _apply(func)
    body = mod["main"].body

    assert len(_collect_init_barrier_calls(body)) == 1
    assert len(_collect_fence_barrier_init(body)) == 1
    assert len(_collect_shuffle_elect(body)) == 1

    init_call = _collect_init_barrier_calls(body)[0]
    # arrive count should be 128
    assert init_call.args[1].value == 128


def test_multiple_barriers():
    """Multiple barriers with different arrive counts."""

    @T.prim_func
    def func():
        with T.Kernel(1, threads=128):
            mbars = T.alloc_barrier([1, 1, 128, 128])  # noqa: F841

    mod = _apply(func)
    body = mod["main"].body

    init_calls = _collect_init_barrier_calls(body)
    assert len(init_calls) == 4

    arrive_counts = sorted([c.args[1].value for c in init_calls])
    assert arrive_counts == [1, 1, 128, 128]

    # Should have exactly one shuffle_elect guard and one fence_barrier_init
    assert len(_collect_shuffle_elect(body)) == 1
    assert len(_collect_fence_barrier_init(body)) == 1

    # Should sync after init
    syncs = _collect_storage_syncs(body)
    assert len(syncs) >= 1


def test_no_barrier_is_noop():
    """Pass should be a no-op when no barrier buffers are present."""

    @T.prim_func
    def func():
        with T.Kernel(1, threads=128):
            buf = T.alloc_shared((16,), T.float16)
            buf[0] = T.float16(0)

    mod = _apply(func)
    body = mod["main"].body

    assert len(_collect_init_barrier_calls(body)) == 0
    assert len(_collect_fence_barrier_init(body)) == 0


def test_plan_update_keeps_barrier_init_with_tcgen05_no_tma():
    """Regression for tcgen05 no-TMA kernels after pass reordering."""

    pass_configs = {tl.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True}

    @T.prim_func
    def func(
        X: T.Tensor((256, 256), T.float16),
        Y: T.Tensor((256, 256), T.float16),
    ):
        with T.Kernel(2, 2, threads=256) as (bx, by):
            A_shared = T.alloc_shared((128, 128), T.float16)
            B_shared = T.alloc_shared((128, 128), T.float16)
            C_tmem = T.alloc_tmem([128, 128], T.float32)
            mbar = T.alloc_barrier(1)
            C_local = T.alloc_fragment((128, 128), T.float32)
            Y_shared = T.alloc_shared((128, 128), T.float16)

            for ko in T.Pipelined(2, num_stages=2):
                T.copy(X[by * 128, ko * 128], A_shared)
                T.copy(X[bx * 128, ko * 128], B_shared)
                T.tcgen05_gemm(
                    A_shared,
                    B_shared,
                    C_tmem,
                    transpose_B=True,
                    mbar=mbar,
                    clear_accum=ko == 0,
                )
                T.mbarrier_wait_parity(mbar, ko % 2)

            T.copy(C_tmem, C_local)
            T.copy(C_local, Y_shared)
            T.copy(Y_shared, Y[by * 128, bx * 128])

    target = tvm.target.Target({"kind": "cuda", "arch": "sm_100"})
    with tvm.transform.PassContext(config=pass_configs), target:
        mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
        mod = CUDAPassPipelineBodyPrologue(mod, target)
        mod = tl.cuda.transform.LowerSharedTmem()(mod)
        mod = tl.transform.IfStmtBinding()(mod)
        mod = tl.transform.PlanAndUpdateBufferAllocationLocation()(mod)

        barrier_blocks = _collect_barrier_blocks(mod["main"].body)
        assert len(barrier_blocks) == 1
        assert "barrier_init" in barrier_blocks[0].annotations

        mod = tl.cuda.transform.LowerSharedBarrier()(mod)

    body = mod["main"].body
    assert len(_collect_init_barrier_calls(body)) == 1
    assert len(_collect_fence_barrier_init(body)) == 1


if __name__ == "__main__":
    tilelang.testing.main()
