#!/usr/bin/env python3
"""AnimGen Public API example, Python 3.10+, standard library only.

Prepare only quotes; run explicitly approves a local preflight check, then creates.
The API recalculates price on create: --approve-credits is NOT a server-side cap.
Keep state private: it contains the image and prompt, but never the API key.
Resume with the same state file, never a newly prepared request.
"""

import argparse
import base64
import hashlib
import json
import os
from pathlib import Path
import random
import re
import sys
import tempfile
import time
import uuid
from datetime import timezone
from email.utils import parsedate_to_datetime
from urllib.error import HTTPError, URLError
from urllib.parse import urlsplit
from urllib.request import HTTPRedirectHandler, Request, build_opener


TERMINAL = {"succeeded", "failed", "cancelled"}
DEFAULT_BASE = "https://api.animgen.com/v1"
MAX_RETRY_AGE = 24 * 60 * 60


class WorkflowError(Exception):
    """Safe messages only: never propagate a URL, body, or credential."""


class NoRedirect(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None  # Never forward API credentials through redirects.


def canonical(value):
    return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False)


def request_hash(body):
    return hashlib.sha256(canonical(body).encode()).hexdigest()


def private_write(path, value, *, new=False):
    path = Path(path)
    payload = (json.dumps(value, ensure_ascii=False, indent=2) + "\n").encode()
    if new:
        fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
        with os.fdopen(fd, "wb") as file:
            file.write(payload)
        return
    # Atomic replacement preserves the original request across interruption.
    name = None
    try:
        with tempfile.NamedTemporaryFile(dir=path.parent, delete=False) as file:
            name = file.name
            file.write(payload)
        os.replace(name, path)
    finally:
        if name and os.path.exists(name):
            os.unlink(name)


def retry_delay(headers, fallback):
    value = headers.get("Retry-After", "")
    try:
        delay = float(value)
        if delay >= 0 and delay < float("inf"):
            return delay
    except (TypeError, ValueError):
        pass
    try:
        date = parsedate_to_datetime(value)
        if date.tzinfo is None:
            date = date.replace(tzinfo=timezone.utc)
        return max(0, date.timestamp() - time.time())
    except (TypeError, ValueError, OverflowError):
        return fallback


def remaining(deadline):
    seconds = deadline - time.monotonic()
    if seconds <= 0:
        raise WorkflowError("Local deadline reached. Resume with the same state file; the remote task may still run.")
    return min(30, seconds)


def wait_for(delay, deadline):
    delay = max(0.05, delay)
    if time.monotonic() + delay >= deadline:
        raise WorkflowError("Retry delay exceeds the local deadline. Resume later with the same state file.")
    time.sleep(delay)


def safe_identifier(value):
    return str(value) if re.fullmatch(r"[A-Za-z0-9_.-]{1,100}", str(value)) else "unknown"


class Client:
    def __init__(self):
        self.base = os.environ.get("ANIMGEN_API_BASE", DEFAULT_BASE).rstrip("/")
        parsed = urlsplit(self.base)
        local = parsed.hostname in {"localhost", "127.0.0.1", "::1"}
        if not parsed.hostname or parsed.username or parsed.password or parsed.query or parsed.fragment:
            raise WorkflowError("Invalid API base URL.")
        if parsed.scheme != "https" and not (parsed.scheme == "http" and local):
            raise WorkflowError("Use HTTPS; HTTP is allowed only for a local mock server.")
        self.key = os.environ.get("ANIMGEN_API_KEY")
        if not self.key:
            raise WorkflowError("Set ANIMGEN_API_KEY privately before running the example.")
        self.opener = build_opener(NoRedirect())

    def request(self, method, path, *, body=None, key=None, deadline=None):
        deadline = deadline or time.monotonic() + 120
        payload = None if body is None else canonical(body).encode()
        headers = {"Authorization": "Bearer " + self.key, "Accept": "application/json"}
        if payload is not None:
            headers["Content-Type"] = "application/json"
        if key:
            headers["Idempotency-Key"] = key
        # The only retried writes here are non-spending quotes and keyed creates.
        safe_retry = method == "GET" or path.endswith("/quote") or bool(key)
        for attempt in range(5):
            response_headers = {}
            try:
                req = Request(self.base + path, data=payload, headers=headers, method=method)
                with self.opener.open(req, timeout=remaining(deadline)) as response:
                    data = json.loads(response.read(16 * 1024 * 1024))
                    return data, response.headers
            except HTTPError as error:
                response_headers = error.headers
                try:
                    detail = json.loads(error.read(1024 * 1024)).get("error", {})
                except (ValueError, AttributeError):
                    detail = {}
                code = safe_identifier(detail.get("code", "HTTP_ERROR"))
                request_id = safe_identifier(detail.get("request_id", "unknown"))
                can_retry = error.code == 429 or (
                    detail.get("retryable") is True and (
                        error.code == 503 or (error.code == 409 and code == "IDEMPOTENCY_IN_PROGRESS")
                    )
                )
                if not safe_retry or not can_retry or attempt == 4:
                    raise WorkflowError(f"API HTTP {error.code}: {code}; request_id={request_id}") from None
            except (URLError, TimeoutError, ConnectionError, OSError):
                if not safe_retry or attempt == 4:
                    raise WorkflowError("API transport failed. Preserve the state and reuse its original key.") from None
            except (ValueError, AttributeError):
                raise WorkflowError("Invalid API response. Preserve the state before retrying.") from None
            delay = retry_delay(response_headers, min(30, 2 ** attempt) + random.random())
            wait_for(delay, deadline)
        raise WorkflowError("Retry limit reached.")


