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

from lib.secrets import get_ark_key
ARK_KEY = get_ark_key()
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"

# 成本估算（元/秒·720p）· 参考价，实际以火山方舟计费为准
COST_PER_SECOND = {"720p": 0.5, "1080p": 1.0}

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()
    estimated_cost = round(duration * COST_PER_SECOND.get("720p", 0.5), 2)
    db.execute("""
        INSERT INTO generation_tasks (id, shot_id, task_type, seedance_task_id, status, prompt_used, references_used, cost_estimated, 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 []), estimated_cost))
    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
    status = result.get("status", "")
    if status in ("running", "queued", "processing"):
        return {"status": "running", "task_id": task["seedance_task_id"]}
    if status != "succeeded":
        return {"error": f"Task {status}: {str(result)[:200]}"}

    # 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', cost_actual = ?, completed_at = datetime('now') WHERE id = ?
        """, (result.get("usage", {}).get("total_tokens", 0), task["id"],))
        # 保存seed
        seed = result.get("seed")
        if seed:
            db.execute("UPDATE shots SET status = 'done', output_path = ?, seed_value = ? WHERE id = ?",
                       (output_path, seed, shot_id))
        else:
            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"}

# ====== 批量并行生成 + 自动重试 ======

def batch_submit(episode_id, max_parallel=3, retry_failed=False):
    """批量提交：自动读库中所有 pending 镜头 → 并行提交 → 并行轮询
    
    支持 Ctrl+C 中断恢复：已提交的任务 ID 都在数据库中，
    中断后用 batch_resume() 继续收集。
    """
    from lib.models import get_db
    db = get_db()

    where = "episode_id = ? AND status IN ('pending', 'failed')" if retry_failed else "episode_id = ? AND status = 'pending'"
    shots = db.execute(f"SELECT * FROM shots WHERE {where} ORDER BY shot_number", (episode_id,)).fetchall()

    if not shots:
        print("⚠️ 没有待生成的镜头")
        return

    print(f"🚀 批量提交 {len(shots)} 个镜头（最多并行 {max_parallel}）")
    print(f"💡 可随时 Ctrl+C 中断，用 batch-resume 恢复\n")

    results = {}
    queue = list(shots)

    while queue or any(r.get("status") == "running" for r in results.values()):
        while len([r for r in results.values() if r.get("status") in ("submitted", "running")]) < max_parallel and queue:
            shot = queue.pop(0)
            shot_id = shot["id"]
            print(f"📤 提交 {shot_id}...")
            refs = db.execute("SELECT sa.role, a.file_path FROM shot_assets sa JOIN assets a ON sa.asset_id = a.id WHERE sa.shot_id = ?", (shot_id,)).fetchall()
            ref_paths = [(r["role"], r["file_path"]) for r in refs] if refs else []
            prompt = shot["prompt"] or shot["description"]
            r = generate_shot_async(shot_id, prompt, ref_paths, int(shot["duration_sec"] or 6))
            if "error" in r:
                results[shot_id] = {"status": "failed", "error": r["error"], "retries": 0}
                print(f"  ❌ {r['error']}")
            else:
                results[shot_id] = {"status": "running", "task_id": r["task_id"], "retries": 0}
                print(f"  ✅ {r['task_id'][:20]}...")

        for shot_id, info in list(results.items()):
            if info["status"] != "running": continue
            result = check_and_collect(shot_id)
            if result.get("status") == "done":
                results[shot_id] = {"status": "done", "output_path": result["output_path"]}
                print(f"✅ {shot_id} 完成")
            elif result.get("error") and "timeout" not in str(result.get("error", "")):
                info["retries"] += 1
                if info["retries"] <= 2:
                    print(f"🔄 {shot_id} 重试 ({info['retries']}/2)")
                    db.execute("UPDATE shots SET status = 'pending' WHERE id = ?", (shot_id,)); db.commit()
                    refs = db.execute("SELECT sa.role, a.file_path FROM shot_assets sa JOIN assets a ON sa.asset_id = a.id WHERE sa.shot_id = ?", (shot_id,)).fetchall()
                    ref_paths = [(r["role"], r["file_path"]) for r in refs]
                    shot = db.execute("SELECT * FROM shots WHERE id = ?", (shot_id,)).fetchone()
                    prompt = shot["prompt"] or shot["description"]
                    new_r = generate_shot_async(shot_id, prompt, ref_paths, int(shot["duration_sec"] or 6))
                    if "error" not in new_r:
                        results[shot_id] = {"status": "running", "task_id": new_r["task_id"], "retries": info["retries"]}
                    else:
                        results[shot_id] = {"status": "failed", "error": new_r["error"], "retries": info["retries"]}
                else:
                    results[shot_id] = {"status": "failed", "retries": info["retries"]}
                    db.execute("UPDATE shots SET status = 'failed', error_log = ? WHERE id = ?", (str(result.get("error", ""))[:200], shot_id)); db.commit()
                    print(f"❌ {shot_id} 重试{info['retries']}次后失败")
        time.sleep(30)

    done = sum(1 for r in results.values() if r["status"] == "done")
    failed = sum(1 for r in results.values() if r["status"] == "failed")
    print(f"\n📊 批量完成: {done}✅  {failed}❌  (共{len(results)})")
    return results


def batch_resume(episode_id):
    """从中断恢复：读数据库中所有 generating 状态的镜头 → 继续轮询收集
    
    适用于 batch_submit 中途 Ctrl+C 后恢复。
    已提交到 Seedance 的任务 ID 保存在 generation_tasks 表中，不会丢失。
    """
    from lib.models import get_db
    db = get_db()

    tasks = db.execute("""
        SELECT gt.shot_id, gt.seedance_task_id, s.shot_number
        FROM generation_tasks gt
        JOIN shots s ON gt.shot_id = s.id
        WHERE s.episode_id = ? AND s.status = 'generating'
        ORDER BY s.shot_number
    """, (episode_id,)).fetchall()

    if not tasks:
        print("⚠️ 没有正在生成的镜头，无需恢复")
        return

    print(f"🔄 恢复 {len(tasks)} 个正在生成的镜头...\n")

    results = {}
    for task in tasks:
        shot_id = task["shot_id"]
        print(f"📥 检查 {shot_id} (task: {task['seedance_task_id'][:20]}...)")
        result = check_and_collect(shot_id)
        if result.get("status") == "done":
            results[shot_id] = "done"
            print(f"  ✅ 已完成")
        elif result.get("status") == "running":
            results[shot_id] = "running"
            print(f"  🔄 仍在生成中")
        else:
            results[shot_id] = "unknown"
            print(f"  ❓ {result}")

    done = sum(1 for v in results.values() if v == "done")
    running = sum(1 for v in results.values() if v == "running")
    print(f"\n📊 恢复结果: {done}✅  {running}🔄  (共{len(results)})")
    if running > 0:
        print(f"💡 还有 {running} 个镜头在生成中，稍后再次运行 batch-resume")
    return results