import functools
import os

import torch
from bench_tirx_prenorm import build_kernel, compare

VALIDATE = os.getenv("MAGI_TIRX_PRENORM_VALIDATE", "0") == "1"
checked = set()
reported = set()


@functools.lru_cache(maxsize=1)
def executable():
    return build_kernel()


def apply(norm, x, dispatcher):
    def fallback():
        return norm(x, modality_dispatcher=dispatcher)

    if (
        torch.is_grad_enabled()
        or norm.training
        or not x.is_cuda
        or x.dtype != torch.bfloat16
        or x.ndim != 2
        or x.shape[1] != 5120
        or not x.is_contiguous()
        or norm.eps != 1e-6
        or norm.num_modality not in (1, 3)
        or norm.weight.dtype != torch.bfloat16
        or norm.weight.device != x.device
        or not norm.weight.is_contiguous()
        or norm.weight.numel() != norm.num_modality * 5120
    ):
        return fallback()

    groups = (
        [x.shape[0]]
        if norm.num_modality == 1
        else [int(n) for n in dispatcher.group_size_cpu]
    )
    if (
        len(groups) != norm.num_modality
        or any(n < 0 for n in groups)
        or sum(groups) != x.shape[0]
        or x.shape[0] == 0
    ):
        return fallback()

    weight = norm.weight
    version = None if torch.is_inference(weight) else weight._version
    cached = getattr(norm, "_tirx_prenorm_scale_cache", None)
    if (
        cached is None
        or cached[0] is not weight
        or cached[1] != version
    ):
        # 保留BF16 weight+1的舍入，固定权重推理中复用。
        scale = (weight.detach() + 1).reshape(
            norm.num_modality, 5120
        ).contiguous()
        norm._tirx_prenorm_scale_cache = (weight, version, scale)
    else:
        scale = cached[2]

    output = torch.empty_like(x)
    ex = executable()
    start = 0
    for i, size in enumerate(groups):
        if size:
            ex(
                x[start:start + size].reshape(-1),
                scale[i],
                output[start:start + size].reshape(-1),
                size,
            )
        start += size

    key = (norm.num_modality, x.device.index)
    if key not in reported:
        print(
            f"TIRX_PRENORM_ACTIVE shape={tuple(x.shape)} "
            f"groups={groups}",
            flush=True,
        )
        reported.add(key)

    if VALIDATE and key not in checked:
        print(
            f"TIRX_PRENORM_CHECK shape={tuple(x.shape)} "
            f"groups={groups}",
            flush=True,
        )
        expected = fallback()
        compare(output, expected)
        checked.add(key)
        print("TIRX_PRENORM_CHECK: 真实输入误差检查通过", flush=True)

    return output


def transform_source(source):
    old = (
        "self.pre_norm(hidden_states, "
        "modality_dispatcher=modality_dispatcher)"
    )
    if source.count(old) != 1:
        raise RuntimeError("Attention pre_norm入口不匹配，未安装补丁")
    return source.replace(
        old,
        "_tirx_prenorm(self.pre_norm, hidden_states, modality_dispatcher)",
        1,
    )


@torch.inference_mode()
def prewarm():
    if os.getenv("MAGI_TIRX_PREWARM", "0") != "1":
        return
    print("TIRX_PREWARM_PRENORM: 开始", flush=True)
    ex = executable()
    x = torch.zeros(8 * 5120, device="cuda", dtype=torch.bfloat16)
    w = torch.ones(5120, device="cuda", dtype=torch.bfloat16)
    y = torch.empty_like(x)
    for _ in range(3):
        ex(x, w, y, 8)
    torch.cuda.synchronize()
    print("TIRX_PREWARM_PRENORM: 完成", flush=True)