#!/usr/bin/env python3
"""Mureka instrumental generator: submit -> poll -> download.

Usage:
    python3 tools/mureka_gen.py --prompt "..." --out /path/to/out_dir \
        [--model mureka-7.6] [--task-id <existing id>] [--format flac|mp3]

Requires env: MUREKA_API_KEY, MUREKA_API_BASE (default https://api.mureka.ai).

With --task-id, skips submission and just polls that task. Downloads all
choices into <out>/<task_id>_<choice_index>.<format> and prints a JSON summary.
"""
import argparse
import json
import os
import sys
import time
import urllib.request
import urllib.error

POLL_INTERVAL = 10
POLL_TIMEOUT = 15 * 60


def api(method, path, body=None):
    base = os.environ.get("MUREKA_API_BASE", "https://api.mureka.ai").rstrip("/")
    key = os.environ.get("MUREKA_API_KEY")
    if not key:
        sys.exit("MUREKA_API_KEY not set")
    req = urllib.request.Request(
        base + path,
        method=method,
        data=json.dumps(body).encode() if body is not None else None,
        headers={
            "Authorization": "Bearer " + key,
            "Content-Type": "application/json",
        },
    )
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.loads(r.read().decode())


def submit(prompt, model):
    resp = api("POST", "/v1/instrumental/generate", {"prompt": prompt, "model": model})
    tid = resp.get("id")
    if not tid:
        sys.exit("submit failed, no id in response: %s" % resp)
    return tid, resp


def query(tid):
    try:
        return api("GET", "/v1/instrumental/query/%s" % tid)
    except urllib.error.HTTPError:
        # fall back to the song endpoint
        return api("GET", "/v1/song/query/%s" % tid)


def poll(tid):
    deadline = time.time() + POLL_TIMEOUT
    while True:
        resp = query(tid)
        status = resp.get("status")
        print("poll %s: status=%s" % (tid, status), flush=True)
        if status == "succeeded":
            return resp
        if status in ("failed", "cancelled"):
            sys.exit("task %s ended with status=%s: %s" % (tid, status, resp))
        if time.time() > deadline:
            sys.exit("task %s timed out after %ds" % (tid, POLL_TIMEOUT))
        time.sleep(POLL_INTERVAL)


def download(url, path):
    with urllib.request.urlopen(url, timeout=120) as r, open(path, "wb") as f:
        while True:
            chunk = r.read(1 << 16)
            if not chunk:
                break
            f.write(chunk)
    return path


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--prompt")
    p.add_argument("--out", required=True, help="output directory")
    p.add_argument("--model", default="mureka-7.6")
    p.add_argument("--task-id", help="poll existing task instead of submitting")
    p.add_argument("--format", choices=["flac", "mp3"], default="flac")
    args = p.parse_args()

    os.makedirs(args.out, exist_ok=True)

    if args.task_id:
        tid = args.task_id
    else:
        if not args.prompt:
            sys.exit("--prompt required when submitting")
        tid, resp = submit(args.prompt, args.model)
        print("submitted task id=%s status=%s" % (tid, resp.get("status")), flush=True)

    result = poll(tid)

    summary = {"task_id": tid, "files": []}
    for i, ch in enumerate(result.get("choices", [])):
        url = ch.get("flac_url") if args.format == "flac" else ch.get("url")
        if not url:
            url = ch.get("url")
        ext = "flac" if "flac" in (url or "") else "mp3"
        path = os.path.join(args.out, "%s_%d.%s" % (tid, i, ext))
        download(url, path)
        summary["files"].append(
            {"path": path, "duration_ms": ch.get("duration"), "choice_id": ch.get("id")}
        )
        print("downloaded %s (%sms)" % (path, ch.get("duration")), flush=True)

    print("SUMMARY " + json.dumps(summary))


if __name__ == "__main__":
    main()