#!/usr/bin/env python
"""Call a project Fish TTS API with a zero-shot voice or saved reference ID."""

from __future__ import annotations

import argparse
import base64
from pathlib import Path

import requests


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--url", default="http://127.0.0.1:8080/v1/tts")
    parser.add_argument("--text", required=True)
    voice = parser.add_mutually_exclusive_group()
    voice.add_argument("--reference-audio", type=Path)
    voice.add_argument("--reference-id")
    parser.add_argument("--reference-text")
    parser.add_argument("--output", type=Path, default=Path("output.wav"))
    parser.add_argument("--api-key")
    parser.add_argument("--seed", type=int)
    parser.add_argument("--max-new-tokens", type=int, default=1024)
    parser.add_argument("--chunk-length", type=int, default=200)
    parser.add_argument("--top-p", type=float, default=0.9)
    parser.add_argument("--temperature", type=float, default=0.9)
    parser.add_argument("--repetition-penalty", type=float, default=1.1)
    parser.add_argument(
        "--normalize",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Enable or disable API loudness normalization",
    )
    parser.add_argument(
        "--use-memory-cache",
        choices=("on", "off"),
        default="on",
        help="Reuse cached reference codes/text, or force a fresh reference encode",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    references = []
    if args.reference_audio is not None:
        if not args.reference_audio.is_file():
            raise SystemExit(f"Reference audio is missing: {args.reference_audio}")
        if not args.reference_text or not args.reference_text.strip():
            raise SystemExit("--reference-text is required with --reference-audio")
        references.append(
            {
                "audio": base64.b64encode(args.reference_audio.read_bytes()).decode(
                    "ascii"
                ),
                "text": args.reference_text,
            }
        )
    elif args.reference_text:
        raise SystemExit("--reference-text requires --reference-audio")

    payload = {
        "text": args.text,
        "references": references,
        "reference_id": args.reference_id,
        "format": "wav",
        "streaming": False,
        "normalize": args.normalize,
        "max_new_tokens": args.max_new_tokens,
        "chunk_length": args.chunk_length,
        "top_p": args.top_p,
        "temperature": args.temperature,
        "repetition_penalty": args.repetition_penalty,
        "seed": args.seed,
        "use_memory_cache": args.use_memory_cache,
    }
    headers = {"Accept": "audio/wav"}
    if args.api_key:
        headers["Authorization"] = f"Bearer {args.api_key}"
    response = requests.post(
        args.url,
        json=payload,
        headers=headers,
        timeout=600,
    )
    if response.status_code != 200:
        raise SystemExit(
            f"TTS request failed ({response.status_code}): {response.text[:2000]}"
        )
    content_type = response.headers.get("content-type", "")
    if "audio/" not in content_type and not response.content.startswith(b"RIFF"):
        raise SystemExit(f"Unexpected response content type: {content_type}")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_bytes(response.content)
    print(f"audio={args.output}")
    print(f"bytes={len(response.content)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())