#!/usr/bin/env python3
"""
修复 CoreX 版 vLLM 0.17.0 的 Llama 4 chunked local attention bug。

症状（第一个真实请求时）：
  TypeError: CommonAttentionMetadata.__init__() missing 2 required positional
  arguments: 'key_start_loc' and 'seq_lens_np'

原因：
  vllm/v1/attention/backends/utils.py 的
  make_local_attention_virtual_batches() 在构造 CommonAttentionMetadata 时
  没有传新版本 dataclass 要求的 key_start_loc / seq_lens_np。

本脚本幂等：已经打过补丁会直接返回 0。
"""

import os
import sys

try:
    import vllm
except ImportError as e:
    print(f"ERROR: 无法 import vllm: {e}", file=sys.stderr)
    sys.exit(1)

TARGET = os.path.join(
    os.path.dirname(vllm.__file__),
    "v1",
    "attention",
    "backends",
    "utils.py",
)

MARKER = "seq_lens_np=seqlens_k_local"

OLD = """        query_start_loc=query_start_loc_cpu.to(device=device, non_blocking=True),
        seq_lens=seq_lens_cpu.to(device=device, non_blocking=True),"""

NEW = """        query_start_loc=query_start_loc_cpu.to(device=device, non_blocking=True),
        key_start_loc=torch.from_numpy(
            np.concatenate(
                (np.zeros(1, dtype=np.int32), np.cumsum(seqlens_k_local, dtype=np.int32))
            )
        ).to(device=device, non_blocking=True),
        seq_lens=seq_lens_cpu.to(device=device, non_blocking=True),
        seq_lens_np=seqlens_k_local,"""


def main() -> int:
    if not os.path.exists(TARGET):
        print(f"ERROR: 找不到目标文件: {TARGET}", file=sys.stderr)
        return 1

    with open(TARGET, "r", encoding="utf-8") as f:
        src = f.read()

    if MARKER in src:
        print(f"patch already applied: {TARGET}")
        return 0

    if OLD not in src:
        print(
            "ERROR: 目标代码块未找到，vLLM 版本可能不同。\n"
            f"请手动检查: {TARGET}",
            file=sys.stderr,
        )
        return 1

    with open(TARGET, "w", encoding="utf-8") as f:
        f.write(src.replace(OLD, NEW, 1))

    print(f"patched: {TARGET}")
    return 0


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