#!/usr/bin/env bash

set -euo pipefail

PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
DATASET_ID="pipecat-ai/smart-turn-data-v3.2-test"
PINNED_DATASET_REVISION="0500378e8ed6d38e37b016e24d261e8e6c6a6859"
DATASET_REVISION="${DATASET_REVISION:-${PINNED_DATASET_REVISION}}"
TARGET_DIR="${PROJECT_ROOT}/data/raw/smart-turn-data-v3.2-test"
EXPECTED_PARQUET_SHARDS=10
EXPECTED_ROWS=31527

if [[ "${DATASET_REVISION}" != "${PINNED_DATASET_REVISION}" ]]; then
  echo "Refusing non-pinned official-test revision: ${DATASET_REVISION}" >&2
  exit 2
fi
if [[ "$#" -ne 0 ]]; then
  echo "This sealed downloader accepts no extra hf CLI arguments." >&2
  exit 2
fi

if [[ -z "${PYTHON_BIN:-}" ]]; then
  if [[ -x "${PROJECT_ROOT}/.venv/bin/python" ]]; then
    PYTHON_BIN="${PROJECT_ROOT}/.venv/bin/python"
  else
    PYTHON_BIN="python3"
  fi
fi

if [[ "${UNSEAL_OFFICIAL_TEST:-}" != "I_HAVE_FROZEN_MODEL_AND_THRESHOLD" ]]; then
  cat >&2 <<'MESSAGE'
The official test set is sealed to prevent tuning leakage.
Freeze the model, preprocessing, controller, and threshold first, then run:

  UNSEAL_OFFICIAL_TEST=I_HAVE_FROZEN_MODEL_AND_THRESHOLD bash scripts/download_test_dataset.sh
MESSAGE
  exit 2
fi

if [[ -z "${FROZEN_MANIFEST:-}" ]]; then
  echo "Set FROZEN_MANIFEST to the output of scripts/freeze_candidate.py." >&2
  exit 2
fi
"${PYTHON_BIN}" "${PROJECT_ROOT}/scripts/verify_freeze.py" \
  --manifest "${FROZEN_MANIFEST}"

if [[ -f "${PROJECT_ROOT}/.env" ]]; then
  set -a
  # shellcheck disable=SC1091
  source "${PROJECT_ROOT}/.env"
  set +a
fi

export HF_TOKEN="${HF_TOKEN:-${hf_token:-}}"
if [[ -z "${HF_TOKEN}" ]]; then
  echo "Missing Hugging Face token. Set HF_TOKEN or hf_token in .env." >&2
  exit 1
fi
unset hf_token

export UV_CACHE_DIR="${UV_CACHE_DIR:-${PROJECT_ROOT}/.cache/uv}"
export UV_TOOL_DIR="${UV_TOOL_DIR:-${PROJECT_ROOT}/.cache/uv-tools}"
export UV_TOOL_BIN_DIR="${UV_TOOL_BIN_DIR:-${PROJECT_ROOT}/.cache/uv-bin}"
export HF_XET_HIGH_PERFORMANCE="${HF_XET_HIGH_PERFORMANCE:-1}"

mkdir -p "${TARGET_DIR}"
if [[ -n "${HF_CLI_BIN:-}" ]]; then
  if [[ ! -x "${HF_CLI_BIN}" ]]; then
    echo "HF_CLI_BIN is not executable: ${HF_CLI_BIN}" >&2
    exit 1
  fi
  hf_command=("${HF_CLI_BIN}")
elif command -v hf >/dev/null 2>&1; then
  hf_command=("$(command -v hf)")
else
  hf_command=(uvx --from 'huggingface_hub>=0.27,<2' hf)
fi

"${hf_command[@]}" download \
  "${DATASET_ID}" \
  --repo-type dataset \
  --revision "${DATASET_REVISION}" \
  --local-dir "${TARGET_DIR}"

"${PYTHON_BIN}" - "${TARGET_DIR}" "${EXPECTED_PARQUET_SHARDS}" "${EXPECTED_ROWS}" <<'PY'
from __future__ import annotations

import hashlib
import sys
from pathlib import Path

