# LIB · Seedance 2.0 视频生成适配器 + 异步任务队列
# D180 · 2026-07-10
"""火山方舟 ARK → Seedance 2.0 · 提交+轮询+下载+写库 全自动"""
import json, os, subprocess, time, tempfile, base64

ARK_KEY = "ark-ddeba9f4-8c5a-449e-b549-9c29ec1e6f8c-a39ea"
ARK_BASE = "https://ark.cn-beijing.volces.com/api/v3"
SEEDANCE_MODEL = "doubao-seedance-2-0-260128"
TASK_URL = f"{ARK_BASE}/contents/generations/tasks"

def _curl(method, url, payload=None):
    """通用 curl 调用"""
    tf = None
    cmd = ["curl", "-s", url, "-H", f"Authorization: Bearer {ARK_KEY}"]
    if payload:
        tf = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False, encoding='utf-8')
        json.dump(payload, tf, ensure_ascii=False)
        tf.close()
        cmd += ["-H", "Content-Type: application/json", "-d", f"@{tf.name}"]
    if method == "POST":
        cmd.insert(1, "-X")
        cmd.insert(2, "POST")
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
        if tf: os.unlink(tf.name)
        return json.loads(r.stdout) if r.stdout else {"error": r.stderr}
    except Exception as e:
        if tf: os.unlink(tf.name)
        return {"error": str(e)}

def submit_task(prompt, references=None, duration=6, resolution="720p"):
    """提交 Seedance 生成任务 → 返回 task_id"""
    content = [{"type": "text", "text": prompt}]
    if references:
        for ref_type, ref_path in references:
            with open(ref_path, "rb") as f:
                fmt = "png" if ref_path.endswith(".png") else "jpeg"
                b64 = base64.b64encode(f.read()).decode()
            content.append({
                "type": "image_url",
                "role": ref_type,
                "image_url": {"url": f"data:image/{fmt};base64,{b64}"}
            })
    payload = {
        "model": SEEDANCE_MODEL,
        "content": content,
        "duration": duration,
        "resolution": resolution
    }
    result = _curl("POST", TASK_URL, payload)
    if "id" in result:
        return result["id"]
    return {"error": result}

def poll_task(task_id, poll_interval=30, max_wait=600):
    """轮询任务状态 → 返回结果"""
    elapsed = 0
    while elapsed < max_wait:
        result = _curl("GET", f"{TASK_URL}/{task_id}")
        if "error" in result:
            return result
        status = result.get("status", "")
        if status == "succeeded":
            return result
        if status == "failed":
            return {"error": f"Task failed: {result}"}
        time.sleep(poll_interval)
        elapsed += poll_interval
    return {"error": "timeout"}

def download_video(video_url, output_path):
    """下载视频"""
    subprocess.run(["curl", "-s", "-o", output_path, video_url], check=False)
    return os.path.exists(output_path) and os.path.getsize(output_path) > 0

def generate_shot(shot_id, prompt, references, duration=6, output_dir="outputs/videos"):
    """一站式生成：提交→轮询→下载→返回结果"""
    print(f"[LIB] 提交 {shot_id}...")
    os.makedirs(output_dir, exist_ok=True)
    task_id = submit_task(prompt, references, duration)
    if isinstance(task_id, dict):
        return {"error": task_id}

    print(f"[LIB] 任务 {task_id[:20]}... 等待完成")
    result = poll_task(task_id)
    if "error" in result:
        return result

    video_url = result.get("content", {}).get("video_url")
    if not video_url:
        return {"error": "no video_url"}

    output_path = os.path.join(output_dir, f"{shot_id}.mp4")
    print(f"[LIB] 下载视频...")
    if download_video(video_url, output_path):
        return {
            "task_id": task_id,
            "output_path": output_path,
            "seed": result.get("seed"),
            "duration": result.get("duration"),
            "resolution": result.get("resolution"),
            "tokens": result.get("usage", {}).get("total_tokens", 0)
        }
    return {"error": "download failed"}

def generate_shot_async(shot_id, prompt, references, duration=6, output_dir="outputs/videos"):
    """异步提交：只提交不等待，返回 task_id 用于后续轮询"""
    task_id = submit_task(prompt, references, duration)
    if isinstance(task_id, dict):
        return task_id
    # 写入任务记录
    from lib.models import get_db
    db = get_db()
    db.execute("""
        INSERT INTO generation_tasks (id, shot_id, task_type, seedance_task_id, status, prompt_used, references_used, started_at)
        VALUES (?, ?, ?, ?, ?, ?, ?, datetime('now'))
    """, (task_id[:20], shot_id, "video", task_id, "running", prompt, json.dumps([r[0] for r in references] if references else [])))
    db.execute("UPDATE shots SET seedance_task_id = ?, status = 'generating' WHERE id = ?", (task_id, shot_id))
    db.commit()
    return {"task_id": task_id, "status": "submitted"}

def check_and_collect(shot_id):
    """检查任务状态，完成后自动下载"""
    from lib.models import get_db
    db = get_db()
    task = db.execute("""
        SELECT * FROM generation_tasks WHERE shot_id = ? ORDER BY created_at DESC LIMIT 1
    """, (shot_id,)).fetchone()
    if not task:
        return {"error": "no task found"}
    result = poll_task(task["seedance_task_id"], poll_interval=0, max_wait=1)
    if "error" in result:
        return result
    if "video_url" not in str(result):
        return {"status": "running", "task_id": task["seedance_task_id"]}

    # Succeeded - download
    video_url = result.get("content", {}).get("video_url")
    output_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "outputs", "videos")
    os.makedirs(output_dir, exist_ok=True)
    output_path = os.path.join(output_dir, f"{shot_id}.mp4")
    if download_video(video_url, output_path):
        db.execute("""
            UPDATE generation_tasks SET status = 'succeeded', completed_at = datetime('now') WHERE id = ?
        """, (task["id"],))
        db.execute("UPDATE shots SET status = 'done', output_path = ? WHERE id = ?", (output_path, shot_id))
        db.commit()
        return {"status": "done", "output_path": output_path, "seed": result.get("seed")}
    return {"error": "download failed"}