#!/usr/bin/env python3
"""Patch ROCm DFlash support into installed vLLM and AITER packages.

The reachable host uses newer `vllm/v1` attention paths than the older
one-off patch script from the learnings. This script applies the same logical
fixes against the actual installed package files inside the runtime container.
It is intentionally idempotent.
"""

from __future__ import annotations

import importlib.util
import re
import sys
from pathlib import Path


def locate_module_file(module_name: str) -> Path:
    spec = importlib.util.find_spec(module_name)
    if spec is None or spec.origin is None:
        raise RuntimeError(f"Could not locate module: {module_name}")
    return Path(spec.origin).resolve()


def first_existing(paths: list[Path]) -> Path:
    for path in paths:
        if path.exists():
            return path
    raise RuntimeError("Could not locate any expected path:\n" + "\n".join(map(str, paths)))


def replace_once(text: str, old: str, new: str, path: Path) -> str:
    if new in text:
        return text
    if old not in text:
        raise RuntimeError(f"Pattern not found in {path}: {old[:120]!r}")
    return text.replace(old, new, 1)


def replace_all_regex(
    text: str,
    pattern: str,
    repl: str,
    path: Path,
    *,
    min_count: int = 1,
) -> str:
    compiled = re.compile(pattern, re.MULTILINE)
    matches = list(compiled.finditer(text))
    marker = re.sub(r"\\(?:g<[^>]+>|[1-9][0-9]*)", "", repl)
    if not matches:
        if repl in text or (marker and marker in text):
            return text
        raise RuntimeError(f"Regex pattern not found in {path}: {pattern}")
    if len(matches) < min_count:
        updated = compiled.sub(repl, text)
        if updated != text and marker and marker in updated:
            return updated
        raise RuntimeError(
            f"Expected at least {min_count} matches in {path}, found {len(matches)}"
        )
    return compiled.sub(repl, text)


def patch_file(path: Path, transform) -> None:
    original = path.read_text()
    updated = transform(original, path)
    if updated == original:
        print(f"[skip] {path}")
        return
    path.write_text(updated)
    print(f"[patch] {path}")


def patch_rocm_aiter_fa(text: str, path: Path) -> str:
    text = replace_once(
        text,
        '    @staticmethod\n    def get_name() -> str:\n        return "FLASH_ATTN"\n\n    @staticmethod\n    def get_impl_cls() -> type["AiterFlashAttentionImpl"]:\n',
        '    @staticmethod\n    def get_name() -> str:\n        return "FLASH_ATTN"\n\n    @classmethod\n    def supports_non_causal(cls) -> bool:\n        return True\n\n    @staticmethod\n    def get_impl_cls() -> type["AiterFlashAttentionImpl"]:\n',
        path,
    )
    # Ensure AiterFlashAttentionMetadata dataclass has causal: bool field
    if "    causal: bool" not in text.split("class AiterFlashAttentionMetadata")[1].split("def ")[0]:
        text = replace_once(
            text,
            "class AiterFlashAttentionMetadata:\n",
            "class AiterFlashAttentionMetadata:\n    causal: bool\n",
            path,
        )

    # Ensure AiterFlashAttentionMetadata() constructors have causal= kwarg.
    # Use a robust approach: find each constructor call, ensure causal= appears
    # exactly once (first arg). This handles all upstream variations.
    def _ensure_causal_in_constructor(m: re.Match) -> str:
        call_text = m.group(0)
        if re.search(r'\bcausal=', call_text):
            return call_text  # already has causal=, leave it
        # Insert causal= as first kwarg after the opening paren
        return call_text.replace(
            "AiterFlashAttentionMetadata(\n",
            "AiterFlashAttentionMetadata(\n            causal=common_attn_metadata.causal,\n",
        )

    text = re.sub(
        r'(?:attn_metadata = |return )AiterFlashAttentionMetadata\([^)]*?\)',
        _ensure_causal_in_constructor,
        text,
        flags=re.DOTALL,
    )

    # Replace hardcoded causal=True with dynamic causal=attn_metadata.causal
    # in flash attention calls. Skip if upstream already uses a dynamic causal.
    text = re.sub(
        r"(softmax_scale=self\.scale,\n)(\s*)causal=True,",
        r"\1\2causal=attn_metadata.causal,",
        text,
        flags=re.MULTILINE,
    )

    # Safety: remove any duplicate causal= kwargs that may have been created
    # by patching an upstream that already had causal= in constructors.
    # This finds lines like `causal=...,` that appear more than once in the
    # same parenthesized call and removes the extras (keeps first occurrence).
    def _dedup_causal(m: re.Match) -> str:
        block = m.group(0)
        causal_lines = [i for i, line in enumerate(block.split('\n')) if re.match(r'\s*causal=', line)]
        if len(causal_lines) <= 1:
            return block
        lines = block.split('\n')
        # Keep only the first causal= line
        seen = False
        result = []
        for line in lines:
            if re.match(r'\s*causal=', line):
                if seen:
                    continue  # drop duplicate
                seen = True
            result.append(line)
        return '\n'.join(result)

    text = re.sub(
        r'AiterFlashAttentionMetadata\([^)]*\)',
        _dedup_causal,
        text,
        flags=re.DOTALL,
    )

    return text