def prepare(args, client):
    if Path(args.state).exists():
        raise WorkflowError("State already exists. Resume it or choose a new file for a deliberate new operation.")
    image_path = Path(args.image)
    mime = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp"}.get(image_path.suffix.lower())
    if not mime or image_path.stat().st_size > 10 * 1024 * 1024:
        raise WorkflowError("Use a PNG/JPEG/WebP image no larger than 10 MiB for this inline example.")
    image = image_path.read_bytes()
    catalog, _ = client.request("GET", "/models")
    model = next((item for item in catalog["data"] if item["id"] == args.model), None)
    if not model or not model["supports_first_frame"] or "first_frame" not in model["modes"]:
        raise WorkflowError("Choose a current model that supports first-frame generation.")
    if model["requires_prompt"] and not args.prompt.strip():
        raise WorkflowError("This model requires a prompt.")
    video = {"model": model["id"], "duration_seconds": model["default_duration_seconds"]}
    for field in ("resolution", "ratio"):
        if model.get("default_" + field):
            video[field] = model["default_" + field]
    body = {
        "input": {"first_frame": {"type": "base64", "media_type": mime, "data": base64.b64encode(image).decode()}},
        "prompt": args.prompt,
        "video": video,
        "selection": {"mode": "full"},
        "export": {"output_formats": ["frames_zip"], "frame_count": 24, "output_width": 512, "output_height": 512},
    }
    account, _ = client.request("GET", "/account")
    quote, _ = client.request("POST", "/animations/quote", body=body)
    balance, _ = client.request("GET", "/credits/balance")
    state = {
        "schema_version": 1, "api_base": client.base, "account_id": account["id"], "api_key_id": account["api_key"]["id"],
        "idempotency_key": str(uuid.uuid4()), "request": body, "request_hash": request_hash(body),
        "task_id": None, "first_create_at": None,
    }
    private_write(args.state, state, new=True)
    print(f"Quoted credits: {quote['credits']}; available: {balance['available']}. No generation started.")
    print("Private state saved. Review the quote, then use run --approve-credits with the SAME state file.")
    return 0


def validate_download_url(url, api_base):
    parsed, api = urlsplit(url), urlsplit(api_base)
    local_mock = api.scheme == "http" and parsed.netloc == api.netloc and parsed.scheme == "http"
    if not parsed.hostname or parsed.username or parsed.password or parsed.fragment:
        raise WorkflowError("Unsafe download URL rejected.")
    if parsed.scheme != "https" and not local_mock:
        raise WorkflowError("Non-HTTPS download rejected.")


def save_asset(client, asset_id, output, deadline):
    asset_id = str(uuid.UUID(asset_id))  # No remote-controlled path components.
    for attempt in range(2):
        asset, _ = client.request("GET", f"/assets/{asset_id}", deadline=deadline)
        url = asset["download_url"]
        validate_download_url(url, client.base)
        suffix = {
            "application/zip": ".zip", "image/png": ".png", "application/json": ".json",
            "video/mp4": ".mp4", "video/webm": ".webm", "video/quicktime": ".mov",
        }.get(asset["mime_type"], ".bin")
        destination = Path(output) / (asset_id + suffix)
        expected = asset["byte_size"]
        if not isinstance(expected, int) or expected < 0:
            raise WorkflowError("Invalid asset byte size.")
        if destination.exists() or destination.is_symlink():
            if not destination.is_symlink() and destination.is_file() and destination.stat().st_size == expected:
                return destination
            raise WorkflowError("An existing output does not match; choose a different output directory.")

        class DownloadRedirect(HTTPRedirectHandler):
            def redirect_request(self, req, fp, code, msg, headers, newurl):
                validate_download_url(newurl, client.base)
                return super().redirect_request(req, fp, code, msg, headers, newurl)

        name = None
        try:
            # A separate request/opener: no Authorization or API headers, including on redirects.
            opener = build_opener(DownloadRedirect())
            with opener.open(Request(url), timeout=remaining(deadline)) as response:
                with tempfile.NamedTemporaryFile(dir=output, delete=False) as file:
                    name, count = file.name, 0
                    while True:
                        remaining(deadline)
                        chunk = response.read(64 * 1024)
                        if not chunk:
                            break
                        count += len(chunk)
                        if count > expected:
                            raise WorkflowError("Download exceeded the asset's declared size.")
                        file.write(chunk)
                if count != expected:
                    raise WorkflowError("Incomplete download; resume to retry this asset.")
            # Exclusive link avoids overwriting an unrelated file in a concurrent run.
            os.link(name, destination)
            return destination
        except HTTPError as error:
            if error.code not in {401, 403, 404} or attempt:
                raise WorkflowError(f"Asset download failed: HTTP {error.code}; URL omitted.") from None
            # Refresh the asset metadata once if a signed URL expired.
        except (URLError, TimeoutError, ConnectionError):
            raise WorkflowError("Asset transfer failed; resume with the same state file.") from None
        finally:
            if name and os.path.exists(name):
                os.unlink(name)
    raise WorkflowError("Could not refresh the asset download.")


