#!/usr/bin/env python3
"""Publish public metrics and bounded visualization snapshots atomically to HF."""

from __future__ import annotations

import argparse
import json
import os
import shutil
import sys
import tempfile
import time
from datetime import datetime, timezone
from pathlib import Path

from dotenv import load_dotenv

REPO_ROOT = Path(__file__).resolve().parents[1]
load_dotenv(REPO_ROOT / ".env")
DEFAULT_RESULTS = REPO_ROOT / "space" / "results"
DEFAULT_REPO = "ThinkcatLab/LiveHouse-TS"
API_RETRIES = 5
UPLOAD_IGNORE_PATTERNS = [
    # Raw inference/task archives stay local. public_snapshots/<model>/<dataset>.json
    # contains the bounded, explicit visualization export and MUST be published.
    "prequential/**",
    "saved_evals/**",
    "*/forecasts/**",
    "**/forecasts/**",
    "evaluation_metrics.jsonl",
    "evaluation_metrics_canonical.jsonl",
    "aggregates/**/evaluation_metrics_canonical.jsonl",
    "aggregates/canonical_all_history/**",
    "aggregates/*.zip",
    ".DS_Store",
]
REMOTE_PRUNE_PATHS = [
    "results/evaluation_metrics.jsonl",
    "results/evaluation_metrics_canonical.jsonl",
    "results/aggregates/results.zip",
    "results/aggregates/aggregates.zip",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo", default=os.getenv("HF_SPACE_REPO", DEFAULT_REPO))
    parser.add_argument("--results", type=Path, default=DEFAULT_RESULTS)
    parser.add_argument("--results-repo", default=os.getenv("HF_RESULTS_REPO"))
    parser.add_argument(
        "--commit-message",
        default="Update leaderboard results from local eval",
    )
    parser.add_argument(
        "--restart-space",
        action="store_true",
        help=(
            "Explicitly restart the Space after pushing results. Disabled by default "
            "to avoid periodic 504s during live leaderboard refreshes."
        ),
    )
    return parser.parse_args()


def load_token() -> str:
    load_dotenv(REPO_ROOT / ".env")
    token = (os.getenv("HF_TOKEN") or "").strip()
    if not token:
        print("HF_TOKEN is required (add to .env)", file=sys.stderr)
        raise SystemExit(1)
    return token


def get_hf_api(token: str):
    from huggingface_hub import HfApi

    endpoint = (os.getenv("HF_ENDPOINT") or "").strip() or None
    return HfApi(token=token, endpoint=endpoint)


def ensure_space_repo(token: str, repo_id: str) -> None:
    try:
        from huggingface_hub.utils import RepositoryNotFoundError
    except ImportError:
        from huggingface_hub.errors import RepositoryNotFoundError

    api = get_hf_api(token)
    try:
        api.repo_info(repo_id, repo_type="space")
    except RepositoryNotFoundError:
        print(f"Creating Space {repo_id}…", file=sys.stderr)
        api.create_repo(
            repo_id=repo_id,
            repo_type="space",
            space_sdk="gradio",
            exist_ok=True,
        )
        print(f"Created Space {repo_id}", file=sys.stderr)


def prune_remote_non_leaderboard_artifacts(token: str, repo_id: str) -> bool:
    """Delete remote-only artifacts that are not read by the Space UI.

    The evaluation metric JSONL files are local aggregation caches. Once they grow beyond
    the Hub git limit, keeping it in the Space repository blocks every refresh.
    """

    api = get_hf_api(token)
    try:
        remote_files = set(api.list_repo_files(repo_id=repo_id, repo_type="space"))
    except Exception as exc:
        print(f"Could not verify remote files before upload: {exc}", file=sys.stderr)
        return False

    stale_paths = sorted(path for path in remote_files if (
        path in REMOTE_PRUNE_PATHS or path.startswith("results/prequential/")
        or path.startswith("results/saved_evals/") or "/forecasts/" in path
    ))
    if not stale_paths:
        return True

    try:
        from huggingface_hub import CommitOperationDelete

        api.create_commit(
            repo_id=repo_id,
            repo_type="space",
            operations=[CommitOperationDelete(path_in_repo=path) for path in stale_paths],
            commit_message="Prune non-leaderboard result artifacts",
        )
        print(f"Pruned {len(stale_paths)} non-leaderboard artifacts from current remote revision")
        return True
    except Exception as exc:
        print(f"Could not prune remote non-leaderboard artifacts: {exc}", file=sys.stderr)
        return False


def update_push_status(results: Path, *, ok: bool, detail: str = "") -> None:
    status_path = results / "online_status.json"
    if not status_path.exists():
        return
    try:
        payload = json.loads(status_path.read_text())
    except (json.JSONDecodeError, OSError):
        return
    payload["push_status"] = "ok" if ok else "failed"
    payload["pushed_at"] = datetime.now(timezone.utc).isoformat()
    if detail:
        payload["push_detail"] = detail
    status_path.write_text(json.dumps(payload, indent=4) + "\n")
    # Evaluation rewrites online_status while rapid cycles wait for publication.
    # Keep the last actual publish attempt independently for the watchdog.
    receipt = results / "publication_status.json"
    temporary = receipt.with_suffix(".tmp")
    temporary.write_text(json.dumps({key: payload[key] for key in ("push_status", "pushed_at")}))
    temporary.replace(receipt)


def restart_space(token: str, repo_id: str) -> None:
    api = get_hf_api(token)
    try:
        api.restart_space(repo_id)
        print(f"Restarted Space {repo_id} to load new results", file=sys.stderr)
    except Exception as exc:
        print(f"Space restart failed (non-fatal): {exc}", file=sys.stderr)


def maybe_restart_space(args: argparse.Namespace, token: str) -> None:
    """Restart only when explicitly requested by a full deployment workflow."""

    if args.restart_space:
        restart_space(token, args.repo)


def stage_public_results(results: Path, destination: Path) -> None:
    from huggingface_hub.utils import filter_repo_objects
    candidates = [p.relative_to(results).as_posix() for p in results.rglob("*") if p.is_file()]
    for rel in filter_repo_objects(candidates, ignore_patterns=UPLOAD_IGNORE_PATTERNS):
        if Path(rel).suffix not in {".csv", ".json", ".tex"}:
            continue
        output = destination / rel
        output.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(results / rel, output)


def push_via_api(args: argparse.Namespace, token: str) -> int:
    try:
        from huggingface_hub import HfApi
    except ImportError:
        print("huggingface_hub not installed; run setup_portable.sh", file=sys.stderr)
        return 1

    api = get_hf_api(token)
    results_repo = getattr(args, "results_repo", None) or args.repo + "-results"
    api.create_repo(results_repo, repo_type="dataset", exist_ok=True)
    last_error = ""

    for attempt in range(1, API_RETRIES + 1):
        try:
            # Refuse an old clone replacing a newer operator's cumulative history.
            from huggingface_hub import hf_hub_download
            from huggingface_hub.utils import EntryNotFoundError
            import pandas as pd
            revision = api.dataset_info(results_repo).sha
            metrics = pd.read_csv(args.results / "aggregates/live_release_metrics.csv")
            history_rows = len(metrics)
            try:
                previous_path = hf_hub_download(results_repo, "results/online_status.json",
                                                repo_type="dataset", revision=revision, token=token)
                previous = json.loads(Path(previous_path).read_text())
            except EntryNotFoundError:
                previous = {}
            if int(previous.get("canonical_history_rows", 0)) > history_rows:
                print("Remote history is newer/larger. Restore the latest operator backup before publishing.", file=sys.stderr)
                return 1
            with tempfile.TemporaryDirectory(prefix="livehouse-results-") as temp:
                stage = Path(temp)
                stage_public_results(args.results, stage)
                # The public status and tables belong to the same Hub commit.
                update_push_status(stage, ok=True, detail="HF Hub API")
                status_path = stage / "online_status.json"
                status = json.loads(status_path.read_text())
                status["canonical_history_rows"] = history_rows
                status["latest_scored_release_time"] = str(metrics["release_time"].max()) if history_rows else None
                status_path.write_text(json.dumps(status, indent=2) + "\n")
                commit = api.upload_folder(
                    folder_path=str(stage), path_in_repo="results", repo_id=results_repo,
                    repo_type="dataset", revision="main", parent_commit=revision, commit_message=args.commit_message,
                )
            print(f"Pushed results → https://huggingface.co/datasets/{results_repo}/commit/{commit.oid}")
            update_push_status(args.results, ok=True, detail="HF Hub API")
            maybe_restart_space(args, token)
            return 0
        except Exception as exc:
            last_error = str(exc)
            print(
                f"HF Hub API upload retry {attempt}/{API_RETRIES}: {exc}",
                file=sys.stderr,
            )
            time.sleep(3 * attempt)

    print(f"HF Hub API upload failed after {API_RETRIES} retries: {last_error}", file=sys.stderr)
    return 1


def main() -> int:
    args = parse_args()
    if not args.results.is_dir():
        print(f"Results directory not found: {args.results}", file=sys.stderr)
        return 1

    token = load_token()
    code = push_via_api(args, token)
    if code != 0:
        update_push_status(args.results, ok=False, detail="HF Hub API upload failed")
    return code


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