def patch_supports_non_causal(text: str, path: Path, backend_name: str) -> str:
    insertion = (
        f'    @staticmethod\n    def get_name() -> str:\n        return "{backend_name}"\n\n'
        "    @classmethod\n    def supports_non_causal(cls) -> bool:\n        return True\n\n"
    )
    current = f'    @staticmethod\n    def get_name() -> str:\n        return "{backend_name}"\n\n'
    if "def supports_non_causal" in text:
        return text
    if current not in text:
        raise RuntimeError(f"Could not find get_name block in {path}")
    return text.replace(current, insertion, 1)


def patch_aiter_wrapper(text: str, path: Path) -> str:
    if "IS_CAUSAL=causal" in text:
        return text.replace(
            '    assert causal, "Only causal attention is supported"\n',
            "",
        )
    text = replace_once(
        text,
        '    assert causal, "Only causal attention is supported"\n',
        "",
        path,
    )
    text = replace_all_regex(
        text,
        r"(ALL_DECODE=ALL_DECODE,\n)(\s*)(\*\*attn_config,)",
        r"\1\2IS_CAUSAL=causal,\n\2\3",
        path,
        min_count=1,
    )
    text = replace_all_regex(
        text,
        r"(ALL_DECODE=ALL_DECODE,\n)(\s*)(\*\*config,)",
        r"\1\2IS_CAUSAL=causal,\n\2\3",
        path,
        min_count=1,
    )
    return text


def patch_aiter_kernel(text: str, path: Path) -> str:
    if "IS_CAUSAL: tl.constexpr = True" in text:
        text = text.replace(
            "num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE)",
            "num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE) if IS_CAUSAL else cdiv_fn(seq_len, TILE_SIZE)",
        )
        text = text.replace(
            "seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1",
            "seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1 if IS_CAUSAL else seq_offset[None, :] < seq_len",
        )
        return text
    text = replace_all_regex(
        text,
        r"(ALL_DECODE: tl\.constexpr = False,  # bool\n)(\):)",
        r"\1    IS_CAUSAL: tl.constexpr = True,  # bool\n\2",
        path,
        min_count=2,
    )
    text = replace_all_regex(
        text,
        r"num_tiles = cdiv_fn\(max_seq_prefix_len, TILE_SIZE\)",
        "num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE) if IS_CAUSAL else cdiv_fn(seq_len, TILE_SIZE)",
        path,
        min_count=2,
    )
    text = replace_all_regex(
        text,
        r"seq_mask = seq_offset\[None, :\] < context_len \+ query_pos\[:, None\] \+ 1",
        "seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1 if IS_CAUSAL else seq_offset[None, :] < seq_len",
        path,
        min_count=2,
    )
    return text


