"""Example client for the OmniVoice ZeroGPU batch endpoint (Gradio API).

Usage:
    pip install gradio_client
    python client_example.py salmanbvps/omnivoice-batch-tts [ref.wav "ref transcript"]

Sends a batch (auto + voice-design, and -- if a ref wav is given -- voice-clone)
in a single API call and writes out_0.wav, out_1.wav, ...
"""

import base64
import json
import sys

from gradio_client import Client


def b64_file(path: str) -> str:
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode("ascii")


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    space = sys.argv[1]  # "owner/space" or full URL

    items = [
        {"text": "Hello from OmniVoice, generated in a single batched call.",
         "language": "en"},
        {"text": "This voice was designed from a text instruction.",
         "language": "en", "instruct": "female, british accent"},
    ]
    if len(sys.argv) >= 3:
        items.append({
            "text": "And this one is cloned from your reference audio.",
            "language": "en",
            "ref_audio_b64": b64_file(sys.argv[2]),
            "ref_text": sys.argv[3] if len(sys.argv) >= 4 else None,
        })

    payload = {"items": items, "num_step": 32, "guidance_scale": 2.0,
               "audio_format": "wav"}

    client = Client(space)
    result = client.predict(json.dumps(payload), api_name="/batch")
    data = result if isinstance(result, dict) else json.loads(result)

    if "error" in data:
        print("Error:", data["error"])
        sys.exit(1)

    for r in data["results"]:
        i = r["index"]
        if r["status"] != "success":
            print(f"[{i}] ERROR: {r.get('error')}")
            continue
        out = f"out_{i}.wav"
        with open(out, "wb") as f:
            f.write(base64.b64decode(r["audio_b64"]))
        print(f"[{i}] {r['duration']}s -> {out}")


if __name__ == "__main__":
    main()