#!/usr/bin/env python3
"""
Seedance 2.0（火山方舟）视频生成 —— 多模态输入，建任务→轮询→下载。
依赖：pip install requests
官方文档：https://www.volcengine.com/docs/82379/1520757

Key 用环境变量：  export ARK_API_KEY=你的新Key
⚠️ 别把 Key 硬写进文件、别贴进聊天。
"""
import os, time, requests

def _load_key():
    k = os.environ.get("ARK_API_KEY", "").strip()       # 优先环境变量
    if k:
        return k
    p = os.path.expanduser("~/seedance/ark_key.txt")     # 否则读这个文件
    if os.path.exists(p):
        for line in open(p, encoding="utf-8"):
            s = line.strip()
            if s and not s.startswith("#"):
                return s
    return ""

ARK_API_KEY = _load_key()
MODEL = "doubao-seedance-2-0-260128"                    # Seedance 2.0
BASE  = "https://ark.cn-beijing.volces.com/api/v3"
HEAD  = {"Authorization": f"Bearer {ARK_API_KEY}", "Content-Type": "application/json"}


def create_task(prompt, images=None, video_url=None, audio_url=None,
                ratio="16:9", duration=10, generate_audio=True, watermark=False):
    """images: 参考图URL列表(含首帧)；video_url/audio_url: 参考视频/音频。
       注意：所有素材必须是公网可访问的 URL（本地文件先传对象存储/图床）。"""
    content = [{"type": "text", "text": prompt}]
    for u in (images or []):
        content.append({"type": "image_url", "image_url": {"url": u}, "role": "reference_image"})
    if video_url:
        content.append({"type": "video_url", "video_url": {"url": video_url}, "role": "reference_video"})
    if audio_url:
        content.append({"type": "audio_url", "audio_url": {"url": audio_url}, "role": "reference_audio"})
    body = {"model": MODEL, "content": content, "ratio": ratio,
            "duration": duration, "generate_audio": generate_audio, "watermark": watermark}
    r = requests.post(f"{BASE}/contents/generations/tasks", headers=HEAD, json=body, timeout=30)
    r.raise_for_status()
    return r.json()["id"]


def wait_video(task_id, interval=5):
    while True:
        r = requests.get(f"{BASE}/contents/generations/tasks/{task_id}", headers=HEAD, timeout=30)
        r.raise_for_status()
        d = r.json(); st = d.get("status"); print("状态:", st)
        if st == "succeeded":
            return d["content"]["video_url"]
        if st in ("failed", "cancelled"):
            raise RuntimeError(d)
        time.sleep(interval)


def download(url, path):
    with requests.get(url, stream=True, timeout=180) as r:
        r.raise_for_status()
        with open(path, "wb") as f:
            for c in r.iter_content(8192):
                f.write(c)
    print("已保存:", path)


if __name__ == "__main__":
    PROMPT = "（粘贴你的像素风开场 prompt）"
    tid = create_task(PROMPT, ratio="16:9", duration=15)
    # 图生/接龙： create_task(PROMPT, images=["https://.../首帧.png"])
    # 锁风格：    create_task(PROMPT, video_url="https://.../参考片.mp4", audio_url="https://.../bgm.mp3")
    print("任务:", tid)
    download(wait_video(tid), os.path.expanduser("~/seedance/输出/开场.mp4"))