def patch_vllm_triton_unified_attention(text: str, path: Path) -> str:
    if "IS_CAUSAL=causal" in text:
        return text.replace(
            '    assert causal, "Only causal attention is supported"\n',
            "",
        )
    text = replace_once(
        text,
        '    assert causal, "Only causal attention is supported"\n',
        "",
        path,
    )
    text = replace_all_regex(
        text,
        r"num_tiles = cdiv_fn\(max_seq_prefix_len, TILE_SIZE\)",
        "num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE) if IS_CAUSAL else cdiv_fn(seq_len, TILE_SIZE)",
        path,
        min_count=2,
    )
    text = replace_all_regex(
        text,
        r"seq_mask = seq_offset\[None, :\] <= query_abs_pos",
        "seq_mask = seq_offset[None, :] <= query_abs_pos if IS_CAUSAL else seq_offset[None, :] < seq_len",
        path,
        min_count=2,
    )
    text = replace_all_regex(
        text,
        r"USE_FP8: tl\.constexpr,  # bool",
        "USE_FP8: tl.constexpr,  # bool\n    IS_CAUSAL: tl.constexpr = True,  # bool",
        path,
        min_count=2,
    )
    text = replace_all_regex(
        text,
        r"(BLOCK_M=BLOCK_M,\n)(\s*)",
        r"\1\2IS_CAUSAL=causal,\n\2",
        path,
        min_count=2,
    )
    return text


def patch_vllm_dflash(text: str, path: Path) -> str:
    return replace_once(
        text,
        '            assert getattr(attn_metadata, "causal", None) is False, (\n',
        '            assert getattr(attn_metadata, "causal", None) in (True, False, None), (\n',
        path,
    )


def patch_vllm_selector(text: str, path: Path) -> str:
    if "AttentionBackendEnum.TRITON_MLA" in text:
        return text
    text = replace_once(
        text,
        "from vllm.v1.attention.backends.registry import (\n"
        "    MAMBA_TYPE_TO_BACKEND_MAP,\n"
        "    MambaAttentionBackendEnum,\n"
        ")\n",
        "from vllm.v1.attention.backends.registry import (\n"
        "    AttentionBackendEnum,\n"
        "    MAMBA_TYPE_TO_BACKEND_MAP,\n"
        "    MambaAttentionBackendEnum,\n"
        ")\n",
        path,
    )
    new = (
        "    speculative_config = vllm_config.speculative_config\n"
        "    hf_config = vllm_config.model_config.hf_config\n"
        "    architectures = list(getattr(hf_config, \"architectures\", []) or [])\n"
        "    is_dflash_draft = any(\n"
        "        str(arch).startswith(\"DFlash\") for arch in architectures\n"
        "    )\n"
        "    use_non_causal = (\n"
        "        speculative_config is not None\n"
        "        and speculative_config.method == \"dflash\"\n"
        "        and is_dflash_draft\n"
        "    )\n"
        "\n"
        "    backend = vllm_config.attention_config.backend\n"
        "    if (\n"
        "        speculative_config is not None\n"
        "        and speculative_config.method == \"dflash\"\n"
        "        and use_mla\n"
        "        and not is_dflash_draft\n"
        "        and backend is None\n"
        "    ):\n"
        "        backend = AttentionBackendEnum.TRITON_MLA\n"
    )
    old_variants = [
        (
        "    speculative_config = vllm_config.speculative_config\n"
        "    use_non_causal = (\n"
        "        speculative_config is not None and speculative_config.method == \"dflash\"\n"
        "    )\n"
        ),
        (
        "    speculative_config = vllm_config.speculative_config\n"
        "    hf_config = vllm_config.model_config.hf_config\n"
        "    architectures = list(getattr(hf_config, \"architectures\", []) or [])\n"
        "    use_non_causal = (\n"
        "        speculative_config is not None\n"
        "        and speculative_config.method == \"dflash\"\n"
        "        and any(str(arch).startswith(\"DFlash\") for arch in architectures)\n"
        "    )\n"
        ),
    ]
    for old in old_variants:
        if old in text:
            text = text.replace(old, new, 1)
            break
    else:
        if new not in text:
            raise RuntimeError(f"Could not find selector speculative block in {path}")
    return replace_once(
        text,
        "        backend=vllm_config.attention_config.backend,\n",
        "        backend=backend,\n",
        path,
    )


