"""Server-side song example. Keep the job file private and reuse it to resume."""
import json
import os
import sys
import time
import uuid
from pathlib import Path
from urllib.parse import quote
from urllib.request import Request, urlopen
from urllib.error import HTTPError

API = "https://api.omnapi.com"
TERMINAL = {"COMPLETED", "FAILED", "CANCELLED"}


def run_song_job(job, save, request, now=time.monotonic, sleep=time.sleep, timeout_seconds=2700):
    if not job.get("taskId"):
        save(job)  # Persist the idempotency key and exact body before a paid call.
        receipt = request("/api/v1/suno/songs", job["input"], job["requestId"])
        if not isinstance(receipt.get("taskId"), str) or not receipt["taskId"]:
            raise RuntimeError("No taskId in receipt; retain the request ID and investigate.")
        job["taskId"] = receipt["taskId"]
        job["generationId"] = receipt.get("generationId")
        save(job)
    deadline = now() + timeout_seconds
    while now() < deadline:
        task = request("/api/v1/tasks/" + quote(job["taskId"], safe=""))
        if not isinstance(task.get("status"), str):
            raise RuntimeError("Invalid task response; retain this job file.")
        if task["status"] in TERMINAL:
            job["result"] = task
            save(job)
            return task
        sleep(min(5, max(0, deadline - now())))
    raise TimeoutError("Polling deadline reached. Run again with the same job file to resume.")


def main():
    api_key = os.environ.get("OMNAPI_KEY")
    if not api_key:
        raise RuntimeError("Set OMNAPI_KEY on this server.")
    filename = Path(sys.argv[1] if len(sys.argv) > 1 else "suno-job.json").resolve()
    try:
        job = json.loads(filename.read_text(encoding="utf-8"))
    except FileNotFoundError:
        job = {"requestId": str(uuid.uuid4()), "input": {
            "mode": "simple", "model": "chirp-hawk",
            "prompt": os.environ.get("OMNAPI_PROMPT", "upbeat city pop, bright guitars, summer night drive"),
        }}
        with os.fdopen(os.open(filename, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "w", encoding="utf-8") as stream:
            json.dump(job, stream, indent=2)
    if not isinstance(job.get("requestId"), str) or not job["requestId"] or not isinstance(job.get("input"), dict):
        raise RuntimeError("Invalid job file. Do not overwrite an uncertain paid request.")

    def save(state):
        temporary = filename.with_name(filename.name + "." + str(uuid.uuid4()) + ".tmp")
        with os.fdopen(os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "w", encoding="utf-8") as stream:
            json.dump(state, stream, indent=2)
        os.replace(temporary, filename)

    def request(path, data=None, request_id=None):
        headers = {"x-api-key": api_key}
        body = None
        if data is not None:
            headers.update({"Content-Type": "application/json", "Idempotency-Key": request_id})
            body = json.dumps(data).encode("utf-8")
        try:
            with urlopen(Request(API + path, data=body, headers=headers), timeout=30) as response:
                return json.load(response)
        except HTTPError as error:
            # Do not print raw API error bodies or automatically retry a paid write.
            retry = "; observe Retry-After before resuming" if error.headers.get("Retry-After") else ""
            raise RuntimeError(f"HTTP {error.code}{retry}. Retain the same job file and request ID.") from None

    result = run_song_job(job, save, request)
    print(f"Task {job['taskId']}: {result['status']}. Result saved in the private job file.")
    return 0 if result["status"] == "COMPLETED" else 1


if __name__ == "__main__":
    try:
        sys.exit(main())
    except Exception as error:
        print(str(error), file=sys.stderr)
        sys.exit(1)