# isort: off
# fmt: off
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass

import triton
from triton_kernels import target_info
from triton_kernels.target_info import get_cdna_version, get_rdna_version, cuda_capability_geq
from triton_kernels.tensor import FP4, FP32, Tensor, torch_dtype_to_dtype
import torch
from triton_kernels.tensor_details.layout_details.hopper_scale import HopperMXScaleLayout
from triton_kernels.tensor_details.layout_details.strided import StridedLayout
from triton_kernels.tensor_details.layout_details.base import Layout
from triton_kernels.tensor_details.layout_details.blackwell_value_shuffled import BlackwellMX4ValueShuffledLayout
from .opt_flags_details import opt_flags_amd, opt_flags_nvidia

_MIN_NVFP4_RAGGED_TRAIN_ROWS = 65_536

@dataclass
class OptFlags:
    block_m: int
    block_n: int
    block_k: int
    num_warps: int
    num_stages: int
    group_m: int
    xcd_swizzle: int
    w_cache_modifier: str
    split_k: int
    is_persistent: bool
    idle_sms: int
    epilogue_subtile: int | None
    arch: str
    occupancy_target: int
    target_kernel_kwargs: dict
    clc: bool = False


def max_allowable_mn(
    max_mn: int,
    m: int,
    n: int,
    split_k: int,
    ) -> int:
    return 1 if m * n >= max_mn else split_k


def all_constraints_satisfied(opt_flags: OptFlags, constraints: dict) -> bool:
    _split_k_constraints = ['split_k', 'max_allowable_mn', 'disable_mx4_block_swap']
    assert all(getattr(opt_flags, ck) == cv for ck, cv in constraints.items() if cv is not None and ck not in _split_k_constraints)
    if constraints.get('split_k') and not constraints.get('max_allowable_mn'):
        assert opt_flags.split_k == constraints['split_k']


