"""Generate Sonilo text-to-music output and save output.m4a.

Run:
    export SONILO_API_KEY=sk_...
    python text-to-music-output.py

This script writes output.m4a only after real audio_chunk bytes are received and
a complete event arrives. It never creates placeholder media files.
"""

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

import requests

API_URL = "https://api.sonilo.com/v1/text-to-music"
OUTPUT_FILE = Path("output.m4a")
PROMPT = (
    "30 seconds of upbeat electronic instrumental background music for a "
    "product launch video, no vocals."
)
USER_AGENT = "SoniloPythonExample/1.0 (+https://platform.sonilo.com/docs)"


def retry_after_seconds(resp, attempt):
    value = resp.headers.get("Retry-After")
    if value:
        try:
            return max(1, int(value))
        except ValueError:
            return 5
    return min(2**attempt, 30)


def classify_error(resp):
    body = resp.text
    code = f"http_{resp.status_code}"
    message = body
    try:
        payload = resp.json()
        code = payload.get("code", code)
        message = payload.get("message", body)
    except ValueError:
        pass

    if resp.status_code == 402 and code == "payment_required":
        return "insufficient_balance", (
            "insufficient_balance: add balance or use an eval key with enough "
            f"credits. API returned payment_required: {message}"
        )
    return code, f"Sonilo request failed with HTTP {resp.status_code}: {message}"


def open_generation(api_key):
    headers = {
        "Authorization": f"Bearer {api_key}",
        "User-Agent": USER_AGENT,
    }
    files = {
        "prompt": (None, PROMPT),
        "duration": (None, "30"),
    }

    for attempt in range(1, 6):
        resp = requests.post(
            API_URL,
            headers=headers,
            files=files,
            stream=True,
            timeout=300,
        )
        if resp.status_code != 429:
            return resp
        delay = retry_after_seconds(resp, attempt)
        resp.close()
        print(f"Rate limited. Retrying after {delay}s...", file=sys.stderr)
        time.sleep(delay)

    raise RuntimeError("rate_limit_exceeded: still receiving 429 after retries")


def main():
    api_key = os.environ.get("SONILO_API_KEY")
    if not api_key:
        raise RuntimeError("Set SONILO_API_KEY before running this script")

    chunks = []
    complete = False
    with open_generation(api_key) as resp:
        if not resp.ok:
            _, message = classify_error(resp)
            raise RuntimeError(message)

        for line in resp.iter_lines(decode_unicode=True):
            if not line:
                continue
            event = json.loads(line)
            event_type = event.get("type")
            if event_type == "audio_chunk":
                chunks.append(base64.b64decode(event["data"]))
            elif event_type == "complete":
                complete = True
                break
            elif event_type == "error":
                code = event.get("code", "unknown_error")
                message = event.get("message", "Sonilo generation failed")
                raise RuntimeError(f"{code}: {message}")

    if not complete:
        raise RuntimeError("stream_incomplete: no complete event received")
    if not chunks:
        raise RuntimeError("empty_audio: complete arrived without audio_chunk data")

    OUTPUT_FILE.write_bytes(b"".join(chunks))
    print(f"Saved {OUTPUT_FILE} ({OUTPUT_FILE.stat().st_size} bytes)")


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        print(f"Error: {exc}", file=sys.stderr)
        sys.exit(1)