def run(args, client):
    state = json.loads(Path(args.state).read_text())
    if state.get("schema_version") != 1 or state.get("api_base") != client.base:
        raise WorkflowError("State version or API base mismatch.")
    if state.get("request_hash") != request_hash(state["request"]):
        raise WorkflowError("The saved request changed. Recover the original task instead of reusing this key.")
    deadline = time.monotonic() + args.timeout
    account, _ = client.request("GET", "/account", deadline=deadline)
    if account["id"] != state["account_id"]:
        raise WorkflowError("This state belongs to a different account. Do not recreate the operation.")
    if not state["task_id"]:
        if account["api_key"]["id"] != state.get("api_key_id"):
            raise WorkflowError("API key identity changed. Idempotency is scoped to the original key; recover the task before proceeding.")
        started = state["first_create_at"]
        if started is not None and time.time() - started >= MAX_RETRY_AGE:
            raise WorkflowError("Uncertain create is older than 24 hours. Inspect existing tasks before any new creation.")
        if started is None:
            quote, _ = client.request("POST", "/animations/quote", body=state["request"], deadline=deadline)
            if quote["credits"] > args.approve_credits:
                raise WorkflowError("Current quote exceeds your local approval. No create was sent.")
            print(f"Preflight quote: {quote['credits']}. This check is not a server-enforced spending cap.")
            state["first_create_at"] = time.time()
            private_write(args.state, state)
        task, _ = client.request("POST", "/animations", body=state["request"], key=state["idempotency_key"], deadline=deadline)
        state["task_id"] = str(uuid.UUID(task["id"]))
        private_write(args.state, state)
    task_id = str(uuid.UUID(state["task_id"]))
    while True:
        task, headers = client.request("GET", f"/animations/{task_id}", deadline=deadline)
        status = task["status"]
        print(f"Task {task_id}: {safe_identifier(status)}")
        if status in TERMINAL:
            break
        if status not in {"queued", "running", "cancelling"}:
            raise WorkflowError("Unknown task status. Preserve the task ID and inspect the response contract.")
        wait_for(retry_delay(headers, 5), deadline)
    Path(args.output).mkdir(parents=True, exist_ok=True)
    failed_downloads = 0
    for asset in task.get("outputs", []):
        try:
            path = save_asset(client, asset["id"], args.output, deadline)
            print(f"Saved asset: {path}")
        except (WorkflowError, OSError, ValueError, KeyError):
            failed_downloads += 1
            print(f"Could not save asset {safe_identifier(asset.get('id'))}; resume to retry.", file=sys.stderr)
    if status != "succeeded":
        code = safe_identifier((task.get("error") or {}).get("code", status))
        print(f"Task incomplete: {code}. Available outputs were processed; no new task was created.", file=sys.stderr)
    return 0 if status == "succeeded" and not failed_downloads else 2


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    prep = commands.add_parser("prepare", help="Discover and quote only; never creates a paid task")
    prep.add_argument("--image", required=True)
    prep.add_argument("--model", required=True)
    prep.add_argument("--prompt", required=True)
    prep.add_argument("--state", required=True)
    execute = commands.add_parser("run", help="Explicitly approve a local check, create once logically, and resume")
    execute.add_argument("--state", required=True)
    execute.add_argument("--approve-credits", required=True, type=int)
    execute.add_argument("--output", required=True)
    execute.add_argument("--timeout", type=float, default=900, help="Overall local deadline in seconds")
    args = parser.parse_args(argv)
    if args.command == "run" and (args.approve_credits < 0 or not 0 < args.timeout < float("inf")):
        parser.error("Use a nonnegative approval and a finite positive timeout.")
    try:
        client = Client()
        return prepare(args, client) if args.command == "prepare" else run(args, client)
    except (WorkflowError, OSError, ValueError, KeyError, TypeError) as error:
        message = str(error) if isinstance(error, WorkflowError) else "Invalid local file, state, or response; details omitted to protect private data."
        print(message, file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())