def make_default_opt_flags_amd(
    out_dtype,
    lhs_dtype,
    rhs_dtype,
    precision_config,
    batch_size,
    m,
    n,
    k,
    ragged_metadata,
    can_use_persistent_tma,
    can_use_split_k,
    enforce_bitwise_invariance,
    epilogue_effective_itemsize,
    x_transpose,
    has_y_acc_in,
    constraints,
):
    constraints_supported = {"block_m", "block_n", "block_k", "split_k", "is_persistent", "epilogue_subtile", "idle_sms", "max_allowable_mn", "num_warps", "disable_mx4_block_swap"}
    unsupported = set(constraints.keys()) - constraints_supported
    assert not unsupported, f"Given unsupported constraint: {unsupported}"
    # tokens per slice
    if ragged_metadata is None:
        slice_size = m
    elif ragged_metadata.expected_slice_size is None:
        slice_size = max(1, m // ragged_metadata.n_slices)
    else:
        slice_size = ragged_metadata.expected_slice_size

    is_cdna4 = get_cdna_version() == 4
    # block_m
    if constraints.get("block_m", None):
        block_m = constraints["block_m"]
    elif enforce_bitwise_invariance:
        block_m = 256 if is_cdna4 else 128
    elif slice_size >= 512 and n >= 2048:
        block_m = 256 if is_cdna4 else 128
    elif is_cdna4 and m >= 512:
        block_m = 128
    elif get_rdna_version() in (3, 4) and m >= 512:
        block_m = 64
    else:
        block_m = max(32, min(triton.next_power_of_2(slice_size), 64))

    if ragged_metadata is not None:
        grid_m = ragged_metadata.n_blocks(ragged_metadata.n_slices, m, block_m)
    else:
        grid_m = triton.cdiv(m, block_m)
    # group_m:
    group_m = 4
    # number of xcds
    num_xcds = 8
    xcd_swizzle = num_xcds
    # block_nk:
    # TODO: Does opt_flags_amd.compute_block_nk need to be refactored?
    block_n, block_k = opt_flags_amd.compute_block_nk(
        n, block_m, grid_m, num_xcds, lhs_dtype, rhs_dtype, precision_config
    )
    is_persistent = constraints.get("is_persistent", False)
    # split_k:
    split_k = 1
    if constraints.get("max_allowable_mn", 0) > 0 and constraints.get("split_k") is not None:
        split_k = max_allowable_mn(constraints["max_allowable_mn"], m, n, constraints.get("split_k"))
    elif constraints.get("split_k", None) is not None:
        split_k = constraints["split_k"]
    elif can_use_split_k and not enforce_bitwise_invariance:
        grid_size = grid_m * ((n + block_n - 1) // block_n)
        n_cu = torch.cuda.get_device_properties(0).multi_processor_count
        split_k = 1 if grid_size == 0 else max(1, n_cu // grid_size)
    # w_cache_modifier:
    w_cache_modifier = ".cg" if block_m <= 32 else None
    # num_warps, num_stages
    num_warps = 2 if (m is not None and m <= 16) else 8
    num_stages = 2
    # AMD-specific
    target_kernel_kwargs = {"waves_per_eu": 0, "matrix_instr_nonkdim": 16, "kpack": 1}
    epilogue_subtile = constraints.get('epilogue_subtile', None)
    if epilogue_subtile is None:
        epilogue_subtile = 1

    # prevents OutOfSharedMemoryError for mxfp8 on CDNA3
    if get_cdna_version() == 3 and rhs_dtype.bitwidth == 8 and precision_config.b_mx_scale is not None:
        num_stages = 1

    if lhs_dtype.bitwidth == 16 and rhs_dtype.bitwidth == 4 and precision_config.b_mx_scale is not None:
        # specific configs for F16 x MXFP4 on CDNA4
        if is_cdna4:
            split_k = 1
            if m <= 1024:
                target_kernel_kwargs["waves_per_eu"] = 3
                block_n = 128
                block_k = 128
                num_warps = 4
            else:
                target_kernel_kwargs["waves_per_eu"] = 0
                block_m = 64
                block_n = 512
                block_k = 256
                num_warps = 8

        # Specific configs for F16 x MXFP4 on RDNA.
        if get_rdna_version() in (3, 4):
            split_k = 1
            target_kernel_kwargs["waves_per_eu"] = 0
            num_stages = 1
            if m <= 512:
                block_m = 16
                block_n = 64
                block_k = 512
                num_warps = 4
            else:
                block_m = 64
                block_n = 128
                block_k = 128
                num_warps = 4

    def replace_with_valid_constraint(k: str, v):
        if constraints.get(k, None) is not None:
            return constraints[k]
        else:
            return v

    ret = OptFlags(
        block_m=replace_with_valid_constraint('block_m', block_m),
        block_n=replace_with_valid_constraint('block_n', block_n),
        block_k=replace_with_valid_constraint('block_k', block_k),
        num_warps=replace_with_valid_constraint('num_warps', num_warps),
        num_stages=num_stages,
        group_m=group_m,
        xcd_swizzle=xcd_swizzle,
        w_cache_modifier=w_cache_modifier,
        split_k=split_k,
        is_persistent=is_persistent,
        clc=False,
        idle_sms=constraints.get("idle_sms", _get_idle_sms()),
        epilogue_subtile=epilogue_subtile,
        arch=None,
        target_kernel_kwargs=target_kernel_kwargs,
        occupancy_target=1,
    )
    # check constraints
    all_constraints_satisfied(ret, constraints)
    return ret

def make_default_opt_flags_nvidia(
    out_dtype,
    lhs_dtype,
    rhs_dtype,
    precision_config,
    batch_size,
    m,
    n,
    k,
    routing_data,
    can_use_persistent_tma,
    can_use_split_k,
    enforce_bitwise_invariance,
    epilogue_effective_itemsize,
    x_transpose,
    has_y_acc_in,
    constraints,
    intermediate_out_dtype,
    x_uses_tma_when_persistent=True,
    w_transpose=False,
    mx_block_size=None,
    epilogue_reduction_n=1,
):
    constraints_supported = {"block_m", "block_n", "block_k", "split_k", "is_persistent", "clc", "epilogue_subtile", "num_stages", "idle_sms", "max_allowable_mn", "num_warps", "disable_mx4_block_swap"}
    unsupported = set(constraints.keys()) - constraints_supported
    assert not unsupported, f"Given unsupported constraint: {unsupported}"
    is_large_ragged_nvfp4 = (
        routing_data is not None
        and m >= _MIN_NVFP4_RAGGED_TRAIN_ROWS
        and lhs_dtype == FP4
        and rhs_dtype == FP4
        and precision_config.a_mx_tensor_scale is not None
        and precision_config.b_mx_tensor_scale is not None
        and cuda_capability_geq(10, 3)
    )
    # tokens per expert
    if routing_data is None or batch_size > 1:
        slice_size = m
    elif routing_data.expected_slice_size is None:
        slice_size = max(1, m // routing_data.n_slices)
    else:
        slice_size = routing_data.expected_slice_size
    # pid swizzling
    group_m = 8
    if lhs_dtype == FP4 and rhs_dtype == FP4:
        group_m = 16
    xcd_swizzle = 1
    # block_m
    if constraints.get("block_m", None):
        block_m = constraints["block_m"]
    elif is_large_ragged_nvfp4:
        block_m = 128
    elif enforce_bitwise_invariance:
        block_m = 128
    else:
        if slice_size <= 64 and routing_data is not None and routing_data.slice_sizes is not None:
            # Ragged and likely memory bound; set the block size higher to minimize loading weights more than once.
            if (
                lhs_dtype.bitwidth == 16
                and rhs_dtype.bitwidth == 4
                and slice_size >= 16
                and torch.cuda.get_device_capability()[0] >= 10
            ):
                block_m = max(16, min(triton.next_power_of_2(8 * slice_size), 128))
            else:
                block_m = max(16, min(triton.next_power_of_2(2 * slice_size), 64))
            if block_m == 64 and precision_config.c_mx_scale is not None and rhs_dtype == FP4 and torch.cuda.get_device_capability()[0] >= 10:
                # when having both fused_activation and mxfp8 downcast in epilogue, block_m=64 causing shared memory overflow
                block_m = 128
        else:
            block_m = max(16, min(triton.next_power_of_2(slice_size), 128))
    # block n
    arch = None
    if constraints.get("block_n", None) is not None:
        # A single block_n constraint should apply consistently to both the
        # regular and persistent/TMA paths.
        block_n = constraints["block_n"]
        block_n_tma = constraints["block_n"]
    elif is_large_ragged_nvfp4:
        block_n, block_n_tma = 256, 256
    else:
        block_n, block_n_tma = opt_flags_nvidia.compute_block_n(n, arch, precision_config)
    # is_persistent
    grid_size_tma = opt_flags_nvidia.compute_grid_size(routing_data, batch_size, m, n, block_m, block_n_tma)
    n_sms = torch.cuda.get_device_properties(0).multi_processor_count
    tiles_per_sm = grid_size_tma / n_sms
    supports_persistent = can_use_persistent_tma and (arch is None or int(arch[2:-1]) >= 9)
    clc = constraints.get("clc", False)
    if clc:
        if not cuda_capability_geq(10, 0):
            raise InapplicableConstraint("cannot enforce `clc=True` constraint before NVIDIA SM100")
        if not supports_persistent:
            raise InapplicableConstraint("cannot enforce `clc=True` constraint without persistent TMA support")
        if constraints.get("is_persistent", True) is False:
            raise InapplicableConstraint("`clc=True` requires `is_persistent=True`")
    a_mx_scale_layout = None if not isinstance(precision_config.a_mx_scale, Tensor) else precision_config.a_mx_scale.storage.layout
    b_mx_scale_layout = None if not isinstance(precision_config.b_mx_scale, Tensor) else precision_config.b_mx_scale.storage.layout

    def _is_layout_strided(layout: Layout | None) -> bool:
        return layout is None or isinstance(layout, StridedLayout)

    requires_persistent = (not _is_layout_strided(a_mx_scale_layout) or not _is_layout_strided(b_mx_scale_layout)) and target_info.has_native_mxfp()
    if clc:
        is_persistent = True
    elif constraints.get("is_persistent", None) is not None:
        if requires_persistent and not constraints["is_persistent"]:
            raise InapplicableConstraint("cannot enforce `is_persistent=False` constraint because persistent kernel is required")
        is_persistent = constraints["is_persistent"]
    elif requires_persistent:
        assert supports_persistent, "persistent kernel required but not supported"
        is_persistent = True
    else:
        has_simple_epilogue = precision_config.max_num_imprecise_acc is None
        is_persistent = supports_persistent and has_simple_epilogue and (tiles_per_sm >= 2.0 or lhs_dtype.bitwidth <= 8) and (out_dtype.bitwidth < 32 or lhs_dtype.bitwidth == 32 or rhs_dtype.bitwidth == 32)
        # TMA is slower for batched matmuls with small m/n/k.
        if m * n * k < 131072:
            is_persistent = False
        if isinstance(b_mx_scale_layout, HopperMXScaleLayout):
            # TODO: persistent kernel is currently slower than non-persistent
            is_persistent = False

    # adjust block_n based on is_persistent signal
    block_n = block_n_tma if is_persistent else block_n
    if (is_persistent and constraints.get("block_n", None) is None
            and cuda_capability_geq(10, 0) and (lhs_dtype == FP32 or rhs_dtype == FP32)
            and not x_uses_tma_when_persistent):
        # Blackwell's fp32/tf32 persistent dot stages an operand in TMEM in
        # addition to the accumulator. A 128x256 accumulator already consumes
        # the full 512-column TMEM budget, so leave headroom for that operand.
        block_n = min(block_n, 128)
    if (is_persistent and constraints.get("block_n", None) is None
            and opt_flags_nvidia.is_blackwell_mx_lhs_dense_rhs(precision_config, lhs_dtype, rhs_dtype)):
        # Native Blackwell MX lhs + dense rhs persistent dots also stage an
        # expanded operand in TMEM, so keep the accumulator tile below the
        # 512-column budget.
        block_n = min(block_n, 128)
    # adjust block_m based on is_persistent signal
    if is_persistent and opt_flags_nvidia.is_x_scale_swizzled(precision_config):
        # a mx scale has been swizzled to BlackwellActMXScaleLayout, enforce block_m=128 to align with swizzling layout
        block_m = 128
    # block k
    if constraints.get("block_k", None) is not None:
        block_k = constraints["block_k"]
    elif is_large_ragged_nvfp4:
        block_k = 256
    else:
        block_k = opt_flags_nvidia.compute_block_k(m, k, is_persistent, lhs_dtype, rhs_dtype, precision_config, has_y_acc_in)
    if block_n == 256 and block_k == 128 and block_m <= 64 and is_persistent and rhs_dtype == FP4 and k >= 4096 and slice_size > 1 and lhs_dtype != torch.bfloat16 and not constraints.get("disable_mx4_block_swap", False):
        # Swap block_n and block_k for mxfp4 weights so that block_k is a full cacheline, so long as K is sufficiently large.
        # TODO: swizzle the HBM layout of the weights instead
        block_n, block_k = block_k, block_n
    # split_k
    split_k = 1
    if constraints.get("max_allowable_mn", 0) > 0 and constraints.get("split_k") is not None:
        split_k = max_allowable_mn(constraints["max_allowable_mn"], m, n, constraints.get("split_k"))
    elif constraints.get("split_k", None) is not None:
        split_k = constraints["split_k"]
    elif can_use_split_k and not enforce_bitwise_invariance:
        estimated_actual_grid_size = opt_flags_nvidia.compute_grid_size(None, batch_size, m, n, block_m, block_n)
        split_k = opt_flags_nvidia.compute_split_k(block_k, k, estimated_actual_grid_size)
    if split_k > 1:
        # Split-K writes full-N scratch and applies fused reductions in the
        # reduce kernel, not in the matmul epilogue.
        epilogue_reduction_n = 1
    compute_num_stages_args = (
        precision_config,
        is_persistent,
        block_m,
        block_n,
        block_k,
        torch_dtype_to_dtype(intermediate_out_dtype) if split_k > 1 else out_dtype,
        lhs_dtype,
        rhs_dtype,
        x_transpose,
        epilogue_effective_itemsize,
        has_y_acc_in,
        mx_block_size,
        epilogue_reduction_n,
    )

    num_warps = opt_flags_nvidia.compute_num_warps(block_m, block_n, is_persistent, precision_config, constraints)
    if is_large_ragged_nvfp4 and constraints.get("num_warps", None) is None:
        # Eight warps overrun GB300 TMEM for the large ragged NVFP4 tile.
        num_warps = 4
    if (constraints.get("num_warps", None) is None
            and is_persistent
            and block_n <= 128
            and block_k >= 256
            and opt_flags_nvidia.is_blackwell_mx_lhs_dense_rhs(precision_config, lhs_dtype, rhs_dtype)):
        num_warps = max(num_warps, 8)

    # Occupancy target and maxnreg (for Hopper)
    occupancy_target = 1
    is_hopper_scale = isinstance(b_mx_scale_layout, HopperMXScaleLayout)
    if is_hopper_scale:
        occupancy_target = 16 // num_warps
        if precision_config.a_mx_scale is not None and precision_config.c_mx_scale is not None:
            # Hopper MXFP4 RHS plus MX input/output needs more than the
            # 128-register cap implied by the default occupancy target.
            occupancy_target = 1
    threads_per_warp = 32
    reg_per_sm = 64 * 1024
    max_reg_per_thread = 256
    is_blackwell_or_newer = cuda_capability_geq(10, 0)
    if is_persistent and not is_blackwell_or_newer:
        maxnreg = reg_per_sm // (num_warps * threads_per_warp * occupancy_target)
        maxnreg = min(max_reg_per_thread, maxnreg)
    else:
        maxnreg = None

    if constraints.get("epilogue_subtile", None) is not None:
        subtiles_to_check = [constraints["epilogue_subtile"]]
    elif is_large_ragged_nvfp4:
        subtiles_to_check = [2]
    else:
        subtiles_to_check = [1] if out_dtype == FP4 else [1, 2, 4]
    num_stages = -1
    for ep in subtiles_to_check:
        ns = opt_flags_nvidia.compute_num_stages(*compute_num_stages_args, epilogue_subtile=ep,
                                                 occupancy_target=occupancy_target,
                                                 w_transpose=w_transpose)
        if ns > num_stages:
            epilogue_subtile, num_stages = ep, ns

    if constraints.get("num_stages", None):
        num_stages = constraints["num_stages"]
    elif is_large_ragged_nvfp4 and not any(
        key in constraints
        for key in (
            "block_m",
            "block_n",
            "block_k",
            "num_warps",
            "is_persistent",
            "split_k",
            "max_allowable_mn",
            "epilogue_subtile",
        )
    ):
        num_stages = 3 if epilogue_reduction_n == 1 else 2
    assert num_stages >= 1
    ret = OptFlags(
        block_m=block_m,
        block_n=block_n,
        block_k=block_k,
        num_warps=num_warps,
        num_stages=num_stages,
        group_m=group_m,
        xcd_swizzle=xcd_swizzle,
        w_cache_modifier=None,
        split_k=split_k,
        is_persistent=is_persistent,
        clc=clc,
        epilogue_subtile=epilogue_subtile,
        arch=arch,
        target_kernel_kwargs=dict(
            maxnreg=maxnreg,
            # For some reason, overlapping the epilogue is slower for hopper bf16 x mxfp4
            FLATTEN_LOOPS=not is_hopper_scale,
        ),
        idle_sms=constraints.get("idle_sms", _get_idle_sms()),
        occupancy_target=occupancy_target,
    )
    # check constraints
    all_constraints_satisfied(ret, constraints)
    return ret

# --------------
# User Interface
# --------------

_opt_flags_constraints: ContextVar[dict | None] = ContextVar("opt_flags_constraints", default=None)
_opt_flags: ContextVar[OptFlags | None] = ContextVar("opt_flags", default=None)
_idle_sms = 0

def _get_idle_sms() -> int:
    return _idle_sms

def set_idle_sms(num_idle_sms: int):
    global _idle_sms
    _idle_sms = num_idle_sms

def _get_opt_flags_constraints() -> dict:
    constraints = _opt_flags_constraints.get()
    return {} if constraints is None else constraints

def update_opt_flags_constraints(constraints: dict[str, int]):
    updated = _get_opt_flags_constraints().copy()
    updated.update(constraints)
    _opt_flags_constraints.set(updated)

def reset_opt_flags_constraints():
    _opt_flags_constraints.set(None)

@contextmanager
def scoped_opt_flags_constraints(constraints):
    updated = _get_opt_flags_constraints().copy()
    updated.update(constraints)
    token = _opt_flags_constraints.set(updated)
    try:
        yield
    finally:
        _opt_flags_constraints.reset(token)

def reset_opt_flags():
    _opt_flags.set(None)

def set_opt_flags(opt_flags: OptFlags):
    assert not _get_opt_flags_constraints(), "setting constraints is incompatible with manual flags override"
    assert _opt_flags.get() is None, "opt_flags already set; please reset to None first"
    _opt_flags.set(opt_flags)

@contextmanager
def scoped_opt_flags(opt_flags: OptFlags):
    token = _opt_flags.set(opt_flags)
    try:
        yield
    finally:
        _opt_flags.reset(token)

class InapplicableConstraint(Exception):
    pass

def make_opt_flags(
    out_dtype,
    lhs_dtype,
    rhs_dtype,
    precision_config,
    batch_size,
    m,
    n,
    k,
    ragged_metadata,
    can_use_persistent_tma,
    can_use_split_k,
    epilogue_effective_itemsize,
    x_transpose,
    has_y_acc_in,
    block_k,
    intermediate_out_dtype,
    mx_block_size=None,
    x_uses_tma_when_persistent=True,
    w_transpose=False,
    rhs_layout=None,
    epilogue_reduction_n=1,
):
    # Empty outputs do not launch a kernel, so persistent TMA constraints are vacuous.
    can_use_persistent_tma = can_use_persistent_tma or batch_size * m * n == 0
    opt_flags_constraints = _get_opt_flags_constraints()
    if opt_flags_constraints.get("is_persistent", False) and not can_use_persistent_tma:
        raise InapplicableConstraint("cannot enforce `is_persistent=True` constraint")
    if opt_flags_constraints.get("split_k") is not None and opt_flags_constraints.get("split_k") > 1 and not can_use_split_k:
        raise InapplicableConstraint("cannot enforce `split_k=True` constraint")
    if opt_flags_constraints.get("max_allowable_mn"):
        if not opt_flags_constraints.get("split_k"):
            raise InapplicableConstraint("split_k also needs to be provided with max_allowable_mn")
    enforce_bitwise_invariance = precision_config.enforce_bitwise_invariance
    opt_flags = _opt_flags.get()
    if opt_flags is not None:
        assert not opt_flags_constraints
        assert block_k is None
        return opt_flags
    if isinstance(rhs_layout, BlackwellMX4ValueShuffledLayout):
        opt_flags_constraints = opt_flags_constraints.copy()
        opt_flags_constraints.setdefault("block_k", rhs_layout.block_k)
        opt_flags_constraints.setdefault("block_n", rhs_layout.block_n)
        opt_flags_constraints.setdefault("disable_mx4_block_swap", True)
        if (
            rhs_layout.block_n == 512
            and not enforce_bitwise_invariance
            and not opt_flags_nvidia.is_x_scale_swizzled(precision_config)
        ):
            opt_flags_constraints.setdefault("block_m", 64)
    if block_k is not None:
        opt_flags_constraints = opt_flags_constraints.copy()
        opt_flags_constraints.update(block_k=block_k, split_k=1)
    args = [out_dtype, lhs_dtype, rhs_dtype, precision_config, batch_size, m, n, k,
            ragged_metadata, can_use_persistent_tma, can_use_split_k,
            enforce_bitwise_invariance, epilogue_effective_itemsize, x_transpose, has_y_acc_in,
            opt_flags_constraints]
    backend = triton.runtime.driver.active.get_current_target().backend
    if backend == "hip":
        return make_default_opt_flags_amd(*args)
    if backend == "cuda":
        return make_default_opt_flags_nvidia(
            *args,
            intermediate_out_dtype,
            x_uses_tma_when_persistent=x_uses_tma_when_persistent,
            w_transpose=w_transpose,
            mx_block_size=mx_block_size,
            epilogue_reduction_n=epilogue_reduction_n,
        )
    assert False
