# SPDX-License-Identifier: MIT
import logging
import torch
from dataclasses import dataclass
from typing import Optional, Dict, Any, Tuple, List

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

class MoeSolutionType:
    MOE_C = "moe_c"
    ASM = "asm"
    TRITON = "triton"
    CK = "ck"


class MoeQuantType:
    """Quantization types supported by get_aiter_moe_config / aiter_moe."""
    W16A16 = "w16a16"
    W4A16 = "w4a16"
    W4A8 = "w4a8"
    W8A8 = "int8_w8a8"
    FP8_W8A8 = "fp8_w8a8"
    INT8_W8A16 = "int8_w8a16"
    FP8_W8A16 = "fp8_w8a16"


@dataclass
class AiterMoeConfig:
    """Config returned by :func:`get_aiter_moe_config`.

    Attributes:
        quant_type: The quantization type this config was obtained for.
        solution_type: Which backend to use (MoeSolutionType constant), or
            None if no solution was found.
        config: Backend-specific config dict (opaque to the caller).
    """
    quant_type: Optional[str] = None
    solution_type: Optional[str] = None
    config: Optional[Dict[str, Any]] = None


def _pick_closest_config(configs: Dict[int, Any], m: int) -> Dict[str, Any]:
    return configs[min(configs.keys(), key=lambda x: abs(x - m))]


def _try_get_moe_c_config(
    quant_type: str,
    m: int,
    e: int,
    n: int,
    block_size: int,
) -> Optional[Dict[str, Any]]:
    try:
        from .fused_moe_c import get_moe_configs_marlin

        if quant_type == MoeQuantType.W4A16:
            configs = get_moe_configs_marlin(
                E=e,
                N=n,
                dtype="int4_w4a16",
                is_bottom=False,
                use_moe_wna16_cuda=True,
            )
        elif quant_type == MoeQuantType.W8A8:
            configs = get_moe_configs_marlin(
                E=e,
                N=n,
                dtype="int8_w8a8",
                is_bottom=False,
                use_moe_wna16_cuda=True,
            )
        elif quant_type == MoeQuantType.FP8_W8A8:
            configs = get_moe_configs_marlin(
                E=e,
                N=n,
                dtype="fp8_w8a8",
                is_bottom=False,
                use_moe_wna16_cuda=True,
            )
        elif quant_type == MoeQuantType.W4A8:
            configs = get_moe_configs_marlin(
                E=e,
                N=n,
                dtype="int8_w4a8",
                block_n = block_size,
                block_k = block_size,
                is_bottom=False,
                use_moe_wna16_cuda=True,
            )
        elif quant_type == MoeQuantType.INT8_W8A16:
            configs = get_moe_configs_marlin(
                E=e,
                N=n,
                dtype="int8_w8a16",
                block_n=0,
                block_k=block_size if block_size else 0,
                is_bottom=False,
                use_moe_wna16_cuda=True,
            )
        else:
            return None

        if configs is None:
            return None
        return _pick_closest_config(configs, m)
    except Exception as exc:
        logger.debug("moe_c config lookup failed for %s: %s", quant_type, exc)
        return None


