"""Package a champion v5 checkpoint, push it on a branch and merge it via a PR.

Usage:
    python scripts/package_champion.py --checkpoint artifacts/training-runs/<run>/epoch-NNNN.ckpt \
        --config <training yaml> --test-eval <choice_evaluate json> --generation 1 \
        --state-dir artifacts/champion-loop [--no-pr]

An ensemble (mean gate margin of several checkpoints) is packaged with
`--manifest <ensemble_evaluate manifest>` instead of `--checkpoint/--config`.
"""

import argparse
import contextlib
import fcntl
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import time
import urllib.error
import urllib.request
from datetime import datetime
from pathlib import Path
from typing import Any
from zoneinfo import ZoneInfo

import yaml

KST = ZoneInfo("Asia/Seoul")
GITHUB_REPO = "shacea/Raina-laya"
HOSTS_FILE = Path.home() / ".config/gh/hosts.yml"
CREDENTIAL_HELPER = (
    '!f() { echo username=x-access-token; echo "password=$GH_TOKEN"; }; f'
)
CO_AUTHOR = "Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>"
PR_FOOTER = "🤖 Generated with [Claude Code](https://claude.com/claude-code)"
CHECKPOINT_NAME = re.compile(r"epoch-(\d{4})\.ckpt")
MERGE_TIMEOUT_SECONDS = 120
POLL_SECONDS = 2


def stamp() -> str:
    """Return the current KST time as text."""
    return datetime.now(KST).strftime("%Y-%m-%d %H:%M:%S KST")


def sha256(path: Path) -> str:
    """Return the hex SHA-256 of a file."""
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while block := handle.read(1 << 20):
            digest.update(block)
    return digest.hexdigest()


def read_token() -> str:
    """Return the github.com oauth token from the gh hosts file (never logged)."""
    in_block = False
    for line in HOSTS_FILE.read_text().splitlines():
        if line and not line[0].isspace():
            in_block = line.strip() == "github.com:"
        elif in_block and (found := re.match(r"\s+oauth_token:\s*(\S+)", line)):
            return found[1]
    msg = f"no github.com oauth_token in {HOSTS_FILE}"
    raise RuntimeError(msg)


def git(cwd: Path, *args: str, auth: bool = False) -> str:
    """Run git in cwd; with auth the token reaches git only via the subprocess env."""
    env = None
    prefix: list[str] = []
    if auth:
        env = {**os.environ, "GH_TOKEN": read_token(), "GIT_TERMINAL_PROMPT": "0"}
        prefix = [
            "-c",
            "credential.helper=",
            "-c",
            f"credential.helper={CREDENTIAL_HELPER}",
        ]
    done = subprocess.run(
        ["git", "-C", str(cwd), *prefix, *args],  # noqa: S607
        env=env,
        capture_output=True,
        text=True,
        check=False,
    )
    if done.returncode:
        msg = f"git {' '.join(args)} failed: {done.stderr.strip()}"
        raise RuntimeError(msg)
    return done.stdout.strip()


def api(method: str, path: str, body: dict[str, Any] | None = None) -> dict[str, Any]:
    """Call the GitHub REST API and return the decoded JSON response."""
    request = urllib.request.Request(  # noqa: S310
        f"https://api.github.com{path}",
        data=None if body is None else json.dumps(body).encode(),
        method=method,
        headers={
            "Authorization": f"Bearer {read_token()}",
            "Accept": "application/vnd.github+json",
            "User-Agent": "raina-laya-package-champion",
        },
    )
    try:
        with urllib.request.urlopen(request, timeout=60) as response:  # noqa: S310
            return json.load(response)
    except urllib.error.HTTPError as error:
        msg = f"{method} {path} -> HTTP {error.code}: {error.read().decode()[:500]}"
        raise RuntimeError(msg) from error


def pct(count: int, total: int) -> str:
    """Format a count over a total as `count/total (xx.xx%)`."""
    return f"{count}/{total} ({100 * count / total:.2f}%)"


