"""Inference-only grouped FP8 expert runtime for RivetCoder.

The Triton kernels in this module are adapted from Hugging Face's
``kernels-community/finegrained-fp8`` package (Apache-2.0).  The adaptation
adds per-output weight scales so separately quantized gate/up projections can
be concatenated without requantizing their checkpoint tensors.

This module is intentionally imported lazily.  Normal BF16 loading, training,
and CPU execution do not require Triton or TorchAO.
"""

from __future__ import annotations

import gc
import os
import shutil
import subprocess
import types
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F
from torch import nn


def ensure_windows_msvc_environment() -> str | None:
    """Populate the MSVC environment Triton's Windows launcher JIT needs."""

    if os.name != "nt":
        return None
    configured = os.environ.get("CC")
    if configured and (Path(configured).is_file() or shutil.which(configured)):
        return configured
    compiler = shutil.which("cl.exe")
    if compiler:
        os.environ["CC"] = compiler
        return compiler

    candidates = [
        Path(os.environ.get("ProgramFiles", r"C:\Program Files"))
        / "Microsoft Visual Studio"
        / "2022"
        / edition
        / "Common7"
        / "Tools"
        / "VsDevCmd.bat"
        for edition in ("Community", "Professional", "Enterprise", "BuildTools")
    ]
    vsdevcmd = next((path for path in candidates if path.is_file()), None)
    if vsdevcmd is None:
        raise RuntimeError(
            "Triton needs an MSVC C compiler on Windows. Install Visual Studio 2022 "
            "C++ Build Tools or launch the server from a Developer PowerShell."
        )

    command = f'call "{vsdevcmd}" -arch=x64 -host_arch=x64 >nul && set'
    completed = subprocess.run(
        command,
        check=True,
        capture_output=True,
        text=True,
        encoding="utf-8",
        errors="replace",
        shell=True,
        executable=os.environ.get("COMSPEC", "cmd.exe"),
    )
    for line in completed.stdout.splitlines():
        name, separator, value = line.partition("=")
        if separator and name:
            os.environ[name] = value
    compiler = shutil.which("cl.exe")
    if compiler is None:
        raise RuntimeError("VsDevCmd completed but cl.exe is still unavailable")
    os.environ["CC"] = compiler
    return compiler


def _load_triton() -> tuple[Any, Any, Any, Any]:
    ensure_windows_msvc_environment()
    try:
        import triton
        import triton.language as tl
        from torch.library import triton_op, wrap_triton
    except ImportError as error:
        raise RuntimeError(
            "Fast FP8 serving requires Triton. Use a PyTorch build that bundles Triton "
            "or install a Windows-compatible Triton package."
        ) from error
    return triton, tl, triton_op, wrap_triton


triton, tl, triton_op, wrap_triton = _load_triton()


@triton.jit
def _fp8_per_row_quant_kernel(x_ptr, q_ptr, scale_ptr, K: tl.constexpr):
    row = tl.program_id(axis=0)
    offsets = tl.arange(0, K)
    values = tl.load(x_ptr + row * K + offsets).to(tl.float32)
    scale = tl.maximum(tl.max(tl.abs(values), axis=0) / 448.0, 1.0e-12)
    quantized = (values / scale).to(tl.float8e4nv)
    tl.store(q_ptr + row * K + offsets, quantized)
    tl.store(scale_ptr + row, scale)


