#!/usr/bin/env python3
"""Studio API example, Python 3 standard library only.

Set STUDIO_ACCESS_KEY before running. Default: local GPU / fast / 5 seconds.
Grok is paid: explicitly set STUDIO_BACKEND=grok-cloud and STUDIO_QUALITY=high.
Running this script creates a project and submits one real generation.
"""
import hashlib
import json
import os
from pathlib import Path
import shutil
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid


def main():
    base = os.environ.get("STUDIO_URL", "https://video-gen.qianyiyule.top").rstrip("/")
    key = os.environ.get("STUDIO_ACCESS_KEY", "")
    if not key:
        raise SystemExit("请先设置 STUDIO_ACCESS_KEY。")
    author = os.environ.get("STUDIO_AUTHOR", "api-member")
    backend = os.environ.get("STUDIO_BACKEND", "win-4060ti")
    quality = os.environ.get("STUDIO_QUALITY", "high" if backend == "grok-cloud" else "fast")
    seconds = int(os.environ.get("STUDIO_SECONDS", "5"))
    project_id = "api-demo-" + uuid.uuid4().hex[:12]

    # Do not forward a key through redirects to another origin.
    class NoRedirect(urllib.request.HTTPRedirectHandler):
        def redirect_request(self, req, fp, code, msg, headers, newurl):
            return None
    opener = urllib.request.build_opener(NoRedirect())

    def request(method, path, body=None, revision=None, idem=None):
        headers = {"Authorization": "Bearer " + key, "Accept": "application/json"}
        data = None
        if body is not None:
            headers["Content-Type"] = "application/json"
            data = json.dumps(body, ensure_ascii=False).encode()
        if revision is not None: headers["If-Match"] = str(revision)
        if idem: headers["Idempotency-Key"] = idem
        req = urllib.request.Request(base + "/api" + path, data=data, headers=headers, method=method)
        try:
            with opener.open(req, timeout=60) as response:
                return json.load(response)
        except urllib.error.HTTPError as error:
            raise RuntimeError(f"HTTP {error.code}: {error.read().decode('utf-8', 'replace')}") from None

    project = request("POST", "/projects", {
        "id": project_id, "title": "API 单镜示例", "template": "single_shot", "aspect": "16:9", "by": author,
    })
    shots = request("PUT", f"/projects/{project_id}/shots", {"by": author, "shots": [{
        "no": 1, "seconds": seconds,
        "prompt": os.environ.get("STUDIO_PROMPT", "A quiet lake at sunrise, gentle ripples, slow camera push in. No dialogue, no music."),
    }]}, revision=project["revision"])
    body = {"shots": [1], "backend": backend, "quality": quality, "by": author}
    idem = str(uuid.uuid4())
    # Save the exact submission before sending, so an uncertain response can be
    # retried with the SAME key, body and If-Match, instead of generating twice.
    state_path = Path(project_id + "-request.json")
    state = {"project_id": project_id, "revision": shots["revision"], "idempotency_key": idem, "render_body": body}
    state_path.write_text(json.dumps(state, ensure_ascii=False, indent=2))
    print("提交记录（不含访问 key）：", state_path, flush=True)
    batch = request("POST", f"/projects/{project_id}/render", body, revision=shots["revision"], idem=idem)
    job_id = batch["jobs"][0]["id"]
    state["job_id"] = job_id
    state_path.write_text(json.dumps(state, ensure_ascii=False, indent=2))
    print("job_id:", job_id, "backend:", backend, flush=True)
    deadline = time.monotonic() + 7200
    while time.monotonic() < deadline:
        job = request("GET", "/jobs/" + job_id)
        print(job["status"], job.get("block_reason") or "", flush=True)
        if job["status"] == "done": break
        if job["status"] in {"failed", "cancelled"}:
            raise SystemExit(job.get("error") or job["status"])
        time.sleep(10)
    else:
        raise SystemExit("等待超过 2 小时，任务未自动取消。请用记录中的 job_id 继续查询。")
    shot = next(s for s in request("GET", f"/projects/{project_id}/shots")["shots"] if s["no"] == job["shot_no"])
    version = next(v for v in shot["versions"] if v["version"] == job["version"])
    logical_path = urllib.parse.quote(version["file"], safe="/")
    url = base + f"/api/projects/{project_id}/files/" + logical_path
    req = urllib.request.Request(url, headers={"Authorization": "Bearer " + key})
    output = Path(project_id + ".mp4")
    partial = output.with_suffix(".mp4.part")
    with opener.open(req, timeout=600) as response, partial.open("wb") as f:
        shutil.copyfileobj(response, f)
    digest = hashlib.sha256()
    with partial.open("rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""): digest.update(chunk)
    if digest.hexdigest() != version["output_hash"]:
        raise SystemExit("下载哈希不符，保留 .part 文件，请重新下载。")
    partial.replace(output)
    print("视频：", output.resolve(), "费用（USD）：", job.get("cost_usd"), flush=True)


if __name__ == "__main__":
    main()