def readme(info: dict[str, Any]) -> str:
    """Render the Korean package README."""
    name, epoch, test, val = info["name"], info["epoch"], info["test"], info["val"]
    counts = test["test_counts"]
    role = test.get("role", "test")
    lines = [
        f"# Raina-Laya v5 최고 모델 패키지 ({name})",
        "",
        (
            "이 디렉터리는 연구 루프가 최고 성능(champion) 모델을 자동으로 패키징한 결과입니다(AGENTS.md §1-1). "
            "이 패키징 작업에서는 학습, 추가 평가, 임계값 변경을 수행하지 않았고 "
            "이미 수행된 test 셋 평가 결과를 그대로 보존했습니다."
        ),
        "",
        "## 모델",
        "",
        f"- 모델 계열: `{test['model_family']}`",
        f"- 세대(generation): `{info['generation']}`",
        f"- 학습 실행(run): `{info['run_id']}`",
        f"- 선택 epoch: `{epoch}`",
        f"- 판단부 체크포인트: `checkpoint/epoch-{epoch:04d}.ckpt` (용량 때문에 git에는 포함하지 않고 로컬에만 보관합니다)",
        "- 클래스 순서: `Fail`, `A`, `S`",
        f"- 패키지 생성 시각: {info['created']}",
        f"- 소스 저장소 커밋: `{info['head']}`",
        "- 설정·선택 근거: `metadata/`, 기계 판독 출처·해시: `provenance.json`",
        "",
        f"## {role} 셋 성능",
        "",
        f"- {role} 셋: `{test['split']}` (곡 {test['test_songs']}개, Fail {counts['Fail']}, A {counts['A']}, S {counts['S']})",
        f"- Fail→Pass(실제 Fail을 A/S로 판정): {pct(test['fail_to_pass_count'], test['fail_count'])}",
        f"- Pass→Fail(실제 A/S를 Fail로 판정): {pct(test['pass_count'] - test['pass_correct_count'], test['pass_count'])}",
        f"- macro F1: {test['macro_f1']:.4f}, S F1: {test['s_f1']:.4f}, A F1: {test['a_f1']:.4f}, Pass 재현율: {test['pass_recall']:.4f}",
        "- 원본 평가 결과: `metadata/test-eval.json`",
        "",
        "## 검증(validation) 성능",
        "",
    ]
    if val:
        matrix = val["confusion_matrix"]
        fail_total = sum(matrix[0])
        pass_total = sum(matrix[1]) + sum(matrix[2])
        lines += [
            f"- Fail→Pass: {pct(sum(matrix[0][1:]), fail_total)}",
            f"- Pass→Fail: {pct(matrix[1][0] + matrix[2][0], pass_total)}",
            f"- macro F1: {val['macro_f1']:.4f}, S F1: {val['s_f1']:.4f}, Pass 재현율: {val['pass_recall']:.4f}",
        ]
    else:
        lines.append("- 이 epoch의 검증 기록이 실행 디렉터리에 없었습니다.")
    lines += [
        "",
        "## 평가 방법",
        "",
        (
            "저장소 루트에서 실행하며, CUDA 전용이므로 물리 GPU 0 또는 3만 사용합니다. "
            "설정이 가리키는 split 파일(`provenance.json`의 `split_file_sha256`와 일치해야 함)과 "
            "MuQ 특징 캐시가 필요합니다. 같은 test 셋의 반복 평가는 포화 위험이 있으므로 "
            "AGENTS.md §1-1의 test 셋 운용 규칙을 따릅니다."
        ),
        "",
        "```bash",
        "CUDA_VISIBLE_DEVICES=0 python -m raina_laya.choice_evaluate \\",
        f"  --config model/{name}/metadata/config.yaml \\",
        f"  --checkpoint model/{name}/checkpoint/epoch-{epoch:04d}.ckpt \\",
        "  --output <출력 json 경로>",
        "```",
        "",
        "## 배포 상태",
        "",
        (
            "`src/raina_laya/features/serving`은 현재 `raina_laya_choice_siglip_v4`만 적재하므로 "
            "이 패키지는 아직 backend에 배포할 수 없습니다. backend 배포는 사용자 승인이 필요합니다."
        ),
        "",
        "## 오프라인 MuQ 캐시",
        "",
        (
            f"`huggingface/`는 `{info['hf_source']}`를 하드링크로 복제한 고정 MuQ 오프라인 캐시이며 "
            "git에는 포함하지 않습니다. `HF_HOME`과 `HF_HUB_CACHE`를 이 디렉터리로 지정하고 "
            "`HF_HUB_OFFLINE=1`로 사용합니다."
        ),
        "",
    ]
    return "\n".join(lines)