def _try_get_asm_config(
    quant_type: str,
    m: int,
    e: int,
    n: int,
    k: int,
    top_k: int,
    block_size: Optional[int],
) -> Optional[Dict[str, Any]]:
    try:
        from .fused_moe_asm_wna16 import get_moe_asm_solution, MoeQuantType as AsmMoeQuantType
        from .jit.utils.chip_info import get_gfx

        arch = get_gfx()

        if quant_type == MoeQuantType.W4A16:
            from .fused_moe_asm_wna16 import decode_sol_w4a16, decode_sol_w4a16_gw32
            if block_size == 32:
                if top_k > 8 or n != 256 or k != 7168:
                    return None
                else:
                    return decode_sol_w4a16_gw32()

            solution = get_moe_asm_solution(
                arch=arch,
                token=m,
                inter_dim=n,
                model_dim=k,
                expert=e,
                topk=top_k,
                quant_type=AsmMoeQuantType.INT4_W4A16,
            )
            if solution == "default":
                return None
            return decode_sol_w4a16(solution)

        if quant_type == MoeQuantType.W8A8:
            from .fused_moe_asm_wna16 import decode_sol_0
            asm_quant_type = AsmMoeQuantType.INT8_W8A8_C if (block_size == 0 or block_size is None) else AsmMoeQuantType.INT8_W8A8
            solution = get_moe_asm_solution(
                arch=arch,
                token=m,
                inter_dim=n,
                model_dim=k,
                expert=e,
                topk=top_k,
                quant_type=asm_quant_type,
            )
            if solution == "default":
                return None
            return decode_sol_0(solution)

        if quant_type == MoeQuantType.FP8_W8A8:
            from .fused_moe_asm_wna16 import decode_sol_0
            asm_quant_type = AsmMoeQuantType.F8_W8A8_C if (block_size == 0 or block_size is None) else AsmMoeQuantType.F8_W8A8
            solution = get_moe_asm_solution(
                arch=arch,
                token=m,
                inter_dim=n,
                model_dim=k,
                expert=e,
                topk=top_k,
                quant_type=asm_quant_type,
            )
            if solution == "default":
                return None
            return decode_sol_0(solution)

        if quant_type == MoeQuantType.W16A16:
            from .fused_moe_asm_wna16 import decode_sol_0

            solution = get_moe_asm_solution(
                arch=arch,
                token=m,
                inter_dim=n,
                model_dim=k,
                expert=e,
                topk=top_k,
                quant_type=AsmMoeQuantType.NO_QUANT,
            )
            if solution == "default":
                return None
            return decode_sol_0(solution)

        return None
    except Exception as exc:
        logger.debug("ASM config lookup failed for %s: %s", quant_type, exc)
        return None


def _try_get_triton_config(
    quant_type: str,
    m: int,
    e: int,
    n: int,
    block_size: int,
) -> Optional[Dict[str, Any]]:
    try:
        from .ops.triton.utils.moe_config_utils import get_moe_configs as triton_get_moe_configs

        if quant_type == MoeQuantType.W16A16:
            return {}  # Non-quantized; no tuned config lookup needed

        dtype_name = {
            MoeQuantType.W4A16: "int4_w4a16",
            MoeQuantType.W8A8: "int8_w8a8",
            MoeQuantType.FP8_W8A8: "fp8_w8a8",
            MoeQuantType.INT8_W8A16: "int8_w8a16",
        }.get(quant_type)
        if dtype_name is None:
            return None

        configs = triton_get_moe_configs(
            E=e,
            N=n,
            dtype=dtype_name,
            block_n=0,
            block_k=block_size if block_size else 0,
            is_bottom=False,
        )
        if configs is None:
            return None
        return _pick_closest_config(configs, m)
    except Exception as exc:
        logger.debug("Triton config lookup failed for %s: %s", quant_type, exc)
        return None


def _try_get_ck_config(
    quant_type: str,
    m: int,
    e: int,
    n: int,
    k: int,
    top_k: int,
    block_shape: Optional[List[int]],
) -> Optional[Dict[str, Any]]:
    try:
        from .fused_moe_ck import get_moe_ck_solution_id, MoeQuantType as CkMoeQuantType
        from .jit.utils.chip_info import get_gfx
        
        if quant_type == MoeQuantType.W16A16:
            ck_quant_type = CkMoeQuantType.NO_QUANT
        elif quant_type == MoeQuantType.W8A8 or quant_type == MoeQuantType.FP8_W8A8:
            ck_quant_type = CkMoeQuantType.INT8_W8A8
        else:
            return None

        arch = get_gfx()
        q_size_n = block_shape[0] if block_shape is not None else 0
        q_size_k = block_shape[1] if block_shape is not None else 0
        solution_id = get_moe_ck_solution_id(
            arch,
            ck_quant_type,
            m,
            n,
            k,
            e,
            top_k,
            q_size_n,
            q_size_k,
        )
        return {"solution_id": solution_id}
    except Exception as exc:
        logger.debug("CK config lookup failed for %s: %s", quant_type, exc)
        return None


