#!/usr/bin/env python3
"""把 R2 chat_template 与 R4 architecture 白名单应用到模型目录。

不访问平台，不启动 vLLM。
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

FORBIDDEN_FRAMEWORK = "vllm_fix_tokenizer"
TOKENIZER_PATCH_KEY = "modelhub_tokenizer_patch"


class PatchError(Exception):
    """模型目录不满足本引擎的修补前提。"""


def engine_root() -> Path:
    return Path(__file__).resolve().parents[1]


def assert_framework(framework: str) -> None:
    if not framework or framework == FORBIDDEN_FRAMEWORK:
        raise PatchError(
            "FRAMEWORK 必须是 vllm_0_17_0_corex_4_4_0，"
            "不能使用 vllm_fix_tokenizer"
        )


def load_whitelist(path: Path) -> set[str]:
    names: set[str] = set()
    for line in path.read_text(encoding="utf-8").splitlines():
        item = line.split("#", 1)[0].strip()
        if item:
            names.add(item)
    if not names:
        raise PatchError(f"架构白名单为空: {path}")
    return names


def load_aliases(path: Path) -> dict[str, str]:
    data = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise PatchError(f"别名文件必须是对象: {path}")
    return {str(key): str(value) for key, value in data.items()}


def preview_default_template(messages: list[dict[str, str]], add_generation_prompt: bool = True) -> str:
    """与 default_chat_template.jinja 同一套角色标记，供无 Jinja 依赖的本地预览。"""
    parts: list[str] = []
    for message in messages:
        role = message["role"]
        content = message["content"]
        parts.append(f"<|{role}|>\n{content}<|end|>\n")
    if add_generation_prompt:
        parts.append("<|assistant|>\n")
    return "".join(parts)


def apply_chat_template(model_dir: Path, template_text: str, force: bool = False) -> dict:
    path = model_dir / "tokenizer_config.json"
    if not path.is_file():
        raise PatchError(f"缺少 tokenizer_config.json: {path}")
    data = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise PatchError("tokenizer_config.json 必须是对象")
    current = data.get("chat_template")
    if isinstance(current, str) and current.strip() and not force:
        return {"injected": False, "reason": "已有 chat_template"}
    data["chat_template"] = template_text
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return {"injected": True, "reason": "已写入 chat_template"}


def apply_architectures(
    model_dir: Path,
    whitelist: set[str],
    aliases: dict[str, str],
    unknown_policy: str = "fail",
) -> dict:
    if unknown_policy not in {"fail", "skip"}:
        raise PatchError("ARCH_UNKNOWN_POLICY 只能是 fail 或 skip")
    path = model_dir / "config.json"
    if not path.is_file():
        raise PatchError(f"缺少 config.json: {path}")
    data = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise PatchError("config.json 必须是对象")
    architectures = data.get("architectures")
    if not isinstance(architectures, list) or not architectures:
        raise PatchError("config.json 缺少 architectures")

    rewritten: list[str] = []
    skipped: list[str] = []
    updated: list[str] = []
    for name in architectures:
        if not isinstance(name, str) or not name:
            raise PatchError("architectures 含有空名称")
        if name in whitelist:
            updated.append(name)
            continue
        target = aliases.get(name)
        if target in whitelist:
            updated.append(target)
            rewritten.append(f"{name}->{target}")
            continue
        if unknown_policy == "skip":
            skipped.append(name)
            updated.append(name)
            continue
        raise PatchError(f"未知架构: {name}")

    changed = updated != architectures
    if changed:
        data["architectures"] = updated
        path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return {
        "architectures": updated,
        "rewritten": rewritten,
        "skipped": skipped,
        "changed": changed,
    }


def apply_tokenizer_patch(model_dir: Path, enabled: bool) -> dict:
    if not enabled:
        return {"applied": False}
    path = model_dir / "tokenizer_config.json"
    if not path.is_file():
        raise PatchError(f"缺少 tokenizer_config.json: {path}")
    data = json.loads(path.read_text(encoding="utf-8"))
    data[TOKENIZER_PATCH_KEY] = "optional-v1"
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return {"applied": True}


def apply_all(
    model_dir: Path,
    template_text: str,
    whitelist: set[str],
    aliases: dict[str, str],
    unknown_policy: str = "fail",
    force_chat_template: bool = False,
    tokenizer_patch: bool = False,
) -> dict:
    chat = apply_chat_template(model_dir, template_text, force=force_chat_template)
    architectures = apply_architectures(model_dir, whitelist, aliases, unknown_policy)
    tokenizer = apply_tokenizer_patch(model_dir, tokenizer_patch)
    return {"chat_template": chat, "architectures": architectures, "tokenizer": tokenizer}


def main(argv: list[str] | None = None) -> int:
    root = engine_root()
    parser = argparse.ArgumentParser(description="应用 R2/R4 模型目录修补")
    parser.add_argument("--model-dir", required=True, type=Path)
    parser.add_argument("--framework", default="vllm_0_17_0_corex_4_4_0")
    parser.add_argument("--unknown-policy", default="fail", choices=("fail", "skip"))
    parser.add_argument("--force-chat-template", action="store_true")
    parser.add_argument("--tokenizer-patch", action="store_true")
    parser.add_argument("--template", type=Path, default=root / "patches/chat_template/default_chat_template.jinja")
    parser.add_argument("--whitelist", type=Path, default=root / "patches/architectures/whitelist.txt")
    parser.add_argument("--aliases", type=Path, default=root / "patches/architectures/aliases.json")
    args = parser.parse_args(argv)
    try:
        assert_framework(args.framework)
        result = apply_all(
            args.model_dir,
            args.template.read_text(encoding="utf-8"),
            load_whitelist(args.whitelist),
            load_aliases(args.aliases),
            unknown_policy=args.unknown_policy,
            force_chat_template=args.force_chat_template,
            tokenizer_patch=args.tokenizer_patch,
        )
    except (PatchError, OSError, json.JSONDecodeError) as exc:
        print(f"修补失败: {exc}", file=sys.stderr)
        return 1
    print(json.dumps(result, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())