def ensemble_readme(info: dict[str, Any]) -> str:
    """Render the Korean README of an ensemble package."""
    name, test, members = info["name"], info["test"], info["members"]
    counts = test["test_counts"]
    role = test.get("role", "test")
    lines = [
        f"# Raina-Laya v5 최고 모델(앙상블) 패키지 ({name})",
        "",
        (
            "이 디렉터리는 연구 루프가 최고 성능(champion) 앙상블을 자동으로 패키징한 결과입니다(AGENTS.md §1-1). "
            "이 패키징 작업에서는 학습, 추가 평가, 임계값 변경을 수행하지 않았고 "
            "이미 수행된 test 셋 평가 결과를 그대로 보존했습니다."
        ),
        "",
        "## 앙상블 규칙",
        "",
        (
            "멤버 체크포인트마다 곡별 게이트 마진(gate margin)을 계산하고, 그 평균(`mean_gate_margin`)으로 "
            "Pass/Fail을 판정합니다. 멤버별 설정과 체크포인트는 `metadata/manifest.json`이 가리킵니다."
        ),
        "",
        "## 모델",
        "",
        f"- 모델 계열: `{test['model_family']}`",
        f"- 세대(generation): `{info['generation']}`",
        f"- 앙상블 매니페스트: `{info['run_id']}` (멤버 {len(members)}개)",
        "- 멤버 체크포인트: `checkpoint/` (용량 때문에 git에는 포함하지 않고 로컬에만 보관합니다)",
        "- 클래스 순서: `Fail`, `A`, `S`",
        f"- 패키지 생성 시각: {info['created']}",
        f"- 소스 저장소 커밋: `{info['head']}`",
        "- 기계 판독 출처·해시: `provenance.json`",
        "",
        "| # | 실행(run) | epoch | 체크포인트 |",
        "| --- | --- | --- | --- |",
        *(
            f"| {i} | `{m['run_id']}` | {m['epoch']} | `checkpoint/{m['packaged_name']}` |"
            for i, m in enumerate(members)
        ),
        "",
        f"## {role} 셋 성능",
        "",
        f"- {role} 셋: `{test['split']}` (곡 {test['test_songs']}개, Fail {counts['Fail']}, A {counts['A']}, S {counts['S']})",
        f"- Fail→Pass(실제 Fail을 A/S로 판정): {pct(test['fail_to_pass_count'], test['fail_count'])}",
        f"- Pass→Fail(실제 A/S를 Fail로 판정): {pct(test['pass_count'] - test['pass_correct_count'], test['pass_count'])}",
        f"- macro F1: {test['macro_f1']:.4f}, S F1: {test['s_f1']:.4f}, A F1: {test['a_f1']:.4f}, Pass 재현율: {test['pass_recall']:.4f}",
        "- 원본 평가 결과: `metadata/test-eval.json`",
        "",
        "## 평가 방법",
        "",
        (
            "저장소 루트에서 실행하며, CUDA 전용이므로 물리 GPU 0 또는 3만 사용합니다. "
            "설정이 가리키는 split 파일(`provenance.json`의 `split_file_sha256`와 일치해야 함)과 "
            "MuQ 특징 캐시가 필요합니다. 같은 test 셋의 반복 평가는 포화 위험이 있으므로 "
            "AGENTS.md §1-1의 test 셋 운용 규칙을 따릅니다."
        ),
        "",
        "```bash",
        "CUDA_VISIBLE_DEVICES=0 python -m raina_laya.ensemble_evaluate \\",
        f"  --manifest model/{name}/metadata/manifest.json \\",
        "  --role test \\",
        "  --output <출력 json 경로>",
        "```",
        "",
        "## 배포 상태",
        "",
        (
            "`src/raina_laya/features/serving`은 현재 `raina_laya_choice_siglip_v4`만 적재하므로 "
            "이 패키지는 아직 backend에 배포할 수 없습니다. backend 배포는 사용자 승인이 필요합니다."
        ),
        "",
        "## 오프라인 MuQ 캐시",
        "",
        (
            f"`huggingface/`는 `{info['hf_source']}`를 하드링크로 복제한 고정 MuQ 오프라인 캐시이며 "
            "git에는 포함하지 않습니다. `HF_HOME`과 `HF_HUB_CACHE`를 이 디렉터리로 지정하고 "
            "`HF_HUB_OFFLINE=1`로 사용합니다."
        ),
        "",
    ]
    return "\n".join(lines)


