#!/usr/bin/env python3
"""Python 3.10+; standard library only. Example client, not a live test record."""
import argparse
import json
import math
import os
from pathlib import Path
import sys
import time
from urllib.error import HTTPError, URLError
from urllib.parse import quote, urlparse
from urllib.request import HTTPRedirectHandler, Request, build_opener

ENDPOINTS = {
    "seedance-2": "https://api.quickrouter.ai/api/v3/contents/generations/tasks",
    "legacy": "https://api.quickrouter.ai/volc/v1/contents/generations/tasks",
}


class NoRedirect(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


class ApiError(RuntimeError):
    def __init__(self, status):
        self.status = status
        super().__init__(f"HTTP {status}; check model permission, quota and endpoint. Request was not resubmitted.")


def request_json(method, url, key, payload=None, timeout=30):
    data = None if payload is None else json.dumps(payload).encode()
    request = Request(url, data=data, method=method, headers={
        "Authorization": f"Bearer {key}", "Content-Type": "application/json", "Accept": "application/json",
    })
    try:
        with build_opener(NoRedirect()).open(request, timeout=timeout) as response:
            result = json.load(response)
    except HTTPError as error:
        raise ApiError(error.code) from None
    except (URLError, TimeoutError):
        raise RuntimeError("Network error; if this was POST, creation may have succeeded. Check task history before creating again.") from None
    if not isinstance(result, dict):
        raise RuntimeError("Expected a JSON object; check the API response.")
    return result


def task_id(value):
    if not isinstance(value, str) or not value or len(value) > 256:
        raise RuntimeError("Missing or invalid task ID. Check task history before creating again.")
    return value


def create_task(endpoint, payload, key, state_path, request=request_json):
    if not isinstance(payload, dict) or not isinstance(payload.get("model"), str) or not isinstance(payload.get("content"), list) or not payload["content"]:
        raise RuntimeError("Payload must contain model and a nonempty content array.")
    # Reserve the state file before the paid POST; never overwrite or auto-resubmit.
    with state_path.open("x", encoding="utf-8") as file:
        json.dump({"endpoint": endpoint, "status": "submission_unconfirmed"}, file)
    response = request("POST", endpoint, key, payload)
    identifier = task_id(response.get("id"))
    print(f"Task created: {identifier}", file=sys.stderr, flush=True)
    state_path.write_text(json.dumps({"endpoint": endpoint, "id": identifier}, indent=2), encoding="utf-8")
    return identifier


def poll_task(endpoint, identifier, key, timeout=600, interval=5, request=request_json, clock=time.monotonic, sleep=time.sleep):
    if not math.isfinite(timeout) or not math.isfinite(interval) or timeout <= 0 or interval <= 0:
        raise RuntimeError("Timeout and interval must be positive finite numbers.")
    deadline = clock() + timeout
    url = endpoint + "/" + quote(task_id(identifier), safe="")
    while clock() < deadline:
        try:
            response = request("GET", url, key, timeout=min(30, max(0.1, deadline - clock())))
        except ApiError as error:
            if error.status != 429 and not 500 <= error.status < 600:
                raise
            print(f"Query HTTP {error.status}; waiting before retry.", file=sys.stderr)
        else:
            status = response.get("status")
            print(f"Task status: {status}", file=sys.stderr)
            if status == "succeeded":
                content = response.get("content")
                video_url = content.get("video_url") if isinstance(content, dict) else None
                if not isinstance(video_url, str) or urlparse(video_url).scheme != "https" or not urlparse(video_url).netloc:
                    raise RuntimeError("Task succeeded but a valid HTTPS content.video_url is missing; inspect the task response.")
                return {"id": identifier, "status": status, "video_url": video_url, "usage": response.get("usage")}
            if status in {"failed", "expired", "cancelled", "canceled"}:
                error = response.get("error")
                code = error.get("code") if isinstance(error, dict) else None
                raise RuntimeError(f"Task ended: {status}; error code: {code}. Check console for details and billing; no resubmission.")
            if status not in {"queued", "running"}:
                raise RuntimeError(f"Unknown task status: {status}; stop and inspect the response.")
        sleep(min(interval, max(0, deadline - clock())))
    raise TimeoutError("Polling deadline reached; the remote task is NOT cancelled. Resume query with the same task ID.")


def main():
    parser = argparse.ArgumentParser(description="Submit once, then poll a Seedance task. Creation may incur charges.")
    parser.add_argument("action", choices=["create", "query"])
    parser.add_argument("--api", choices=ENDPOINTS, default="seedance-2")
    parser.add_argument("--payload", type=Path)
    parser.add_argument("--state", type=Path, help="New state filename for create, existing filename for query")
    parser.add_argument("--task-id")
    parser.add_argument("--timeout", type=float, default=600)
    parser.add_argument("--interval", type=float, default=5)
    args = parser.parse_args()
    if not all(math.isfinite(v) and v > 0 for v in (args.timeout, args.interval)):
        parser.error("--timeout and --interval must be positive finite numbers")
    key = os.environ.get("QUICKROUTER_API_KEY")
    if not key:
        parser.error("Set QUICKROUTER_API_KEY in your environment")
    endpoint = ENDPOINTS[args.api]
    if args.action == "create":
        if not args.payload or not args.state or args.task_id:
            parser.error("create requires --payload and a NEW --state file, without --task-id")
        identifier = create_task(endpoint, json.loads(args.payload.read_text(encoding="utf-8")), key, args.state)
    else:
        if args.payload or bool(args.state) == bool(args.task_id):
            parser.error("query requires exactly one of --state or --task-id, without --payload")
        if args.state:
            state = json.loads(args.state.read_text(encoding="utf-8"))
            endpoint = state.get("endpoint")
            if endpoint not in ENDPOINTS.values():
                raise RuntimeError("Unrecognized endpoint in state file")
            identifier = task_id(state.get("id"))
        else:
            identifier = task_id(args.task_id)
    result = poll_task(endpoint, identifier, key, args.timeout, args.interval)
    print(json.dumps(result, ensure_ascii=False, indent=2))


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