@triton_op("rivet_fp8::per_row_quant", mutates_args=())
def _fp8_per_row_quant(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    if x.ndim != 2 or not x.is_contiguous():
        raise ValueError("FP8 activation input must be a contiguous 2D tensor")
    if x.shape[1] <= 0 or x.shape[1] & (x.shape[1] - 1):
        raise ValueError("FP8 activation width must be a positive power of two")
    quantized = torch.empty_like(x, dtype=torch.float8_e4m3fn)
    scales = torch.empty(x.shape[0], device=x.device, dtype=torch.float32)
    wrap_triton(_fp8_per_row_quant_kernel)[(x.shape[0],)](
        x,
        quantized,
        scales,
        K=x.shape[1],
    )
    return quantized, scales


@triton.autotune(
    configs=[
        triton.Config({}, num_warps=warps, num_stages=stages)
        for warps in (2, 4, 8, 16)
        for stages in (2, 3, 4, 5)
    ],
    key=["N", "K", "BLOCK_M"],
)
@triton.jit
def _grouped_fp8_linear_kernel(
    A,
    B,
    C,
    AScales,
    BScales,
    Offsets,
    TileOffsets,
    S,
    N: tl.constexpr,
    K: tl.constexpr,
    stride_am,
    stride_ak,
    stride_be,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    stride_bs_e,
    stride_bs_n,
    NUM_EXPERTS: tl.constexpr,
    BLOCK_N: tl.constexpr,
    BLOCK_K: tl.constexpr,
    BLOCK_M: tl.constexpr,
    SEARCH_STEPS: tl.constexpr,
):
    tile_m = tl.program_id(axis=0)
    tile_n = tl.program_id(axis=1)
    total_tiles = tl.load(TileOffsets + NUM_EXPERTS - 1)
    if tile_m >= total_tiles:
        return

    low = 0
    high = NUM_EXPERTS
    for _ in tl.static_range(SEARCH_STEPS):
        middle = (low + high) >> 1
        middle_value = tl.load(TileOffsets + middle)
        move_right = middle_value <= tile_m
        low = tl.where(move_right, middle + 1, low)
        high = tl.where(move_right, high, middle)
    expert = low.to(tl.int64)

    previous = tl.maximum(expert - 1, 0)
    expert_start = tl.where(expert == 0, 0, tl.load(Offsets + previous))
    expert_end = tl.load(Offsets + expert)
    expert_rows = expert_end - expert_start
    expert_tile_start = tl.where(expert == 0, 0, tl.load(TileOffsets + previous))
    local_row_start = (tile_m - expert_tile_start) * BLOCK_M

    row_offsets = local_row_start + tl.arange(0, BLOCK_M)
    valid_rows = row_offsets < expert_rows
    global_rows = expert_start + row_offsets
    output_offsets = tile_n * BLOCK_N + tl.arange(0, BLOCK_N)
    k_offsets = tl.arange(0, BLOCK_K)

    a_ptrs = A + global_rows[:, None] * stride_am + k_offsets[None, :] * stride_ak
    b_ptrs = (
        B
        + expert * stride_be
        + output_offsets[None, :] * stride_bn
        + k_offsets[:, None] * stride_bk
    )
    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for _ in range(0, tl.cdiv(K, BLOCK_K)):
        a = tl.load(a_ptrs, mask=valid_rows[:, None], other=0.0)
        b = tl.load(b_ptrs)
        accumulator += tl.dot(a, b)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    activation_scale = tl.load(
        AScales + global_rows,
        mask=valid_rows,
        other=0.0,
    )
    weight_scale = tl.load(
        BScales + expert * stride_bs_e + output_offsets * stride_bs_n,
    )
    accumulator *= activation_scale[:, None] * weight_scale[None, :]

    if C.dtype.element_ty == tl.bfloat16:
        result = accumulator.to(tl.bfloat16)
    elif C.dtype.element_ty == tl.float16:
        result = accumulator.to(tl.float16)
    else:
        result = accumulator
    c_ptrs = C + global_rows[:, None] * stride_cm + output_offsets[None, :] * stride_cn
    tl.store(c_ptrs, result, mask=valid_rows[:, None])


@triton_op("rivet_fp8::grouped_linear", mutates_args=())
def _grouped_fp8_linear(
    activations: torch.Tensor,
    weights: torch.Tensor,
    weight_scales: torch.Tensor,
    offsets: torch.Tensor,
    tokens_per_expert: torch.Tensor,
) -> torch.Tensor:
    if activations.ndim != 2 or not activations.is_contiguous():
        raise ValueError("activations must be contiguous [routes, hidden]")
    if weights.ndim != 3 or not weights.is_contiguous():
        raise ValueError("weights must be contiguous [experts, output, hidden]")
    if weights.dtype != torch.float8_e4m3fn:
        raise TypeError("weights must use torch.float8_e4m3fn")
    experts, output_size, hidden_size = weights.shape
    if activations.shape[1] != hidden_size:
        raise ValueError("activation/weight hidden dimensions do not match")
    if output_size % 128 or hidden_size % 128:
        raise ValueError("grouped FP8 output and hidden dimensions must be divisible by 128")
    if weight_scales.shape != (experts, output_size):
        raise ValueError("weight_scales must have shape [experts, output]")
    if offsets.shape != (experts,) or tokens_per_expert.shape != (experts,):
        raise ValueError("offset/count tensors must have one value per expert")

    # TorchAO's checkpoint config uses dynamic PerTensor activation scaling.
    # Match that per routed expert (rather than per row) so prefill follows the
    # same quantization semantics as the original expert-by-expert calls.
    expert_ids = torch.repeat_interleave(
        torch.arange(experts, device=activations.device),
        tokens_per_expert.to(torch.long),
        output_size=activations.shape[0],
    )
    row_max = activations.abs().amax(dim=-1)
    expert_max = torch.zeros(experts, device=activations.device, dtype=activations.dtype)
    expert_max.scatter_reduce_(0, expert_ids, row_max, reduce="amax", include_self=True)
    expert_scales = (expert_max / 448.0).float().clamp_min(1.0e-12)
    activation_scales = expert_scales.index_select(0, expert_ids).contiguous()
    quantized = (
        activations.float()
        .div(activation_scales.unsqueeze(-1))
        .clamp(min=-448.0, max=448.0)
        .to(torch.float8_e4m3fn)
    )
    output = activations.new_empty((activations.shape[0], output_size))
    block_m = min(max(triton.next_power_of_2((activations.shape[0] + experts - 1) // experts), 16), 128)
    tiles_per_expert = (tokens_per_expert + block_m - 1) // block_m
    tile_offsets = torch.cumsum(tiles_per_expert, dim=0, dtype=torch.int32)
    max_m_tiles = triton.cdiv(activations.shape[0], block_m) + experts
    grid = (max_m_tiles, triton.cdiv(output_size, 128))
    wrap_triton(_grouped_fp8_linear_kernel)[grid](
        quantized,
        weights,
        output,
        activation_scales,
        weight_scales,
        offsets,
        tile_offsets,
        activations.shape[0],
        output_size,
        hidden_size,
        quantized.stride(0),
        quantized.stride(1),
        weights.stride(0),
        weights.stride(2),
        weights.stride(1),
        output.stride(0),
        output.stride(1),
        weight_scales.stride(0),
        weight_scales.stride(1),
        NUM_EXPERTS=experts,
        BLOCK_N=128,
        BLOCK_K=128,
        BLOCK_M=block_m,
        SEARCH_STEPS=experts.bit_length(),
    )
    return output


def grouped_fp8_linear(
    activations: torch.Tensor,
    weights: torch.Tensor,
    weight_scales: torch.Tensor,
    offsets: torch.Tensor,
    tokens_per_expert: torch.Tensor,
) -> torch.Tensor:
    return torch.ops.rivet_fp8.grouped_linear(
        activations,
        weights,
        weight_scales,
        offsets,
        tokens_per_expert,
    )


def _float8_parts(linear: nn.Linear) -> tuple[torch.Tensor, torch.Tensor]:
    weight = linear.weight
    qdata = getattr(weight, "qdata", None)
    scale = getattr(weight, "scale", None)
    if qdata is None or scale is None or "Float8" not in type(weight).__name__:
        raise TypeError("fast serving requires TorchAO Float8Tensor Linear weights")
    if qdata.dtype != torch.float8_e4m3fn or scale.numel() != 1:
        raise TypeError("fast serving currently supports per-tensor E4M3 TorchAO weights")
    return qdata, scale.reshape(())


def _direct_fp8_linear_forward(linear: nn.Linear, hidden_states: torch.Tensor) -> torch.Tensor:
    """TorchAO-compatible PerTensor FP8 Linear without tensor-subclass dispatch."""

    qdata, weight_scale = _float8_parts(linear)
    original_shape = hidden_states.shape
    flattened = hidden_states.reshape(-1, original_shape[-1]).contiguous()
    activation_scale = (flattened.abs().amax() / 448.0).float().reshape(1, 1)
    activation_scale = activation_scale.clamp_min(1.0e-12)
    quantized = (
        flattened.float()
        .div(activation_scale)
        .clamp(min=-448.0, max=448.0)
        .to(torch.float8_e4m3fn)
    )
    output = torch._scaled_mm(
        quantized,
        qdata.t(),
        activation_scale,
        weight_scale.reshape(1, 1),
        out_dtype=linear.weight.dtype,
        use_fast_accum=True,
    )
    if linear.bias is not None:
        output = output + linear.bias
    return output.reshape(*original_shape[:-1], linear.out_features)


def _install_direct_fp8_linears(model: nn.Module) -> int:
    installed = 0
    for module in model.modules():
        if not isinstance(module, nn.Linear) or getattr(module, "_rivet_direct_fp8", False):
            continue
        weight = module.weight
        if (
            "Float8" not in type(weight).__name__
            or getattr(weight, "qdata", None) is None
            or getattr(weight, "scale", None) is None
            or weight.scale.numel() != 1
        ):
            continue
        module.forward = types.MethodType(_direct_fp8_linear_forward, module)
        module._rivet_direct_fp8 = True
        installed += 1
    return installed


class PackedFp8ExpertBank(nn.Module):
    """One layer's 16 experts packed into two grouped FP8 projections."""

    def __init__(
        self,
        gate_up_qdata: torch.Tensor,
        gate_up_scales: torch.Tensor,
        down_qdata: torch.Tensor,
        down_scales: torch.Tensor,
        *,
        gate_clamp_max: float,
        up_clamp_min: float,
        up_clamp_max: float,
    ) -> None:
        super().__init__()
        self.register_buffer("gate_up_qdata", gate_up_qdata, persistent=False)
        self.register_buffer("gate_up_scales", gate_up_scales, persistent=False)
        self.register_buffer("down_qdata", down_qdata, persistent=False)
        self.register_buffer("down_scales", down_scales, persistent=False)
        self.num_experts = int(gate_up_qdata.shape[0])
        self.intermediate_size = int(gate_up_qdata.shape[1] // 2)
        self.hidden_size = int(gate_up_qdata.shape[2])
        self.gate_clamp_max = float(gate_clamp_max)
        self.up_clamp_min = float(up_clamp_min)
        self.up_clamp_max = float(up_clamp_max)

    @classmethod
    @torch.no_grad()
    def from_experts(
        cls,
        experts: nn.ModuleList,
        *,
        gate_clamp_max: float,
        up_clamp_min: float,
        up_clamp_max: float,
    ) -> "PackedFp8ExpertBank":
        if not experts:
            raise ValueError("cannot pack an empty expert list")
        first_gate, _ = _float8_parts(experts[0].gate_proj)
        first_down, _ = _float8_parts(experts[0].down_proj)
        num_experts = len(experts)
        intermediate_size, hidden_size = first_gate.shape
        if tuple(first_down.shape) != (hidden_size, intermediate_size):
            raise ValueError("unexpected down projection shape")
        device = first_gate.device
        gate_up_qdata = torch.empty(
            (num_experts, 2 * intermediate_size, hidden_size),
            device=device,
            dtype=torch.float8_e4m3fn,
        )
        gate_up_scales = torch.empty(
            (num_experts, 2 * intermediate_size), device=device, dtype=torch.float32
        )
        down_qdata = torch.empty(
            (num_experts, hidden_size, intermediate_size),
            device=device,
            dtype=torch.float8_e4m3fn,
        )
        down_scales = torch.empty((num_experts, hidden_size), device=device, dtype=torch.float32)

        for index, expert in enumerate(experts):
            gate_qdata, gate_scale = _float8_parts(expert.gate_proj)
            up_qdata, up_scale = _float8_parts(expert.up_proj)
            down_expert_qdata, down_scale = _float8_parts(expert.down_proj)
            if tuple(gate_qdata.shape) != (intermediate_size, hidden_size):
                raise ValueError("expert gate projection shapes are inconsistent")
            if tuple(up_qdata.shape) != (intermediate_size, hidden_size):
                raise ValueError("expert up projection shapes are inconsistent")
            if tuple(down_expert_qdata.shape) != (hidden_size, intermediate_size):
                raise ValueError("expert down projection shapes are inconsistent")
            gate_up_qdata[index, :intermediate_size].copy_(gate_qdata)
            gate_up_qdata[index, intermediate_size:].copy_(up_qdata)
            gate_up_scales[index, :intermediate_size].copy_(gate_scale.expand(intermediate_size))
            gate_up_scales[index, intermediate_size:].copy_(up_scale.expand(intermediate_size))
            down_qdata[index].copy_(down_expert_qdata)
            down_scales[index].copy_(down_scale.expand(hidden_size))

        return cls(
            gate_up_qdata,
            gate_up_scales,
            down_qdata,
            down_scales,
            gate_clamp_max=gate_clamp_max,
            up_clamp_min=up_clamp_min,
            up_clamp_max=up_clamp_max,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        selected_indices: torch.Tensor,
        selected_weights: torch.Tensor,
    ) -> torch.Tensor:
        token_count = hidden_states.shape[0]
        route_experts = selected_indices.reshape(-1)
        route_tokens = torch.arange(token_count, device=hidden_states.device).repeat_interleave(
            selected_indices.shape[-1]
        )
        order = torch.argsort(route_experts, stable=True)
        sorted_experts = route_experts.index_select(0, order)
        sorted_tokens = route_tokens.index_select(0, order)
        sorted_hidden = hidden_states.index_select(0, sorted_tokens).contiguous()
        sorted_route_weights = selected_weights.reshape(-1).index_select(0, order)
        tokens_per_expert = torch.bincount(
            sorted_experts, minlength=self.num_experts
        ).to(dtype=torch.int32)
        offsets = torch.cumsum(tokens_per_expert, dim=0, dtype=torch.int32)

        gate_up = grouped_fp8_linear(
            sorted_hidden,
            self.gate_up_qdata,
            self.gate_up_scales,
            offsets,
            tokens_per_expert,
        )
        gate, up = gate_up.split(self.intermediate_size, dim=-1)
        intermediate = F.silu(gate.clamp(max=self.gate_clamp_max)) * up.clamp(
            min=self.up_clamp_min,
            max=self.up_clamp_max,
        )
        routed = grouped_fp8_linear(
            intermediate.contiguous(),
            self.down_qdata,
            self.down_scales,
            offsets,
            tokens_per_expert,
        )
        weighted = routed * sorted_route_weights.to(routed.dtype).unsqueeze(-1)
        # ``order`` is a permutation, so restoring route order and reducing a
        # contiguous [token, top_k, hidden] view avoids duplicate-index atomics.
        # This makes decode deterministic and is friendlier to CUDA graphs.
        inverse_order = torch.argsort(order)
        route_outputs = weighted.index_select(0, inverse_order).reshape(
            token_count, selected_indices.shape[-1], self.hidden_size
        )
        # The reference dispatcher visits experts in ascending expert-index
        # order.  Preserve that BF16 accumulation order to minimize long-stack
        # drift across 30 residual layers.
        expert_order = torch.argsort(selected_indices, dim=-1, stable=True)
        route_outputs = route_outputs.gather(
            1,
            expert_order.unsqueeze(-1).expand(-1, -1, self.hidden_size),
        )
        output = torch.zeros_like(hidden_states)
        for route_slot in range(selected_indices.shape[-1]):
            output = output + route_outputs[:, route_slot]
        return output


@torch.no_grad()
def install_fast_fp8_runtime(model: nn.Module) -> dict[str, Any]:
    """Pack every expert layer and enable the inference-only fast path.

    The transformation releases the original per-expert modules to avoid
    duplicating their FP8 storage.  It is intentionally one-way for the current
    process; reload the checkpoint to recover trainable/module-list form.
    """

    if model.training:
        raise RuntimeError("call model.eval() before enabling fast FP8 serving")
    if not torch.cuda.is_available():
        raise RuntimeError("fast FP8 serving requires CUDA")
    ensure_windows_msvc_environment()

    wrappers = tuple(model.fusion_layers())
    packed_layers = 0
    released_experts = 0
    packed_bytes = 0
    for wrapper in wrappers:
        if getattr(wrapper, "fast_expert_bank", None) is not None:
            continue
        bank = PackedFp8ExpertBank.from_experts(
            wrapper.experts,
            gate_clamp_max=wrapper.experts[0].gate_clamp_max,
            up_clamp_min=wrapper.experts[0].up_clamp_min,
            up_clamp_max=wrapper.experts[0].up_clamp_max,
        )
        released_experts += len(wrapper.experts)
        packed_bytes += sum(buffer.numel() * buffer.element_size() for buffer in bank.buffers())
        wrapper.fast_expert_bank = bank
        wrapper.experts = nn.ModuleList()
        wrapper.serving_mode = True
        wrapper.last_router_state = None
        wrapper.last_router_diagnostics = None
        packed_layers += 1
    # TorchAO tensor subclasses can participate in reference cycles.  Explicit
    # collection is necessary before the allocator can release the unpacked
    # per-expert qdata that the packed banks replaced.
    gc.collect()
    torch.cuda.empty_cache()
    direct_fp8_linears = _install_direct_fp8_linears(model)

    return {
        "backend": "triton-grouped-fp8",
        "packed_layers": packed_layers,
        "released_experts": released_experts,
        "packed_bytes": packed_bytes,
        "direct_fp8_linears": direct_fp8_linears,
        "cuda_allocated_bytes": torch.cuda.memory_allocated(),
        "compiler": os.environ.get("CC"),
    }