#!/usr/bin/env python3
"""Upload an already-complete local evaluation result to SwanLab."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import swanlab


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--result-dir", type=Path, required=True)
    parser.add_argument("--run-name", required=True)
    parser.add_argument("--project", default="unified-model")
    args = parser.parse_args()
    result_dir = args.result_dir.resolve()
    payload = json.loads((result_dir / "metrics.json").read_text(encoding="utf-8"))
    metadata = payload["metadata"]
    metrics = payload["metrics"]
    benchmark = str(metadata["benchmark"])
    run = swanlab.init(
        project=args.project,
        name=args.run_name,
        log_dir=str(result_dir / "swanlab_retry"),
        config=metadata,
        tags=[benchmark, "evaluation", "checkpoint-7000", "8GPU"],
        description=f"Judge-free exact-match evaluation on {benchmark}; uploaded from validated local artifacts.",
    )
    swanlab.log(metrics, step=int(metadata["checkpoint_step"]))
    status = {
        "enabled": True,
        "retry_upload": True,
        "run_id": getattr(run, "id", None),
        "run_url": str(getattr(run, "url", "") or "") or None,
    }
    swanlab.finish()
    (result_dir / "swanlab.json").write_text(
        json.dumps(status, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(status, sort_keys=True))


if __name__ == "__main__":
    main()