#!/usr/bin/env python3
"""Validate and publish this bundle without storing or printing HF_TOKEN."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path


BUNDLE = Path(__file__).resolve().parent
REQUIRED = {
    "README.md",
    "ATTRIBUTIONS.md",
    "SHA256SUMS",
    "encoder/config.json",
    "evaluation_results.json",
    "model.safetensors",
    "provenance.json",
    "rl_agent_config.json",
    "tokenizer/tokenizer.json",
    "tokenizer/tokenizer_config.json",
    "training.json",
}


def digest(path: Path) -> str:
    value = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            value.update(chunk)
    return value.hexdigest()


def validate() -> None:
    missing = sorted(name for name in REQUIRED if not (BUNDLE / name).is_file())
    if missing:
        raise SystemExit(f"Missing release files: {', '.join(missing)}")
    for name in ("encoder/config.json", "evaluation_results.json", "provenance.json", "rl_agent_config.json", "tokenizer/tokenizer_config.json", "training.json"):
        with (BUNDLE / name).open(encoding="utf-8") as handle:
            json.load(handle)
    lines = (BUNDLE / "SHA256SUMS").read_text(encoding="utf-8").splitlines()
    expected = {}
    for line in lines:
        checksum, name = line.split("  ", 1)
        expected[name] = checksum
    for name, checksum in expected.items():
        path = BUNDLE / name
        if not path.is_file() or digest(path) != checksum:
            raise SystemExit(f"Checksum mismatch: {name}")
    if not (BUNDLE / "README.md").read_text(encoding="utf-8").startswith("---\n"):
        raise SystemExit("README.md has no Hugging Face YAML metadata")
    print(f"Validated {len(expected)} checksums in {BUNDLE}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--check-only", action="store_true")
    parser.add_argument("--repo-id", help="Defaults to <authenticated-user>/laya-pt-es-nli")
    args = parser.parse_args()
    validate()
    if args.check_only:
        return
    token = os.environ.get("HF_TOKEN")
    if not token:
        raise SystemExit("HF_TOKEN is required for upload")
    from huggingface_hub import HfApi

    api = HfApi(token=token)
    identity = api.whoami()
    owner = str(identity["name"])
    repo_id = args.repo_id or f"{owner}/laya-pt-es-nli"
    requested_owner, separator, name = repo_id.partition("/")
    allowed = {owner, *(str(org["name"]) for org in identity.get("orgs", []))}
    if not separator or name != "laya-pt-es-nli" or requested_owner not in allowed:
        raise SystemExit("repo-id must be <authenticated-user-or-org>/laya-pt-es-nli")
    api.create_repo(repo_id=repo_id, repo_type="model", private=False, exist_ok=True)
    api.upload_folder(
        repo_id=repo_id,
        repo_type="model",
        folder_path=BUNDLE,
        commit_message="Publish Laya PT/ES NLI checkpoint",
        ignore_patterns=["__pycache__/**", "*.pyc"],
    )
    print(f"Published https://huggingface.co/{repo_id}")


if __name__ == "__main__":
    main()