#!/usr/bin/env python3
"""Upload a PP-OCR checkpoint prefix and its training metadata to Hugging Face."""

from __future__ import annotations

import argparse
import json
import os
from datetime import datetime, timezone
from pathlib import Path

from huggingface_hub import HfApi


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo-id", required=True)
    parser.add_argument("--checkpoint-prefix", type=Path, required=True)
    parser.add_argument("--tag", required=True)
    parser.add_argument("--metric", type=float, required=True)
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--dictionary", type=Path, required=True)
    parser.add_argument("--token-file", type=Path, default=Path("/root/.hf_token"))
    args = parser.parse_args()

    token = args.token_file.read_text().strip()
    api = HfApi(token=token)
    api.create_repo(args.repo_id, repo_type="model", exist_ok=True)
    destination = f"checkpoints/{args.tag}"

    uploaded = []
    for suffix in (".pdparams", ".states"):
        source = Path(str(args.checkpoint_prefix) + suffix)
        if source.exists():
            api.upload_file(
                path_or_fileobj=source,
                path_in_repo=f"{destination}/{source.name}",
                repo_id=args.repo_id,
                repo_type="model",
            )
            uploaded.append(source.name)
    for source in (args.config, args.dictionary):
        api.upload_file(
            path_or_fileobj=source,
            path_in_repo=f"{destination}/{source.name}",
            repo_id=args.repo_id,
            repo_type="model",
        )
        uploaded.append(source.name)

    receipt = {
        "tag": args.tag,
        "metric": args.metric,
        "checkpoint_prefix": args.checkpoint_prefix.name,
        "uploaded": uploaded,
        "uploaded_at": datetime.now(timezone.utc).isoformat(),
    }
    api.upload_file(
        path_or_fileobj=json.dumps(receipt, indent=2).encode(),
        path_in_repo="latest-checkpoint.json",
        repo_id=args.repo_id,
        repo_type="model",
    )
    print(json.dumps(receipt))


if __name__ == "__main__":
    main()