def patch_mla_dynamic_fp8_scales(text: str, path: Path) -> str:
    """Patch MLA do_kv_cache_update to compute FP8 scales from actual data.

    Without calibrated scales, FP8 KV cache clips values and produces garbage.
    This patch computes the scale dynamically from each batch of KV data
    before writing to cache, ensuring the full FP8 E4M3 range is used.
    """
    marker = "# [DFlash patch] dynamic FP8 scale"
    if marker in text:
        return text

    old = (
        "    def do_kv_cache_update(\n"
        "        self,\n"
        "        kv_c_normed: torch.Tensor,\n"
        "        k_pe: torch.Tensor,\n"
        "        kv_cache: torch.Tensor,\n"
        "        slot_mapping: torch.Tensor,\n"
        "        kv_cache_dtype: str,\n"
        "        k_scale: torch.Tensor,\n"
        "    ) -> None:\n"
        "        if kv_cache.numel() == 0:\n"
        "            return\n"
        "        from vllm import _custom_ops as ops\n"
        "\n"
        "        ops.concat_and_cache_mla(\n"
        "            kv_c_normed,\n"
        "            k_pe.squeeze(1),\n"
        "            kv_cache,\n"
        "            slot_mapping.flatten(),\n"
        "            kv_cache_dtype=kv_cache_dtype,\n"
        "            scale=k_scale,\n"
        "        )"
    )

    new = (
        "    _fp8_call_count = 0\n"
        "    _fp8_scale_log = {}\n"
        "\n"
        "    def do_kv_cache_update(\n"
        "        self,\n"
        "        kv_c_normed: torch.Tensor,\n"
        "        k_pe: torch.Tensor,\n"
        "        kv_cache: torch.Tensor,\n"
        "        slot_mapping: torch.Tensor,\n"
        "        kv_cache_dtype: str,\n"
        "        k_scale: torch.Tensor,\n"
        "    ) -> None:\n"
        "        if kv_cache.numel() == 0:\n"
        "            return\n"
        "        from vllm import _custom_ops as ops\n"
        "\n"
        "        " + marker + "\n"
        "        if kv_cache_dtype in ('fp8', 'fp8_e4m3', 'fp8_e5m2'):\n"
        "            import os as _os, json as _json\n"
        "            import torch as _torch\n"
        "            _fp8_max = _torch.finfo(_torch.float8_e4m3fnuz).max\n"
        "            _abs_max = max(\n"
        "                kv_c_normed.abs().max().item(),\n"
        "                k_pe.abs().max().item(),\n"
        "                1e-12,\n"
        "            )\n"
        "            _scale = _abs_max / _fp8_max\n"
        "            _static_file = _os.environ.get('FP8_KV_SCALES_FILE')\n"
        "            if _static_file and _os.path.exists(_static_file):\n"
        "                if not hasattr(self, '_fp8_static_scales'):\n"
        "                    self._fp8_static_scales = _json.load(open(_static_file))\n"
        "                _s = self._fp8_static_scales.get('global_scale', _scale)\n"
        "                k_scale.fill_(_s)\n"
        "            else:\n"
        "                k_scale.copy_(_torch.max(k_scale, _torch.tensor(_scale, device=k_scale.device)))\n"
        "                type(self)._fp8_call_count += 1\n"
        "                _cur = k_scale.item()\n"
        "                type(self)._fp8_scale_log[id(k_scale)] = _cur\n"
        "                if type(self)._fp8_call_count % 5000 == 0:\n"
        "                    _vals = list(type(self)._fp8_scale_log.values())\n"
        "                    _dump = {'global_scale': max(_vals), 'num_layers': len(_vals),\n"
        "                             'min_scale': min(_vals), 'max_scale': max(_vals),\n"
        "                             'call_count': type(self)._fp8_call_count}\n"
        "                    _out = '/tmp/fp8_kv_scales.json'\n"
        "                    with open(_out, 'w') as _f:\n"
        "                        _json.dump(_dump, _f, indent=2)\n"
        "\n"
        "        ops.concat_and_cache_mla(\n"
        "            kv_c_normed,\n"
        "            k_pe.squeeze(1),\n"
        "            kv_cache,\n"
        "            slot_mapping.flatten(),\n"
        "            kv_cache_dtype=kv_cache_dtype,\n"
        "            scale=k_scale,\n"
        "        )"
    )

    if old not in text:
        if marker in text:
            return text
        raise RuntimeError(f"Could not find do_kv_cache_update in {path}")
    return text.replace(old, new, 1)