def copy_hf_cache(source: Path, dest: Path) -> None:
    """Hardlink-copy the MuQ cache, falling back to a full copy across filesystems."""
    try:
        _ = subprocess.run(
            ["cp", "-al", str(source), str(dest)],  # noqa: S607
            check=True,
            capture_output=True,
        )
    except subprocess.CalledProcessError:
        shutil.rmtree(dest, ignore_errors=True)
        _ = subprocess.run(
            ["cp", "-a", str(source), str(dest)],  # noqa: S607
            check=True,
        )


def build_ensemble(args: argparse.Namespace, pkg: Path, info: dict[str, Any]) -> None:
    """Write member checkpoints, configs, manifest, provenance and README."""
    name, test = info["name"], info["test"]
    prefix = f"model/{name}"
    (pkg / "checkpoint").mkdir(parents=True)
    (pkg / "metadata/members").mkdir(parents=True)
    teachers: dict[str, str] = {}  # teacher sha256 -> packaged path
    manifest: list[dict[str, str]] = []
    for index, member in enumerate(info["members"]):
        member["packaged_name"] = (
            f"{index:02d}-{member['run_id']}-{member['checkpoint'].name}"
        )
        _ = shutil.copy2(
            member["checkpoint"], pkg / "checkpoint" / member["packaged_name"]
        )
        config_name = f"metadata/members/{index:02d}-config.yaml"
        config = yaml.safe_load(member["config"].read_text())
        teacher = config.get("gate_soft_targets")
        if teacher is None:
            _ = shutil.copy2(member["config"], pkg / config_name)
        else:
            # choice_evaluate compares the teacher file's sha256, so ship the file.
            source = args.repo / teacher
            if not source.is_file():
                msg = f"gate_soft_targets teacher file is missing: {source}"
                raise FileNotFoundError(msg)
            digest = sha256(source)
            if digest not in teachers:
                teacher_name = f"metadata/members/teacher-{digest[:12]}.jsonl"
                teachers[digest] = f"{prefix}/{teacher_name}"
                _ = shutil.copy2(source, pkg / teacher_name)
            config["gate_soft_targets"] = teachers[digest]
            (pkg / config_name).write_text(yaml.safe_dump(config, sort_keys=False))
        member["config_sha256"] = sha256(member["config"])
        manifest.append(
            {
                "config": f"{prefix}/{config_name}",
                "checkpoint": f"{prefix}/checkpoint/{member['packaged_name']}",
            }
        )
    (pkg / "metadata/manifest.json").write_text(
        json.dumps({"members": manifest}, indent=2) + "\n"
    )
    _ = shutil.copy2(args.test_eval, pkg / "metadata/test-eval.json")
    copy_hf_cache(info["hf_source"], pkg / "huggingface")
    provenance = {
        "schema_version": 1,
        "package": name,
        "model_family": test["model_family"],
        "generation": info["generation"],
        "run_id": info["run_id"],
        "ensemble": "mean_gate_margin",
        "created_kst": info["created"],
        "source_repo": str(args.repo),
        "source_repo_head": info["head"],
        "source_manifest": str(args.manifest),
        "checkpoint_sha256": info["checkpoint_sha256"],
        "manifest_sha256": test.get("manifest_sha256"),
        "packaged_manifest_sha256": sha256(pkg / "metadata/manifest.json"),
        "members": [
            {
                "run_id": m["run_id"],
                "epoch": m["epoch"],
                "source_checkpoint": str(m["checkpoint"]),
                "checkpoint_sha256": m["sha256"],
                "config_sha256": m["config_sha256"],
                **entry,
            }
            for m, entry in zip(info["members"], manifest, strict=True)
        ],
        "test_eval_sha256": sha256(args.test_eval),
        "eval_role": test.get("role", "test"),
        "split": info["split"],
        "split_file_sha256": info["split_sha256"],
        "split_digest": test["split_digest"],
        "cache_inventory_digest": test["cache_inventory_digest"],
        "huggingface_source": str(info["hf_source"]),
        "test_metrics": {
            key: test[key]
            for key in (
                "fail_to_pass_count",
                "fail_count",
                "pass_correct_count",
                "pass_count",
                "macro_f1",
                "s_f1",
            )
        },
        # Validation-chosen reject band and decided-song metrics (None before them).
        "reject_band": test.get("reject_band"),
        "decided": test.get("decided"),
    }
    if teachers:
        provenance["gate_soft_targets"] = [
            {"path": path, "sha256": digest} for digest, path in teachers.items()
        ]
    (pkg / "provenance.json").write_text(
        json.dumps(provenance, indent=2, sort_keys=True) + "\n"
    )
    (pkg / "README.md").write_text(ensemble_readme(info))