def get_aiter_moe_config(
    M: int,     # Number of tokens (input sequence length)
    E: int,     # Number of experts
    N1: int,    # GEMM1 output dimension: gated = (intermediate_size * 2), non-gated = intermediate_size
    N2: int,    # GEMM2 output dimension, typically equal to hidden_size
    K: int,     # GEMM1 input dimension, typically equal to hidden_size;  for GEMM2, K typically equal to (moe_intermediate_size / TP)
    top_k: int,
    block_size: int,
    dtype: torch.dtype,
    quant_type: str,
    activation: str = "silu",  # "silu"/"gelu"/"relu2"/"swigluoai"/"swiglustep"...
    gated: Optional[bool] = None,  # True=GLU-gated (N1=2*inter), False=non-gated (N1=inter); None=auto from activation
) -> Tuple[bool, AiterMoeConfig]:
    """Get the best backend config for a MOE problem.

    Currently supported quant types:
    - ``MoeQuantType.W16A16`` (non-quantized)
    - ``MoeQuantType.W4A16``
    - ``MoeQuantType.W8A8`` (int8)
    - ``MoeQuantType.FP8_W8A8`` (fp8)
    - ``MoeQuantType.W4A8``
    - ``MoeQuantType.INT8_W8A16`` (int8 weight, fp16/bf16 activation)

    Backend priority:
    - ``w16a16``: asm > triton
    - ``w4a16``: moe_c > asm > triton
    - ``w8a8``: asm > moe_c > triton > ck
    - ``fp8_w8a8``: asm > moe_c > triton > ck
    - ``w4a8``: moe_c
    - ``int8_w8a16``: moe_c > triton (ASM kernel not available)
    - ``fp8_w8a16``: not yet implemented (raises NotImplementedError)

    For non-gated MOE (e.g. Nemotron with ReLU² activation), pass
    ``gated=False`` (or let it auto-detect from ``activation="relu2"``)
    and set ``N1 = intermediate_size`` (not ``2 * intermediate_size``).
    """
    # Determine gating: explicit > auto-detect from activation
    if gated is None:
        gated = activation in ("silu", "gelu", "swigluoai", "swiglustep")

    # For gated (GLU): N1 = 2 * intermediate_size, n = N1 // 2
    # For non-gated:   N1 = intermediate_size,     n = N1
    n = N1 // 2 if gated else N1
    block_shape = [0, block_size] if block_size else None

    if quant_type == MoeQuantType.W4A16:
        if dtype == torch.float16:
            candidates = [
                (MoeSolutionType.MOE_C, lambda: _try_get_moe_c_config(quant_type, M, E, n, block_size)),
                (MoeSolutionType.TRITON, lambda: _try_get_triton_config(quant_type, M, E, n, block_size)),
            ]
        elif dtype == torch.bfloat16:
            candidates = [
                (MoeSolutionType.ASM, lambda: _try_get_asm_config(quant_type, M, E, n, K, top_k, block_size)),
                (MoeSolutionType.TRITON, lambda: _try_get_triton_config(quant_type, M, E, n, block_size)),
                (MoeSolutionType.MOE_C, lambda: _try_get_moe_c_config(quant_type, M, E, n, block_size)),
            ]
        else:
            raise ValueError(f"Unsupported dtype: {dtype}")
    elif quant_type in (MoeQuantType.W8A8, MoeQuantType.FP8_W8A8):
        if block_size is None or block_size == 0: # Channel wise
            candidates = [
                (MoeSolutionType.ASM, lambda: _try_get_asm_config(quant_type, M, E, n, K, top_k, block_size)),
                (MoeSolutionType.MOE_C, lambda: _try_get_moe_c_config(quant_type, M, E, n, block_size)),
                (MoeSolutionType.TRITON, lambda: _try_get_triton_config(quant_type, M, E, n, block_size)),
                # (MoeSolutionType.CK, lambda: _try_get_ck_config(quant_type, M, E, n, K, top_k, block_shape)),
            ]
        else: # Block wise choose ASM
            candidates = [
                (MoeSolutionType.ASM, lambda: _try_get_asm_config(quant_type, M, E, n, K, top_k, block_size)),
                (MoeSolutionType.TRITON, lambda: _try_get_triton_config(quant_type, M, E, n, block_size)),
            ]

    elif quant_type == MoeQuantType.W4A8:
        candidates = [
            (MoeSolutionType.MOE_C, lambda: _try_get_moe_c_config(quant_type, M, E, n, block_size)),
            (MoeSolutionType.TRITON, lambda: _try_get_triton_config(quant_type, M, E, n, block_size)),
            # (MoeSolutionType.ASM, lambda: _try_get_asm_config(quant_type, M, E, n, K, top_k)),
        ]
    elif quant_type == MoeQuantType.INT8_W8A16:
        # ASM backend currently has no W8A16 kernel/CSV; skip ASM and use moe_c -> triton.
        candidates = [
            (MoeSolutionType.MOE_C, lambda: _try_get_moe_c_config(quant_type, M, E, n, block_size)),
            (MoeSolutionType.TRITON, lambda: _try_get_triton_config(quant_type, M, E, n, block_size)),
        ]
    elif quant_type == MoeQuantType.FP8_W8A16:
        # No backend currently implements FP8 weight + 16-bit activation MoE.
        raise NotImplementedError(
            "MoeQuantType.FP8_W8A16 is not yet supported by any aiter MOE backend (asm/moe_c/triton)."
        )
    elif quant_type == MoeQuantType.W16A16:
        candidates = [
            (MoeSolutionType.ASM, lambda: _try_get_asm_config(quant_type, M, E, n, K, top_k, None)),
            (MoeSolutionType.TRITON, lambda: _try_get_triton_config(quant_type, M, E, n, block_size)),
            # (MoeSolutionType.CK, lambda: _try_get_ck_config(quant_type, M, E, n, K, top_k, block_shape)),
        ]
    else:
        raise ValueError(f"Unsupported quant_type: {quant_type}")

    for solution_type, get_config in candidates:
        config = get_config()
        if config is not None:
            return True, AiterMoeConfig(
                quant_type=quant_type,
                solution_type=solution_type,
                config=config,
            )

    return False, AiterMoeConfig(quant_type=quant_type)


