#!/usr/bin/env bash
# One-shot: pull prebuilt vidaio compression image from Hugging Face and start the lab.
#
#   export HF_TOKEN=hf_xxx          # required if the HF repo is private
#   bash setup.sh
#
# Env:
#   EVAL_HF_IMAGE_REPO   default Realfencer/vidaio-compression-eval
#   EVAL_PORT            default 8081
#   EVAL_DIR             install dir (default: ./vidaio-compression-lab)
#   SKIP_SUITES=1        skip downloading challenge suites
#   EVAL_HF_REPO         suite pack repo (default Realfencer/test)
set -euo pipefail

REPO="${EVAL_HF_IMAGE_REPO:-Realfencer/vidaio-comp-eval-img}"
SUITE_REPO="${EVAL_HF_REPO:-Realfencer/test}"
PORT="${EVAL_PORT:-8081}"
DIR="${EVAL_DIR:-$PWD/vidaio-compression-lab}"
TAG="${EVAL_IMAGE_TAG:-latest}"
ARCHIVE="vidaio-compression-eval-${TAG}.tar.gz"

echo "==> Vidaio compression lab setup"
echo "    image repo : $REPO"
echo "    suite repo : $SUITE_REPO"
echo "    install dir: $DIR"
echo "    port       : $PORT"

if ! command -v docker >/dev/null 2>&1; then
  echo "ERROR: docker is required. Install Docker Engine, then re-run." >&2
  exit 1
fi

PY="${PYTHON:-python3}"
if ! "$PY" -c "import huggingface_hub" 2>/dev/null; then
  echo "==> Installing huggingface_hub"
  "$PY" -m pip install -q "huggingface_hub>=0.23.0"
fi

mkdir -p "$DIR/data/suites" "$DIR/data/evals"
cd "$DIR"

echo "==> Downloading image bundle from Hugging Face"
"$PY" - <<PY
from huggingface_hub import hf_hub_download
import os
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
for name in ("$ARCHIVE", "docker-compose.yml"):
    path = hf_hub_download(
        repo_id="$REPO",
        filename=name,
        repo_type="model",
        local_dir=".",
        token=token,
    )
    print("got", path)
PY

if [[ ! -f "$ARCHIVE" ]]; then
  echo "ERROR: missing $ARCHIVE after download" >&2
  exit 1
fi

echo "==> Loading docker image (this can take a few minutes)"
gunzip -c "$ARCHIVE" | docker load

if ! docker image inspect vidaio-compression-eval:latest >/dev/null 2>&1; then
  echo "ERROR: vidaio-compression-eval:latest not present after docker load" >&2
  exit 1
fi

if [[ "${SKIP_SUITES:-0}" != "1" ]]; then
  echo "==> Pulling challenge suites from $SUITE_REPO"
  "$PY" - <<PY
from huggingface_hub import snapshot_download
import os, shutil
from pathlib import Path
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
dest = Path("data/suites")
dest.mkdir(parents=True, exist_ok=True)
# Prefer suites/ subfolder if present in the pack repo.
try:
    cached = snapshot_download(
        repo_id="$SUITE_REPO",
        repo_type="model",
        token=token,
        allow_patterns=["suites/**", "**/manifest.json", "**/*_ref.mp4", "**/*.json"],
    )
except Exception as e:
    print("suite download warning:", e)
    raise
cached_p = Path(cached)
# Copy suite-* dirs into data/suites
copied = 0
for root in [cached_p / "suites", cached_p]:
    if not root.exists():
        continue
    for suite in root.glob("suite-*"):
        if not suite.is_dir():
            continue
        target = dest / suite.name
        if target.exists():
            shutil.rmtree(target)
        shutil.copytree(suite, target)
        copied += 1
print(f"suites ready: {copied}")
if copied < 1:
    raise SystemExit("no suites found in HF pack — set SKIP_SUITES=1 to skip")
PY
fi

echo "==> Starting compression-eval on :$PORT"
EVAL_PORT="$PORT" docker compose up -d

echo ""
echo "Ready → http://127.0.0.1:${PORT}"
echo "Pick a suite, set AV1 · CRF · VMAF≥85, Run eval."
echo "Encoder + ML model are baked into the image."