def build(args: argparse.Namespace, pkg: Path, info: dict[str, Any]) -> None:
    """Write checkpoint, metadata, provenance and README into the package dir."""
    if args.manifest is not None:
        build_ensemble(args, pkg, info)
        return
    epoch, test = info["epoch"], info["test"]
    run_dir = args.checkpoint.parent
    (pkg / "checkpoint").mkdir(parents=True)
    (pkg / "metadata").mkdir()
    _ = shutil.copy2(args.checkpoint, pkg / "checkpoint" / args.checkpoint.name)
    config = yaml.safe_load(args.config.read_text())
    teacher = config.get("gate_soft_targets")
    if teacher is None:
        _ = shutil.copy2(args.config, pkg / "metadata/config.yaml")
    else:
        # choice_evaluate compares the teacher file's sha256, so ship the file.
        source = args.repo / teacher
        if not source.is_file():
            msg = f"gate_soft_targets teacher file is missing: {source}"
            raise FileNotFoundError(msg)
        packaged = f"model/{info['name']}/metadata/teacher.jsonl"
        _ = shutil.copy2(source, pkg / "metadata/teacher.jsonl")
        config["gate_soft_targets"] = packaged
        (pkg / "metadata/config.yaml").write_text(
            yaml.safe_dump(config, sort_keys=False)
        )
    _ = shutil.copy2(args.test_eval, pkg / "metadata/test-eval.json")
    for name in (
        "selection.json",
        "pareto.json",
        f"epoch-{epoch:04d}.json",
        f"validation-epoch-{epoch:04d}.json",
    ):
        if (run_dir / name).exists():
            _ = shutil.copy2(run_dir / name, pkg / "metadata" / name)
    copy_hf_cache(info["hf_source"], pkg / "huggingface")
    val_path = run_dir / f"validation-epoch-{epoch:04d}.json"
    info["val"] = json.loads(val_path.read_text()) if val_path.exists() else None
    provenance = {
        "schema_version": 1,
        "package": info["name"],
        "model_family": test["model_family"],
        "generation": info["generation"],
        "run_id": info["run_id"],
        "epoch": epoch,
        "created_kst": info["created"],
        "source_repo": str(args.repo),
        "source_repo_head": info["head"],
        "source_checkpoint": str(args.checkpoint),
        "checkpoint_sha256": info["checkpoint_sha256"],
        "config_sha256": sha256(args.config),
        "test_eval_sha256": sha256(args.test_eval),
        "eval_role": test.get("role", "test"),
        "split": info["split"],
        "split_file_sha256": info["split_sha256"],
        "split_digest": test["split_digest"],
        "cache_inventory_digest": test["cache_inventory_digest"],
        "huggingface_source": str(info["hf_source"]),
        "test_metrics": {
            key: test[key]
            for key in (
                "fail_to_pass_count",
                "fail_count",
                "pass_correct_count",
                "pass_count",
                "macro_f1",
                "s_f1",
            )
        },
        # Validation-chosen reject band and decided-song metrics (None before them).
        "reject_band": test.get("reject_band"),
        "decided": test.get("decided"),
    }
    if teacher is not None:
        provenance["gate_soft_targets"] = {
            "path": packaged,
            "sha256": sha256(pkg / "metadata/teacher.jsonl"),
        }
    (pkg / "provenance.json").write_text(
        json.dumps(provenance, indent=2, sort_keys=True) + "\n"
    )
    (pkg / "README.md").write_text(readme(info))