def aiter_moe(
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    moe_config: AiterMoeConfig,
    inplace: Optional[bool] = False,
    activation: str = "silu",
    w1_scale: Optional[torch.Tensor] = None,
    w2_scale: Optional[torch.Tensor] = None,
    w1_zp: Optional[torch.Tensor] = None,
    w2_zp: Optional[torch.Tensor] = None,
    a1_scale: Optional[torch.Tensor] = None,
    a2_scale: Optional[torch.Tensor] = None,
    block_shape: Optional[list] = None,
    global_num_experts: int = -1,
    expert_map: Optional[torch.Tensor] = None,
    routed_scaling_factor: Optional[float] = 1.0,
    use_weight_shuffle: bool = False,
    output_dtype: Optional[torch.dtype] = None,
    gemm1_alpha: Optional[float] = None,
    gemm1_limit: Optional[float] = None
) -> torch.Tensor:
    """Execute MOE using the backend and quant type described by *moe_config*."""
    if moe_config.solution_type is None or moe_config.quant_type is None:
        raise ValueError(
            "moe_config has no valid solution_type/quant_type. "
            "Call get_aiter_moe_config first and check the status."
        )
    
    if output_dtype is None:
        output_dtype = hidden_states.dtype

    use_int4_w4a16 = moe_config.quant_type == MoeQuantType.W4A16
    use_int8_w8a8 = moe_config.quant_type == MoeQuantType.W8A8
    use_fp8_w8a8 = moe_config.quant_type == MoeQuantType.FP8_W8A8
    use_int8_w4a8 = moe_config.quant_type == MoeQuantType.W4A8
    use_int8_w8a16 = moe_config.quant_type == MoeQuantType.INT8_W8A16

    if moe_config.solution_type == MoeSolutionType.MOE_C:
        from .fused_moe_c import moe_c_fused_experts

        return moe_c_fused_experts(
            hidden_states,
            w1,
            w2,
            topk_weights,
            topk_ids,
            inplace=inplace,
            use_int4_w4a16=use_int4_w4a16,
            use_int8_w8a8=use_int8_w8a8,
            use_fp8_w8a8=use_fp8_w8a8,
            use_int8_w4a8=use_int8_w4a8,
            use_int8_w8a16=use_int8_w8a16,
            activation=activation,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            w1_zp=w1_zp,
            w2_zp=w2_zp,
            a1_scale=a1_scale,
            a2_scale=a2_scale,
            block_shape=block_shape,
            routed_scaling_factor=routed_scaling_factor,
            gemm1_alpha=gemm1_alpha,
            gemm1_limit=gemm1_limit
        )

    if moe_config.solution_type == MoeSolutionType.ASM:
        from .fused_moe_asm_wna16 import fused_experts_asm_impl
        per_channel_quant = True if block_shape is None else False
        cfg = moe_config.config
        solution_id = f"{cfg['SOL_ID1']}+{cfg['SOL_ID2']}"
        return fused_experts_asm_impl(
            hidden_states,
            w1,
            w2,
            topk_weights,
            topk_ids,
            dtype=output_dtype,
            inplace=inplace,
            use_int4_w4a16=use_int4_w4a16,
            use_int8_w8a8=use_int8_w8a8,
            use_fp8_w8a8=use_fp8_w8a8,
            activation=activation,
            per_channel_quant = per_channel_quant,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            w1_zp=w1_zp,
            w2_zp=w2_zp,
            a1_scale=a1_scale,
            a2_scale=a2_scale,
            block_shape=block_shape,
            use_shuffle=use_weight_shuffle,
            solution_id=solution_id,
            routed_scaling_factor=routed_scaling_factor,
            gemm1_alpha=gemm1_alpha,
            gemm1_limit=gemm1_limit
        )

    if moe_config.solution_type == MoeSolutionType.TRITON:
        from .ops.triton.fused_moe import fused_experts_impl

        # W8A8 / W8A16 channel-wise (block_shape=None) requires per_channel_quant=True
        per_channel_quant = (use_int8_w8a8 or use_fp8_w8a8 or use_int8_w8a16) and block_shape is None

        return fused_experts_impl(
            hidden_states,
            w1,
            w2,
            topk_weights,
            topk_ids,
            output_dtype=output_dtype,
            inplace=inplace,
            use_int4_w4a16=use_int4_w4a16,
            use_int8_w8a8=use_int8_w8a8,
            use_fp8_w8a8=use_fp8_w8a8,
            use_int8_w8a16=use_int8_w8a16,
            activation=activation,
            per_channel_quant=per_channel_quant,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            w1_zp=w1_zp,
            w2_zp=w2_zp,
            a1_scale=a1_scale,
            a2_scale=a2_scale,
            block_shape=block_shape,
            routed_scaling_factor=routed_scaling_factor,
            gemm1_alpha=gemm1_alpha,
            gemm1_limit=gemm1_limit
        )

    if moe_config.solution_type == MoeSolutionType.CK:
        from .fused_moe_ck import run_fused_experts_ck_impl

        solution_id = moe_config.config["solution_id"]
        return run_fused_experts_ck_impl(
            hidden_states,
            w1,
            w2,
            topk_weights,
            topk_ids,
            odtype=output_dtype,
            inplace=inplace,
            use_int8_w8a8=use_int8_w8a8,
            use_fp8_w8a8=use_fp8_w8a8,
            activation=activation,
            global_num_experts=global_num_experts,
            expert_map=expert_map,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
            w1_zp=w1_zp,
            w2_zp=w2_zp,
            a1_scale=a1_scale,
            a2_scale=a2_scale,
            block_shape=block_shape,
            use_shuffle=use_weight_shuffle,
            routed_scaling_factor=routed_scaling_factor,
            solution_id=solution_id,
        )

    raise ValueError(f"Unknown solution_type: {moe_config.solution_type}")


