"""Create a new Seedance scene from a photo via Affogato's public API.

pip install requests
Set AFFOGATO_API_KEY, then run:
python complete_loop.py --photo noor-street.webp --submit
Run without --submit to upload, prepare and estimate only.
"""
import argparse
import json
import os
from pathlib import Path
import time
import uuid

import requests

BASE = "https://api.affogato.ai/v1"
MODEL = "seedance-2-5-ref2v-modelark"
PROMPT = (
    "Create a new six-second video starring Noor, the adult woman in #Image1. "
    "Use the photo only for her identity: preserve her face, curly updo, loose "
    "face-framing curls and hoop earrings. Compose a new scene, not the street "
    "photograph. Noor wears a floor-length satin evening gown at a black-tie "
    "reception. She walks toward the camera, pauses beside a tall floral "
    "arrangement and smiles at someone just beside the camera. Guests in "
    "tuxedos and evening dresses converse behind her. One continuous shot, "
    "camera moving gently backward, her face clearly visible. Quiet reception "
    "ambience, no distinct dialogue, music, text or overlays."
)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--photo", type=Path, default=Path("noor-street.webp"))
    parser.add_argument("--state", type=Path, default=Path("api-example-state.json"))
    parser.add_argument("--output", type=Path, default=Path("output.mp4"))
    parser.add_argument("--submit", action="store_true")
    parser.add_argument("--max-credits", type=int, default=1000)
    parser.add_argument("--wait-minutes", type=float, default=15)
    args = parser.parse_args()
    client = requests.Session()
    client.headers["Authorization"] = "Bearer " + os.environ["AFFOGATO_API_KEY"]

    def api(method, route, **kwargs):
        response = client.request(method, BASE + route, timeout=40, **kwargs)
        if not response.ok:
            # Error bodies contain the useful API error code, never the key.
            raise RuntimeError(f"HTTP {response.status_code}: {response.text[:1200]}")
        return response.json()

    state = json.loads(args.state.read_text()) if args.state.exists() else {}

    def save():
        temp = args.state.with_suffix(".tmp")
        temp.write_text(json.dumps(state, indent=2))
        temp.replace(args.state)

    if not state.get("request"):
        model = api("GET", "/models/" + MODEL)
        if "reference_image" not in model["inputs"]:
            raise RuntimeError("Model input changed; inspect /models before continuing")
        state["model_schema"] = model
        if not state.get("upload"):
            with args.photo.open("rb") as photo:
                state["upload"] = api("POST", "/files", files={
                    "file": (args.photo.name, photo, "image/webp")
                })
            save()
        state["request"] = {
            "model": MODEL,
            "input": {
                "reference_image": state["upload"]["id"],
                "prompt": PROMPT,
                "duration": 6,
                "resolution": "720p",
                "aspect_ratio": "9:16",
                "generate_audio": True,
            },
        }
        state["idempotency_key"] = str(uuid.uuid4())
        save()

    if not state.get("generation_id"):
        state["estimate"] = api("POST", "/generations/estimate", json=state["request"])
        save()
        quote = state["estimate"]
        print("Estimated credits:", quote["credits"], flush=True)
        if not args.submit:
            print("Estimate only. Add --submit to create the paid generation.", flush=True)
            return
        if (quote.get("credits_max") or quote["credits"]) > args.max_credits:
            raise RuntimeError("Estimate exceeds --max-credits")
        # Same body + same idempotency key safely retries an ambiguous POST.
        # Do not replay an ambiguous submission past the API's 24-hour window.
        if state.get("submitted_at") and time.time() - state["submitted_at"] > 23 * 3600:
            raise RuntimeError("Inspect generation history before retrying this old submission")
        state.setdefault("submitted_at", time.time())
        save()
        job = api("POST", "/generations", json=state["request"],
                  headers={"Idempotency-Key": state["idempotency_key"]})
        state["created_response"] = job
        state["generation_id"] = job["id"]
        save()
        print("Generation ID:", job["id"], flush=True)

    print("Following:", state["generation_id"], flush=True)
    deadline = time.monotonic() + args.wait_minutes * 60
    while time.monotonic() < deadline:
        job = api("GET", "/generations/" + state["generation_id"] + "?wait=25")
        state["last_response"] = job
        save()
        print("Status:", job["status"], flush=True)
        if job["status"] in {"failed", "cancelled"}:
            raise RuntimeError(json.dumps(job.get("error") or {"status": job["status"]}))
        if job["status"] == "succeeded":
            videos = [o for o in job["outputs"] if o.get("status") == "ready"
                      and o.get("url") and (o.get("mime_type") or "").startswith("video/")]
            if not videos:
                raise RuntimeError("No ready video output; inspect the saved response")
            # Do not send the API Authorization header to the download host.
            with requests.get(videos[0]["url"], stream=True, timeout=60) as response:
                response.raise_for_status()
                with args.output.open("wb") as output:
                    for chunk in response.iter_content(1024 * 1024):
                        output.write(chunk)
            state["downloaded_at"] = time.time()
            save()
            print("Saved:", str(args.output), "Credits charged:", job["credits_charged"], flush=True)
            return
        time.sleep(3)
    print("Still running remotely. Run again with the same state file to resume.", flush=True)


if __name__ == "__main__":
    main()