def commit_package(root: Path, name: str, info: dict[str, Any]) -> str:
    """Stage only non-ignored package files, commit and return the commit sha."""
    test = info["test"]
    files = git(
        root, "ls-files", "--others", "--exclude-standard", "--", f"model/{name}"
    ).splitlines()
    leaked = [f for f in files if f.endswith(".ckpt") or "/huggingface/" in f]
    if leaked or not files:
        msg = f".gitignore does not keep the package clean: leaked={leaked[:3]}"
        raise RuntimeError(msg)
    _ = git(root, "add", "--", *files)
    _ = git(
        root,
        "commit",
        "-m",
        f"feat(model): package v5 champion {name}",
        "-m",
        f"{test.get('role', 'test')} set {test['split']}: Fail->Pass {info['f2p']}, Pass->Fail {info['p2f']}.\n"
        f"macro F1 {test['macro_f1']:.4f}, S F1 {test['s_f1']:.4f}. "
        f"{info['run_text']}.\n"
        "Checkpoint and MuQ cache are gitignored and stay local.",
        "-m",
        CO_AUTHOR,
    )
    return git(root, "rev-parse", "HEAD")


def open_and_merge(
    branch: str, name: str, head_sha: str, info: dict[str, Any], record: dict[str, Any]
) -> None:
    """Open the PR, wait until it is mergeable, then merge it with a merge commit."""
    test = info["test"]
    body = "\n".join(
        [
            f"Automated v5 champion package (AGENTS.md section 1-1): `{name}`.",
            "",
            "| Metric | Value |",
            "| --- | --- |",
            f"| Evaluation set ({test.get('role', 'test')}) | `{test['split']}` ({test['test_songs']} songs) |",
            f"| Fail->Pass | {info['f2p']} |",
            f"| Pass->Fail | {info['p2f']} |",
            f"| macro F1 | {test['macro_f1']:.4f} |",
            f"| S F1 | {test['s_f1']:.4f} |",
            f"| Run / epoch / generation | {info['run_cell']} |",
            "",
            (
                "The checkpoint and MuQ cache are gitignored and stay local. "
                "Serving still loads only v4, so this package is not deployed."
            ),
            "",
            PR_FOOTER,
        ]
    )
    pulls = f"/repos/{GITHUB_REPO}/pulls"
    pr = api(
        "POST",
        pulls,
        {
            "title": f"Package v5 champion {name}",
            "head": branch,
            "base": "main",
            "body": body,
        },
    )
    number = pr["number"]
    record["pr_number"], record["pr_url"] = number, pr["html_url"]
    deadline = time.monotonic() + MERGE_TIMEOUT_SECONDS
    # GitHub reports mergeable=null while it computes the merge state.
    while (pr := api("GET", f"{pulls}/{number}"))["mergeable"] is not True:
        if pr["mergeable"] is False or time.monotonic() > deadline:
            msg = f"PR #{number} is not mergeable: {pr['mergeable']!r}"
            raise RuntimeError(msg)
        time.sleep(POLL_SECONDS)
    merged = api(
        "PUT",
        f"{pulls}/{number}/merge",
        {"merge_method": "merge", "sha": head_sha},
    )
    record["merged"], record["merge_sha"] = bool(merged["merged"]), merged["sha"]