try:
    import pyarrow.parquet as parquet
except ImportError as exc:
    raise SystemExit("Snapshot verification requires pyarrow; install the project first.") from exc

target = Path(sys.argv[1]).resolve()
expected_shard_count = int(sys.argv[2])
expected_rows = int(sys.argv[3])

# Immutable upstream Git-LFS SHA-256 oids and byte sizes for revision
# 0500378e8ed6d38e37b016e24d261e8e6c6a6859. These are content hashes, not
# locally generated expectations.
expected = {
    "data/train-00000-of-00010.parquet": (
        486_502_678,
        "a87c75806b814ee7379998b6f9dc65a6433c01bfec2875e62c5d1ccd2b37257a",
    ),
    "data/train-00001-of-00010.parquet": (
        489_429_742,
        "c408bd3b31cc3cb907280fa5d3186f0f5ba08c6beb84532a3828d685b967b7d8",
    ),
    "data/train-00002-of-00010.parquet": (
        479_920_042,
        "2b50ff3346f8aecc6c4b0b706b593c5b3b174b73ce84c795b515b68e6abb3788",
    ),
    "data/train-00003-of-00010.parquet": (
        486_565_988,
        "0afd86b7d1cdf03ffb804a00278fb76477d29ef008b71ad1d7a746bb1b25850c",
    ),
    "data/train-00004-of-00010.parquet": (
        477_831_330,
        "4c600774512010880f72dd86cc8abe5d46df594615e472351bd270cc94e1ff66",
    ),
    "data/train-00005-of-00010.parquet": (
        495_015_303,
        "b4537a3b96498481b98c5d60b8d84ad05109ae0acbeafc18759b15ff1d0d9335",
    ),
    "data/train-00006-of-00010.parquet": (
        481_054_312,
        "ef0eb0085b55e05fc5594f035c8011afb78d039e1b3e50c579a96e9295d9acec",
    ),
    "data/train-00007-of-00010.parquet": (
        478_266_840,
        "eddd1db2f95fff2f08e18ca1fe73fe0c5d9eb8030bea97da2d109297a0158b67",
    ),
    "data/train-00008-of-00010.parquet": (
        479_859_534,
        "96314dc8bb77515a5d1d02e8cb7c1410da54169d33d6f73ca2f87738ec1269f0",
    ),
    "data/train-00009-of-00010.parquet": (
        483_305_860,
        "769283c79bea4ae6eebdfe7d09fe481f8a154c53b5e3c29fe2c71d7b55dfc862",
    ),
}
if len(expected) != expected_shard_count:
    raise SystemExit("Internal error: expected shard manifest has the wrong length.")

actual = {
    path.relative_to(target).as_posix()
    for path in target.rglob("*.parquet")
    if path.is_file()
}
expected_names = set(expected)
if actual != expected_names:
    missing = sorted(expected_names - actual)
    unexpected = sorted(actual - expected_names)
    raise SystemExit(
        "Incomplete test snapshot: exact shard inventory mismatch; "
        f"missing={missing}, unexpected={unexpected}."
    )

total_rows = 0
for relative, (expected_bytes, expected_sha256) in expected.items():
    path = target / relative
    if path.is_symlink() or not path.is_file():
        raise SystemExit(f"Invalid test shard (must be a regular file): {relative}")
    if path.stat().st_size != expected_bytes:
        raise SystemExit(f"Test shard byte-size mismatch: {relative}")
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    if digest.hexdigest() != expected_sha256:
        raise SystemExit(f"Test shard SHA-256 mismatch: {relative}")
    try:
        total_rows += int(parquet.ParquetFile(path).metadata.num_rows)
    except Exception as exc:
        raise SystemExit(f"Cannot read Parquet metadata for {relative}: {exc}") from exc

if total_rows != expected_rows:
    raise SystemExit(
        f"Incomplete test snapshot: expected {expected_rows} rows, found {total_rows}."
    )
print(
    f"Verified {len(expected)} exact Parquet shards, {total_rows} rows, "
    "and all immutable SHA-256 hashes."
)
PY

echo "Verified sealed test snapshot at revision ${DATASET_REVISION}."