#!/usr/bin/env python3
"""BHuman one-recipient example. No automatic paid POST retries."""
import argparse
import base64
import getpass
import json
import os
from pathlib import Path
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid

BASE = "https://studio.bhuman.ai/api/ai_studio"
STATE = Path("bhuman-generation.json")


def request(path, payload=None, authorization=None):
    headers = {"Accept": "application/json"}
    if authorization:
        headers["Authorization"] = authorization
    body = None
    if payload is not None:
        headers["Content-Type"] = "application/json"
        body = json.dumps(payload).encode()
    req = urllib.request.Request(BASE + path, data=body, headers=headers)
    # Never retry or follow redirects on a paid submission, or forward its secret.
    class NoRedirect(urllib.request.HTTPRedirectHandler):
        def redirect_request(self, *args, **kwargs):
            return None
    opener = urllib.request.build_opener(NoRedirect)
    with opener.open(req, timeout=60) as response:
        return json.load(response)


def save_ids(ids):
    # IDs give access to individual results: private file, no credentials.
    fd = os.open(STATE, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    with os.fdopen(fd, "w") as file:
        json.dump({"generation_ids": ids}, file)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--resume", action="store_true", help="Poll saved IDs; never submit.")
    args = parser.parse_args()
    if args.resume:
        ids = json.loads(STATE.read_text())["generation_ids"]
    else:
        if STATE.exists():
            raise SystemExit("Saved generation exists. Use --resume to check it before new work.")
        campaign_id = str(uuid.UUID(input("Owned campaign UUID: ").strip()))
        client_id = getpass.getpass("REST Client ID (hidden): ")
        client_secret = getpass.getpass("REST Client Secret (hidden): ")
        if not client_id or not client_secret:
            raise SystemExit("Both REST credential fields are required.")
        payload = {"campaign_id": campaign_id, "generation_mode": "keep_original",
                   "variables": ["first_name"], "names": [["Alex"]]}
        print("Requires a ready campaign with exactly the first_name variable.")
        print("One production generation may consume credits. Review Plan & usage first.")
        if input("Type GENERATE to submit exactly once: ").strip() != "GENERATE":
            raise SystemExit("No generation submitted.")
        authorization = "Basic " + base64.b64encode(
            (client_id + ":" + client_secret).encode()).decode()
        try:
            accepted = request("/pipeline/campaign", payload, authorization)
        except (urllib.error.URLError, TimeoutError, ValueError):
            raise SystemExit("Submission outcome is uncertain. Check the campaign in BHuman before resending. No automatic retry was made.")
        ids = accepted.get("result")
        if accepted.get("code") != 200 or not isinstance(ids, list) or not ids:
            raise SystemExit("No expected acceptance envelope. Check the campaign before resubmitting.")
        ids = [str(uuid.UUID(value)) for value in ids]
        save_ids(ids)
        print("Accepted. Private IDs saved in bhuman-generation.json. Keep this file out of shared logs and version control.")
    pending = set(str(uuid.UUID(value)) for value in ids)
    deadline = time.monotonic() + 1800
    while pending and time.monotonic() < deadline:
        for generation_id in list(pending):
            try:
                response = request("/generated_video_by_id?" + urllib.parse.urlencode({"id": generation_id}))
                video = response.get("result", {})
                if not isinstance(video, dict):
                    raise ValueError("Unexpected polling envelope")
            except (urllib.error.URLError, TimeoutError, ValueError):
                print("Status check unavailable. Backing off; no generation retry.")
                continue
            status = video.get("status")
            print("Generation status:", status)
            if status == "failed":
                raise SystemExit("Render failed. Keep the saved generation ID for support; review balance before a new submission.")
            if status == "succeeded" and video.get("url"):
                print("Video ready. Open this private URL and verify video/audio playback:")
                print(video["url"])
                pending.remove(generation_id)
        if pending:
            time.sleep(10)
    if pending:
        raise SystemExit("Polling deadline reached. Run with --resume later; do not submit again.")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        raise SystemExit("Interrupted. If submission began, check BHuman before resending. If IDs were saved, use --resume.")