def inspect_ensemble(args: argparse.Namespace) -> dict[str, Any]:
    """Validate the manifest members, test eval, split and cache source."""
    args.manifest = args.manifest.resolve()
    entries = json.loads(args.manifest.read_text())["members"]
    if not entries:
        msg = "manifest has no members"
        raise ValueError(msg)
    members: list[dict[str, Any]] = []
    for entry in entries:
        checkpoint, config = Path(entry["checkpoint"]).resolve(), Path(entry["config"])
        match = CHECKPOINT_NAME.fullmatch(checkpoint.name)
        if not match:
            msg = f"member checkpoint must be named epoch-NNNN.ckpt: {checkpoint.name}"
            raise ValueError(msg)
        members.append(
            {
                "run_id": checkpoint.parent.name,
                "epoch": int(match[1]),
                "checkpoint": checkpoint,
                "config": config,
                "sha256": sha256(checkpoint),
            }
        )
    test = json.loads(args.test_eval.read_text())
    combined = hashlib.sha256(
        "".join(m["sha256"] for m in members).encode()
    ).hexdigest()
    if test["checkpoint_sha256"] != combined:
        msg = "test eval does not describe these member checkpoints (sha256 mismatch)"
        raise ValueError(msg)
    splits = {yaml.safe_load(m["config"].read_text())["split"] for m in members}
    if len(splits) != 1:
        msg = f"member configs use different splits: {sorted(splits)}"
        raise ValueError(msg)
    (split,) = splits
    split_sha = sha256(args.repo / split)
    if test["split_file_sha256"] != split_sha:
        msg = f"split file {split} differs from the one the test eval used"
        raise ValueError(msg)
    hf_source = args.repo / args.hf_source
    if not hf_source.is_dir():
        msg = f"MuQ cache source is not a directory: {hf_source}"
        raise ValueError(msg)
    run_id = args.manifest.stem
    name = re.sub(
        r"[^a-z0-9._-]+",
        "-",
        f"raina-laya-v5-g{args.generation:03d}-ens-{run_id}".lower(),
    )
    return {
        "name": name,
        "epoch": None,
        "run_id": run_id,
        "run_text": f"Ensemble {run_id} of {len(members)} members, generation {args.generation}",
        "run_cell": f"`{run_id}` / {len(members)} members / {args.generation}",
        "generation": args.generation,
        "members": members,
        "test": test,
        "created": stamp(),
        "head": git(args.repo, "rev-parse", "HEAD"),
        "checkpoint_sha256": combined,
        "split": split,
        "split_sha256": split_sha,
        "hf_source": hf_source,
        "f2p": pct(test["fail_to_pass_count"], test["fail_count"]),
        "p2f": pct(test["pass_count"] - test["pass_correct_count"], test["pass_count"]),
    }


def inspect_inputs(args: argparse.Namespace) -> dict[str, Any]:
    """Validate the checkpoint, test eval, split and cache source; return the facts."""
    if args.manifest is not None:
        return inspect_ensemble(args)
    args.checkpoint = args.checkpoint.resolve()
    match = CHECKPOINT_NAME.fullmatch(args.checkpoint.name)
    if not match:
        msg = f"checkpoint must be named epoch-NNNN.ckpt: {args.checkpoint.name}"
        raise ValueError(msg)
    epoch, run_id = int(match[1]), args.checkpoint.parent.name
    test = json.loads(args.test_eval.read_text())
    checkpoint_sha = sha256(args.checkpoint)
    if test["checkpoint_sha256"] != checkpoint_sha or test["epoch"] != epoch:
        msg = "test eval does not describe this checkpoint (sha256 or epoch mismatch)"
        raise ValueError(msg)
    split = yaml.safe_load(args.config.read_text())["split"]
    split_sha = sha256(args.repo / split)
    if test["split_file_sha256"] != split_sha:
        msg = f"split file {split} differs from the one the test eval used"
        raise ValueError(msg)
    hf_source = args.repo / args.hf_source
    if not hf_source.is_dir():
        msg = f"MuQ cache source is not a directory: {hf_source}"
        raise ValueError(msg)
    name = re.sub(
        r"[^a-z0-9._-]+",
        "-",
        f"raina-laya-v5-g{args.generation:03d}-{run_id}-e{epoch:04d}".lower(),
    )
    return {
        "name": name,
        "epoch": epoch,
        "run_id": run_id,
        "run_text": f"Run {run_id}, epoch {epoch}, generation {args.generation}",
        "run_cell": f"`{run_id}` / {epoch} / {args.generation}",
        "generation": args.generation,
        "test": test,
        "created": stamp(),
        "head": git(args.repo, "rev-parse", "HEAD"),
        "checkpoint_sha256": checkpoint_sha,
        "split": split,
        "split_sha256": split_sha,
        "hf_source": hf_source,
        "f2p": pct(test["fail_to_pass_count"], test["fail_count"]),
        "p2f": pct(test["pass_count"] - test["pass_correct_count"], test["pass_count"]),
    }