def main() -> int:
    vllm_root = locate_module_file("vllm").parent
    site_packages = vllm_root.parent

    rocm_aiter_fa = vllm_root / "v1" / "attention" / "backends" / "rocm_aiter_fa.py"
    rocm_attn = vllm_root / "v1" / "attention" / "backends" / "rocm_attn.py"
    rocm_aiter_unified = (
        vllm_root / "v1" / "attention" / "backends" / "rocm_aiter_unified_attn.py"
    )
    triton_attn = vllm_root / "v1" / "attention" / "backends" / "triton_attn.py"
    selector_path = vllm_root / "v1" / "attention" / "selector.py"
    vllm_triton_ops = (
        vllm_root / "v1" / "attention" / "ops" / "triton_unified_attention.py"
    )
    vllm_dflash = vllm_root / "v1" / "spec_decode" / "dflash.py"
    aiter_wrapper = first_existing(
        [
            site_packages / "aiter" / "ops" / "triton" / "unified_attention.py",
            site_packages
            / "aiter"
            / "ops"
            / "triton"
            / "attention"
            / "unified_attention.py",
        ]
    )
    aiter_kernel = first_existing(
        [
            site_packages
            / "aiter"
            / "ops"
            / "triton"
            / "_triton_kernels"
            / "unified_attention.py",
            site_packages
            / "aiter"
            / "ops"
            / "triton"
            / "_triton_kernels"
            / "attention"
            / "unified_attention.py",
        ]
    )

    mla_backend = vllm_root / "v1" / "attention" / "backend.py"
    patch_file(mla_backend, patch_mla_dynamic_fp8_scales)
    patch_file(rocm_aiter_fa, patch_rocm_aiter_fa)
    patch_file(
        rocm_attn,
        lambda text, path: patch_supports_non_causal(text, path, "ROCM_ATTN"),
    )
    patch_file(
        rocm_aiter_unified,
        lambda text, path: patch_supports_non_causal(
            text, path, "ROCM_AITER_UNIFIED_ATTN"
        ),
    )
    patch_file(
        triton_attn,
        lambda text, path: patch_supports_non_causal(text, path, "TRITON_ATTN"),
    )
    patch_file(selector_path, patch_vllm_selector)
    patch_file(vllm_triton_ops, patch_vllm_triton_unified_attention)
    patch_file(vllm_dflash, patch_vllm_dflash)
    patch_file(aiter_wrapper, patch_aiter_wrapper)
    patch_file(aiter_kernel, patch_aiter_kernel)
    print("[done] ROCm DFlash patch applied")
    return 0


if __name__ == "__main__":
    sys.exit(main())