def get_aiter_moe_config_w4a16(
    M: int,
    E: int,
    N1: int,
    N2: int,
    K: int,
    top_k: int,
    block_size: int,
    dtype: torch.dtype,
) -> Tuple[bool, AiterMoeConfig]:
    """Backward-compatible wrapper for w4a16 config lookup."""
    return get_aiter_moe_config(M, E, N1, N2, K, top_k, block_size, dtype, MoeQuantType.W4A16)


def aiter_moe_w4a16(
    hidden_states: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    moe_config: AiterMoeConfig,
    w1_scale: Optional[torch.Tensor] = None,
    w2_scale: Optional[torch.Tensor] = None,
    w1_zp: Optional[torch.Tensor] = None,
    w2_zp: Optional[torch.Tensor] = None,
    a1_scale: Optional[torch.Tensor] = None,
    a2_scale: Optional[torch.Tensor] = None,
    block_shape: Optional[list] = None,
    global_num_experts: int = -1,
    expert_map: Optional[torch.Tensor] = None,
    activation: str = "silu",
) -> torch.Tensor:
    """Backward-compatible wrapper for w4a16 execution."""
    return aiter_moe(
        hidden_states=hidden_states,
        w1=w1,
        w2=w2,
        topk_weights=topk_weights,
        topk_ids=topk_ids,
        moe_config=moe_config,
        activation=activation,
        w1_scale=w1_scale,
        w2_scale=w2_scale,
        w1_zp=w1_zp,
        w2_zp=w2_zp,
        a1_scale=a1_scale,
        a2_scale=a2_scale,
        block_shape=block_shape,
        global_num_experts=global_num_experts,
        expert_map=expert_map,
    )