def run(args: argparse.Namespace, record: dict[str, Any]) -> None:
    """Build the package in the package worktree and publish it."""
    info = inspect_inputs(args)
    name, root = info["name"], args.package_root
    branch, pkg = f"package/{name}", root / "model" / name
    record.update(name=name, package_path=str(pkg), branch=branch)

    # ponytail: the very first worktree creation is outside the lock, because
    # `git worktree add` needs an empty target; only a first-ever race can fail.
    if not (root / ".git").exists():
        _ = git(args.repo, "worktree", "add", "--detach", str(root), "origin/main")
    with (root / ".package.lock").open("w") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        _ = git(root, "fetch", "origin", auth=True)
        _ = git(root, "checkout", "--detach", "origin/main")
        if pkg.exists():
            msg = f"package already exists: {pkg}"
            raise FileExistsError(msg)
        try:
            _ = git(root, "checkout", "-B", branch)
            build(args, pkg, info)
            record["commit_sha"] = commit_package(root, name, info)
            # Force: the branch name is deterministic and only this script writes it.
            _ = git(root, "push", "--force", "origin", f"{branch}:{branch}", auth=True)
            if not args.no_pr:
                open_and_merge(branch, name, record["commit_sha"], info, record)
                _ = git(root, "fetch", "origin", auth=True)
                _ = git(root, "checkout", "--detach", "origin/main")
                _ = git(root, "branch", "-D", branch)
        except Exception:
            if not record["merged"]:
                # Leave a retryable worktree: only ignored/untracked leftovers of ours.
                with contextlib.suppress(RuntimeError):
                    _ = git(root, "checkout", "--detach", "-f", "origin/main")
                shutil.rmtree(pkg, ignore_errors=True)
            raise


def main(argv: list[str] | None = None) -> int:
    """Package one champion, record the outcome in packages.jsonl, return exit code."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", type=Path)
    parser.add_argument("--config", type=Path)
    parser.add_argument("--manifest", type=Path)
    parser.add_argument("--test-eval", type=Path, required=True)
    parser.add_argument("--generation", type=int, required=True)
    parser.add_argument("--state-dir", type=Path, required=True)
    parser.add_argument("--repo", type=Path)
    parser.add_argument("--package-root", type=Path)
    parser.add_argument("--no-pr", action="store_true")
    parser.add_argument(
        "--hf-source",
        type=Path,
        default=Path("model/raina-laya-best-20261001/huggingface"),
    )
    args = parser.parse_args(argv)
    if args.manifest is None:
        if args.checkpoint is None or args.config is None:
            parser.error("--checkpoint and --config are required without --manifest")
    elif args.checkpoint is not None or args.config is not None:
        parser.error("--manifest excludes --checkpoint and --config")
    if args.repo is None:
        args.repo = Path(git(Path.cwd(), "rev-parse", "--show-toplevel"))
    args.repo = args.repo.resolve()
    if args.package_root is None:
        args.package_root = args.repo.parent / "Raina-laya-packages"
    args.package_root = args.package_root.resolve()

    record: dict[str, Any] = {
        "started_kst": stamp(),
        "name": None,
        "package_path": None,
        "commit_sha": None,
        "pr_number": None,
        "pr_url": None,
        "merged": False,
        "merge_sha": None,
        "error": None,
    }
    try:
        run(args, record)
    except Exception as error:  # noqa: BLE001
        record["error"] = f"{type(error).__name__}: {error}"
    record["finished_kst"] = stamp()
    args.state_dir.mkdir(parents=True, exist_ok=True)
    with (args.state_dir / "packages.jsonl").open("a") as handle:
        _ = handle.write(json.dumps(record, sort_keys=True) + "\n")
        handle.flush()
        os.fsync(handle.fileno())
    sys.stdout.write(json.dumps(record, sort_keys=True) + "\n")
    return 1 if record["error"] else 0


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