#!/usr/bin/env python3
"""
Fetch the (private) eval data into a local dir before Streamlit starts serving.

The Space repo is PUBLIC, so the eval data is deliberately NOT committed to it.
Instead the data lives in a private HF dataset and is pulled at container
startup using the `HF_TOKEN` Space secret. The download is targeted — only the
two file types the dashboard reads, for the one subfolder — so it is a small
(~tens of MB) fetch that completes in seconds, well before the first visitor.

Because this runs in the Docker CMD *before* `streamlit run`, no visitor ever
sees a partial download: by the time the app serves, DATA_DIR is fully on disk
and `_resolve_eval_dir()` reads it as a local path (no network on the request
path). On a cold restart this re-fetches, but the targeted snapshot is fast.

Env / secrets:
  HF_TOKEN             — required to read the private dataset
  HF_REPO_ID           — dataset repo (default below)
  HF_DEFAULT_BRANCH    — branch/revision to pull
  HF_DEFAULT_SUBFOLDER — subfolder within the branch to pull
  DATA_DIR             — where to write the data (default /app/data)
"""

import os
import shutil
import sys


def main() -> None:
    repo_id = os.environ.get("HF_REPO_ID", "srmdtranslations/model_gu_en_evaluations")
    branch = os.environ.get("HF_DEFAULT_BRANCH", "main")
    subfolder = os.environ.get("HF_DEFAULT_SUBFOLDER", "")
    token = os.environ.get("HF_TOKEN") or None
    data_dir = os.environ.get("DATA_DIR", "/app/data")

    if not token:
        print("[prewarm] HF_TOKEN not set — skipping fetch (app will try at runtime).",
              flush=True)
        return
    if not subfolder:
        print("[prewarm] HF_DEFAULT_SUBFOLDER not set — skipping fetch.", flush=True)
        return

    try:
        from huggingface_hub import snapshot_download

        print(f"[prewarm] fetching {repo_id}@{branch}/{subfolder} → {data_dir} …",
              flush=True)
        # Only the two file types the dashboard reads — keeps the fetch small/fast.
        local_root = snapshot_download(
            repo_id=repo_id,
            repo_type="dataset",
            revision=branch,
            token=token,
            allow_patterns=[
                f"{subfolder}/**/*.holistic.jsonl",
                f"{subfolder}/**/*.intent_entity.jsonl",
            ],
        )
        src = os.path.join(local_root, subfolder)
        if not os.path.isdir(src):
            print(f"[prewarm] subfolder missing after download: {src}", flush=True)
            return

        # Copy the matching files into DATA_DIR (resolve symlinks so the dir is
        # self-contained and readable without the HF cache layout).
        if os.path.exists(data_dir):
            shutil.rmtree(data_dir)
        n = 0
        for root, _, files in os.walk(src):
            for fname in files:
                if not (fname.endswith(".holistic.jsonl")
                        or fname.endswith(".intent_entity.jsonl")):
                    continue
                rel = os.path.relpath(os.path.join(root, fname), src)
                out = os.path.join(data_dir, rel)
                os.makedirs(os.path.dirname(out), exist_ok=True)
                shutil.copy(os.path.realpath(os.path.join(root, fname)), out)
                n += 1
        print(f"[prewarm] staged {n} files into {data_dir}.", flush=True)
    except Exception as exc:  # noqa: BLE001 — prewarm must never block startup
        print(f"[prewarm] skipped ({exc})", flush=True)


if __name__ == "__main__":
    main()