"""Run audio clips through the Inter-2 Audio playground's /generate API with one of its task presets.

  Emotion   one primary emotion out of seven; free output by default (--constrain for strict JSON)
  Combined  engagement, social signals, emotion and delivery in one answer; constrained by default

The prompts and each task's schema default come from the Space's own tasks.py, so a run sends
exactly what the UI sends. One JSON line per clip is appended to --output; re-running skips
clips already there, so a run that hits your daily ZeroGPU quota can be resumed.

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_inter2_audio.py --task Combined --output combined.jsonl clips/ one.wav
"""

from __future__ import annotations

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

from gradio_client import Client, handle_file

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))  # the Space folder, for tasks.py
from tasks import TASKS  # noqa: E402

SPACE = "Interhuman/inter2-audio-playground"
AUDIO_SUFFIXES = {".wav", ".mp3", ".flac", ".ogg", ".m4a"}


def audio_files(paths: list[str]) -> list[Path]:
    files: list[Path] = []
    for raw in paths:
        path = Path(raw)
        if path.is_dir():
            files += sorted(p for p in path.rglob("*") if p.suffix.lower() in AUDIO_SUFFIXES)
        elif path.is_file():
            files.append(path)
        else:
            raise SystemExit(f"Not found: {path}")
    return files


def one_line(task: str, parsed: dict | None) -> str:
    """A short console summary of the parsed answer."""
    if not isinstance(parsed, dict):
        return "(unparsed)"
    if task == "Emotion":
        return f"{parsed.get('primary_emotion', '?')} ({parsed.get('confidence', '?')})"
    emotion = (parsed.get("emotions") or [{}])[0]
    signals = ", ".join(f"{s.get('name')} ({s.get('probability')})" for s in parsed.get("signals") or []) or "none"
    return (
        f"{parsed.get('engagement_status', '?')} | emotion {emotion.get('name', '?')} | "
        f"delivery {parsed.get('delivery', '?')} | signals {signals}"
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("audio", nargs="+", help="audio files and/or folders (searched recursively)")
    parser.add_argument("--task", choices=list(TASKS), default="Emotion")
    parser.add_argument("--output", required=True, help="JSONL file; existing rows are kept and skipped")
    parser.add_argument(
        "--constrain",
        action=argparse.BooleanOptionalAction,
        default=None,
        help="force the task's JSON schema (default: the task's own default, as in the UI)",
    )
    parser.add_argument("--max-new-tokens", type=int, default=768)
    parser.add_argument("--space", default=SPACE)
    args = parser.parse_args()

    task = TASKS[args.task]
    constrain = task.constrain if args.constrain is None else args.constrain
    out = Path(args.output)
    done = {json.loads(line)["file"] for line in out.read_text().splitlines()} if out.exists() else set()
    files = [f for f in audio_files(args.audio) if str(f) not in done]
    print(f"{args.task} (constrain={constrain}): {len(done)} clips already done, {len(files)} 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, path in enumerate(files, 1):
            start = time.time()
            _, parsed_json, raw = client.predict(
                handle_file(str(path)), task.system_prompt, task.user_prompt, True, constrain,
                args.max_new_tokens, 0.0, 1.0, 1.0, args.task,
                api_name="/generate",
            )  # fmt: skip
            parsed = json.loads(parsed_json) if parsed_json else None
            record = {"file": str(path), "task": args.task, "constrained": constrain, "parsed": parsed}
            f.write(json.dumps({**record, "raw_output": raw}, ensure_ascii=False) + "\n")
            f.flush()
            print(f"[{i}/{len(files)}] {path.name}: {one_line(args.task, parsed)}  ({time.time() - start:.1f}s)")


if __name__ == "__main__":
    main()