#!/usr/bin/env python3
"""
检测模型的head_size，如果不在vLLM支持列表中，则patch到vllm代码中
"""
import os
import json
import re
import shutil
from typing import Optional, List

# vLLM默认支持的head_size列表
DEFAULT_SUPPORTED_HEAD_SIZES = [64, 80, 96, 112, 120, 128, 192, 256]

# vLLM需要patch的文件路径
VLLM_PAGED_ATTN_PATH = "/usr/local/corex/lib64/python3/dist-packages/vllm/attention/ops/paged_attn.py"


def get_model_head_size(model_dir: str) -> Optional[int]:
    """
    从模型config.json中获取head_size

    Args:
        model_dir: 模型目录路径

    Returns:
        head_size，如果无法获取则返回None
    """
    config_path = os.path.join(model_dir, "config.json")

    if not os.path.exists(config_path):
        print(f"[detect_head_size] config.json not found in {model_dir}")
        return None

    try:
        with open(config_path, 'r') as f:
            config = json.load(f)

        # 尝试多种方式获取head_size
        # 1. 直接从head_dim字段获取
        if 'head_dim' in config:
            head_size = config['head_dim']
            print(f"[detect_head_size] Found head_dim in config: {head_size}")
            return head_size

        # 2. 从hidden_size和num_attention_heads计算
        if 'hidden_size' in config and 'num_attention_heads' in config:
            hidden_size = config['hidden_size']
            num_heads = config['num_attention_heads']
            head_size = hidden_size // num_heads
            print(f"[detect_head_size] Calculated from hidden_size({hidden_size}) / num_attention_heads({num_heads}) = {head_size}")
            return head_size

        # 3. 对于GPTJ等模型，可能使用n_embd和n_head
        if 'n_embd' in config and 'n_head' in config:
            n_embd = config['n_embd']
            n_head = config['n_head']
            head_size = n_embd // n_head
            print(f"[detect_head_size] Calculated from n_embd({n_embd}) / n_head({n_head}) = {head_size}")
            return head_size

        # 4. 对于T5等模型，可能使用d_kv和num_heads
        if 'd_kv' in config:
            head_size = config['d_kv']
            print(f"[detect_head_size] Found d_kv in config: {head_size}")
            return head_size

        print(f"[detect_head_size] Cannot determine head_size from config.json")
        return None

    except Exception as e:
        print(f"[detect_head_size] Error reading config.json: {e}")
        return None


def patch_vllm_head_sizes(new_head_size: int) -> bool:
    """
    将新的head_size patch到vLLM的get_supported_head_sizes方法中

    Args:
        new_head_size: 需要添加的新head_size

    Returns:
        是否成功patch
    """
    if not os.path.exists(VLLM_PAGED_ATTN_PATH):
        print(f"[patch] vLLM file not found: {VLLM_PAGED_ATTN_PATH}")
        return False

    try:
        # 备份原文件
        backup_path = VLLM_PAGED_ATTN_PATH + ".backup"
        if not os.path.exists(backup_path):
            shutil.copy2(VLLM_PAGED_ATTN_PATH, backup_path)
            print(f"[patch] Backed up original file to {backup_path}")

        # 读取文件内容
        with open(VLLM_PAGED_ATTN_PATH, 'r') as f:
            content = f.read()

        # 找到get_supported_head_sizes方法并patch
        pattern = r'(@staticmethod\s+def get_supported_head_sizes\(\)\s*->\s*List\[int\]:\s*return\s*\[)([^\]]+)(\])'
        match = re.search(pattern, content, re.MULTILINE | re.DOTALL)

        if not match:
            print(f"[patch] Could not find get_supported_head_sizes method in {VLLM_PAGED_ATTN_PATH}")
            return False

        prefix = match.group(1)
        current_list = match.group(2)
        suffix = match.group(3)

        # 解析当前列表
        current_sizes = [int(x.strip()) for x in current_list.split(',') if x.strip()]

        # 添加新的head_size并排序
        if new_head_size not in current_sizes:
            current_sizes.append(new_head_size)
            current_sizes.sort()

            # 生成新的列表字符串
            new_list_str = ', '.join(str(size) for size in current_sizes)

            # 替换原方法
            new_method = f"{prefix}{new_list_str}{suffix}"
            new_content = re.sub(pattern, new_method, content, flags=re.MULTILINE | re.DOTALL)

            # 写回文件
            with open(VLLM_PAGED_ATTN_PATH, 'w') as f:
                f.write(new_content)

            print(f"[patch] Successfully added head_size {new_head_size} to supported list: {current_sizes}")
            return True
        else:
            print(f"[patch] head_size {new_head_size} already in list, skipping")
            return True

    except Exception as e:
        print(f"[patch] Error patching vLLM file: {e}")
        return False


def main():
    """主函数"""
    try:
        model_dir = os.environ.get("MODEL_DIR", "/model")

        print(f"[detect_head_size] Checking model in: {model_dir}")

        # 获取head_size
        head_size = get_model_head_size(model_dir)

        if head_size is None:
            print(f"[detect_head_size] Could not determine head_size, skipping")
            return

        print(f"[detect_head_size] Model head_size: {head_size}")

        # 检查是否需要patch
        if head_size in DEFAULT_SUPPORTED_HEAD_SIZES:
            print(f"[detect_head_size] head_size {head_size} is already supported by vLLM, skipping patch")
            return

        print(f"[detect_head_size] head_size {head_size} is NOT in default supported list: {DEFAULT_SUPPORTED_HEAD_SIZES}")
        print(f"[detect_head_size] Attempting to patch vLLM...")

        # 执行patch
        success = patch_vllm_head_sizes(head_size)

        if success:
            print(f"[detect_head_size] Successfully patched vLLM to support head_size {head_size}")
        else:
            print(f"[detect_head_size] Failed to patch vLLM for head_size {head_size}")

    except Exception as e:
        print(f"[detect_head_size] Error during head_size detection/patch: {e}")
        print(f"[detect_head_size] Continuing with vLLM startup anyway...")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"[detect_head_size] Unexpected error: {e}")
        print(f"[detect_head_size] This will not prevent vLLM from starting")
        import sys
        sys.exit(0)  # 总是成功退出，不阻止vLLM启动