"""Run VocalAffectBench (besimple-ai/vocal-affect-bench) through the Inter-2 Audio playground's /generate API.

Downloads the benchmark at a pinned revision, sends each clip with the Emotion task's prompt (the
benchmark's own prompt, audio before text) and appends one JSON line per clip to --output, then
prints the accuracy overall and per emotion. Re-running skips clips already in --output, so a run
that hits your daily ZeroGPU quota can be resumed the next day.

--constrain (default) forces the benchmark's answer object: the label limited to its 7 emotions,
compact JSON, confidence in quarters and evidence left empty (only the label is scored).
--no-constrain lets the model write freely.

Usage (from a download of this Space's files; HF_TOKEN is your Hugging Face token):
    pip install gradio_client huggingface_hub  # in a fresh virtual environment
    HF_TOKEN=hf_... python examples/run_vab_bench.py --output vab.jsonl [--limit 10]
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from collections import Counter
from pathlib import Path

from gradio_client import Client, handle_file
from huggingface_hub import snapshot_download

SPACE = "Interhuman/inter2-audio-playground"
DATASET = "besimple-ai/vocal-affect-bench"
REVISION = "8d4165f514befc282dc20360b2d28522d9523413"
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))  # the Space folder, for tasks.py
from tasks import TASKS  # noqa: E402

PROMPT = TASKS["Emotion"].user_prompt


def load_rows() -> list[dict]:
    root = snapshot_download(
        DATASET, repo_type="dataset", revision=REVISION, allow_patterns=["data/metadata.jsonl", "audio/*.wav"]
    )
    rows = [json.loads(line) for line in Path(root, "data", "metadata.jsonl").read_text(encoding="utf-8").splitlines()]
    for row in rows:
        row["audio_path"] = str(Path(root, row.get("file_name") or f"audio/{row['audio_id']}.wav"))
    return rows


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--output", required=True, help="JSONL file; existing rows are kept and skipped")
    parser.add_argument("--space", default=SPACE)
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="only the first N clips (sorted by emotion, so a smoke test, not a sample)",
    )
    parser.add_argument("--constrain", action=argparse.BooleanOptionalAction, default=True)
    args = parser.parse_args()

    out = Path(args.output)
    done = {json.loads(line)["audio_id"] for line in out.read_text().splitlines()} if out.exists() else set()
    rows = [r for r in load_rows() if r["audio_id"] not in done][: args.limit]
    print(f"{len(done)} clips already done, {len(rows)} to run", flush=True)

    client = Client(args.space, token=os.getenv("HF_TOKEN"), verbose=False)
    with out.open("a", encoding="utf-8") as f:
        for i, row in enumerate(rows, 1):
            start = time.time()
            _, parsed, raw = client.predict(
                handle_file(row["audio_path"]), "", PROMPT, True, args.constrain, 768, 0.0, 1.0, 1.0, "Emotion",
                api_name="/generate",
            )  # fmt: skip
            try:
                label = str(json.loads(parsed).get("primary_emotion", "")).strip().lower()
            except (json.JSONDecodeError, AttributeError):
                label = ""
            record = {"audio_id": row["audio_id"], "required_emotion": row["required_emotion"], "predicted": label}
            f.write(json.dumps({**record, "raw_output": raw}) + "\n")
            f.flush()
            mark = "ok" if label == row["required_emotion"] else "x "
            print(f"[{i}/{len(rows)}] {mark} {row['audio_id']}: {row['required_emotion']} -> {label or '?'}"
                  f"  ({time.time() - start:.1f}s)", flush=True)  # fmt: skip

    results = [json.loads(line) for line in out.read_text().splitlines()]
    correct = sum(r["predicted"] == r["required_emotion"] for r in results)
    print(f"\nAccuracy: {correct}/{len(results)} = {correct / len(results):.1%}")
    per_class = Counter(r["required_emotion"] for r in results if r["predicted"] == r["required_emotion"])
    totals = Counter(r["required_emotion"] for r in results)
    print("  " + "  ".join(f"{e} {per_class[e]}/{n}" for e, n in sorted(totals.items())))


if __name__ == "__main__":
    main()