#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""H3 视频云 Python SDK（v5.2，与 MiniMax 官方 MiniMax-H3 V2 接口一致）

用法：
    from h3sdk import H3Client
    cli = H3Client("sk-h3-你的密钥")
    # 文生视频
    task = cli.generate("一个少年在雨后的山路上奔跑，电影感", ratio="16:9")
    # 图生视频（首帧）
    task = cli.generate("人物从门口走向窗边", first_frame_image="https://你的oss.com/first.jpg")
    # 多模态参考
    task = cli.generate("照这个人和这段动作来", reference_images=["https://.../a.jpg"],
                        reference_videos=["https://.../ref.mp4"], duration=10)
    print(task["task"]["content"]["url"])
"""
import json, time, urllib.request, urllib.error

DEFAULT_BASE = "https://www.nexusvoidaigc.com"


class H3Error(Exception):
    def __init__(self, code, msg, http_status=None):
        super().__init__(f"[{code}] {msg}")
        self.code, self.msg, self.http_status = code, msg, http_status


def _el_image(url, role):
    return {"type": "image_url", "image_url": {"url": url}, "role": role}


class H3Client:
    def __init__(self, api_key, base_url=DEFAULT_BASE, timeout=120):
        self.api_key, self.base_url, self.timeout = api_key, base_url.rstrip("/"), timeout

    def _req(self, method, path, body=None, timeout=None):
        req = urllib.request.Request(
            self.base_url + path,
            data=json.dumps(body).encode() if body is not None else None,
            headers={"Authorization": "Bearer " + self.api_key, "Content-Type": "application/json"},
            method=method)
        try:
            with urllib.request.urlopen(req, timeout=timeout or self.timeout) as r:
                return json.loads(r.read().decode())
        except urllib.error.HTTPError as e:
            try:
                br = json.loads(e.read().decode()).get("base_resp", {})
                raise H3Error(str(br.get("status_code", e.code)), br.get("status_msg", ""), e.code)
            except json.JSONDecodeError:
                raise H3Error("http_error", str(e), e.code)

    def generate(self, prompt, *, first_frame_image=None, last_frame_image=None,
                 reference_images=None, reference_videos=None, reference_audios=None,
                 duration=5, resolution="768P", ratio=None,
                 callback_url=None, prompt_optimizer=False, aigc_watermark=False,
                 wait=True, poll_interval=10, timeout=1800, idempotency_key=None,
                 subject_reference=None, seconds=None, size=None):
        """提交生成（content[] 与官方 H3 V2 一致）；wait=True 时阻塞到定局并返回任务详情（含 task.content.url）。

        first_frame_image / last_frame_image：首帧 / 尾帧图（公网 URL，与 reference_* 互斥）
        reference_images / reference_videos / reference_audios：多模态参考（URL 列表）
        ratio：文生视频必填（非 adaptive）；带图/参考任务默认 adaptive
        兼容旧参数：subject_reference→reference_images、seconds→duration、size→resolution。
        """
        if seconds is not None: duration = seconds
        if size is not None: resolution = size
        if subject_reference and not reference_images:
            sr = subject_reference if isinstance(subject_reference, list) else [subject_reference]
            reference_images = [img for item in sr
                                for img in (item.get("image", []) if isinstance(item, dict) else [item])]
        content = [{"type": "text", "text": prompt}]
        if first_frame_image: content.append(_el_image(first_frame_image, "first_frame"))
        if last_frame_image: content.append(_el_image(last_frame_image, "last_frame"))
        for u in (reference_images or []): content.append(_el_image(u, "reference_image"))
        for u in (reference_videos or []):
            content.append({"type": "video_url", "video_url": {"url": u}, "role": "reference_video"})
        for u in (reference_audios or []):
            content.append({"type": "audio_url", "audio_url": {"url": u}, "role": "reference_audio"})
        body = {"model": "MiniMax-H3", "content": content,
                "duration": int(duration), "resolution": resolution,
                "prompt_optimizer": bool(prompt_optimizer), "aigc_watermark": bool(aigc_watermark)}
        if ratio: body["ratio"] = ratio
        if callback_url: body["callback_url"] = callback_url
        if idempotency_key: body["idempotency_key"] = idempotency_key
        r = self._req("POST", "/v2/video_generation", body)
        tid = r.get("task_id")
        if not tid:
            br = r.get("base_resp", {})
            raise H3Error(str(br.get("status_code", "unknown")), br.get("status_msg", "create failed"))
        if not wait:
            return {"task_id": tid}
        return self.wait_task(tid, poll_interval=poll_interval, timeout=timeout)

    def query(self, task_id):
        return self._req("GET", f"/v2/query/video_generation/{task_id}")

    def wait_task(self, task_id, poll_interval=10, timeout=1800):
        deadline = time.time() + timeout
        while time.time() < deadline:
            r = self.query(task_id)
            st = (r.get("task") or {}).get("status")
            if st in ("succeeded", "failed", "cancelled"):
                return r
            time.sleep(poll_interval)
        raise H3Error("timeout", f"任务 {task_id} 超过 {timeout}s 未定局")


if __name__ == "__main__":
    import os, sys
    cli = H3Client(os.environ["H3_API_KEY"])
    prompt = sys.argv[1] if len(sys.argv) > 1 else "一只橙色小猫在窗台上看向镜头"
    print(json.dumps(cli.generate(prompt, ratio="16:9"), ensure_ascii=False, indent=2))