"""HamuBox async media client. Requires Python 3.10+, no third-party packages.

Submit: HAMUBOX_API_KEY=... python client.py image banana-generate.json
Resume: HAMUBOX_API_KEY=... python client.py video --task-id task_...
"""
import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request

BASE_URL = "https://api.hamubox.com"


def request(method, path, api_key, payload=None):
    body = None if payload is None else json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        BASE_URL + path, data=body, method=method,
        headers={"Authorization": "Bearer " + api_key, "Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=120) as response:
        return json.load(response)


def poll(kind, task_id, api_key, max_wait):
    deadline = time.monotonic() + max_wait
    errors = 0
    while time.monotonic() < deadline:
        delay = 5
        try:
            task = request("GET", task_path(kind) + "/" + task_id, api_key)
            errors = 0
            status = task.get("status", "unknown")
            print(f"{task_id}: {status} ({task.get('progress', 0)}%)", file=sys.stderr)
            if status == "failed":
                raise RuntimeError(json.dumps(task.get("error", {}), ensure_ascii=False))
            if status == "completed":
                if kind == "video":
                    urls = [item.get("url") for item in task.get("outputs", [])]
                    if not urls:
                        urls = [task.get("metadata", {}).get("url")]
                elif kind == "novelai":
                    result = task.get("result") or {}
                    if result.get("encoded_vibe"):
                        return task
                    urls = [item.get("url") for item in result.get("images", [])]
                else:
                    if task.get("operation") == "describe" and task.get("descriptions"):
                        return task
                    urls = [item.get("url") for item in task.get("outputs", [])]
                if not urls or any(not url for url in urls):
                    raise RuntimeError("Completed without a result URL; keep the task ID and contact support.")
                return task
        except urllib.error.HTTPError as error:
            if error.code not in (429, 500, 502, 503, 504):
                raise
            errors += 1
            retry_after = error.headers.get("Retry-After", "")
            delay = min(60, 5 * (2 ** min(errors, 4)))
            if retry_after.isdigit():
                delay = max(delay, int(retry_after))
            if errors >= 5:
                raise RuntimeError(f"Polling unavailable. Resume later with task ID {task_id}.") from error
        except (urllib.error.URLError, TimeoutError) as error:
            errors += 1
            if errors >= 5:
                raise RuntimeError(f"Network unavailable. Resume later with task ID {task_id}.") from error
            delay = min(60, 5 * (2 ** errors))
        remaining = deadline - time.monotonic()
        if remaining > 0:
            time.sleep(min(delay, remaining))
    raise TimeoutError(f"Client wait limit reached; this does not cancel the task. Resume with {task_id}.")


def task_path(kind):
    return "/v1/novelai/tasks" if kind == "novelai" else f"/v1/{kind}/generations"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("kind", choices=["image", "video", "novelai"])
    parser.add_argument("request_file", nargs="?")
    parser.add_argument("--task-id")
    parser.add_argument("--max-wait", type=int, default=1800)
    args = parser.parse_args()
    key = os.environ.get("HAMUBOX_API_KEY", "").strip()
    if not key:
        parser.error("Set HAMUBOX_API_KEY first.")
    if bool(args.request_file) == bool(args.task_id):
        parser.error("Provide either request_file or --task-id, not both.")
    if args.max_wait <= 0:
        parser.error("--max-wait must be positive.")
    task_id = args.task_id
    if not task_id:
        with open(args.request_file, encoding="utf-8") as source:
            payload = json.load(source)
        # Never automatically retry POST: an interrupted response can still create a paid task.
        submitted = request("POST", task_path(args.kind), key, payload)
        task_id = submitted.get("id") or submitted.get("task_id")
        if not task_id:
            raise RuntimeError("No task ID returned: " + json.dumps(submitted, ensure_ascii=False))
        print(f"TASK_ID={task_id}", flush=True)
        with open(f"{task_id}.json", "w", encoding="utf-8") as destination:
            json.dump({"id": task_id, "kind": args.kind, "model": payload.get("model")}, destination)
    result = poll(args.kind, task_id, key, args.max_wait)
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    try:
        main()
    except urllib.error.HTTPError as error:
        print(f"HTTP {error.code}: " + error.read(4096).decode("utf-8", errors="replace"), file=sys.stderr)
        sys.exit(1)
    except (RuntimeError, TimeoutError, urllib.error.URLError) as error:
        print(str(error), file=sys.stderr)
        sys.exit(1)