#!/usr/bin/env python3
"""
route_shots.py — BC-006 逐镜路由分配器
读取分镜JSON，调豆包AI为每镜分配最优生产路由：
  - STILL_HOLD   → 单帧出图（¥0）
  - LOCAL_MOTION → 出图+FFmpeg运镜（¥0）
  - LAYERED_2_5D → 多层拆解+伪3D运镜（¥0）
  - LOCAL_LIPSYNC → LivePortrait口型同步（¥0）
  - AI_I2V       → LTX 2.3 图生视频（唯一主力 · 本地 ¥0）

用法:
  python tools/route_shots.py <分镜JSON> [-o 输出JSON]
  python tools/route_shots.py <分镜JSON> --dry-run
"""
import sys, os, json, argparse

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from lib.doubao_chat import chat

ROUTE_DESCRIPTION = """
为每个镜头分配最优生产路由，可选路由类型及其含义：

STILL_HOLD: 静态画面/过渡镜头。角色无运动，仅背景。直接单帧出图即可。成本最低。
  - 适合: 空镜、过渡、环境展示、静态对话
  - 成本: ¥0 (仅出图电费)

LOCAL_MOTION: 画面有简单运动（角色小幅度动作、镜头缓慢推拉摇移）。出图后用FFmpeg做缩放/平移/旋转模拟运镜。
  - 适合: 缓慢摇镜、推近拉远、轻微角色动作
  - 成本: ¥0 (FFmpeg运镜)

LAYERED_2_5D: 画面有前后景分层需求。角色在画面中需要与背景分离做视差效果。
  - 适合: 角色在前景有复杂动作但背景固定、需要2.5D伪3D效果
  - 成本: ¥0 (分层合成)

LOCAL_LIPSYNC: 角色有说话/表情需求，需要口型同步。用LivePortrait处理。
  - 适合: 角色特写说话、重要表情表演
  - 成本: ¥0 (LivePortrait本地)

AI_I2V: 画面有复杂运动（角色全身运动、物体移动、镜头剧烈运动）。必须用AI图生视频。
  - 适合: 角色行走奔跑、打斗动作、物体运动、镜头跟随
  - 成本: ¥0 (LTX 2.3 本地 · 唯一主力)
"""

def extract_json(text):
    """从AI回复提取JSON"""
    if '```json' in text:
        return text.split('```json')[1].split('```')[0].strip()
    if '```' in text:
        return text.split('```')[1].split('```')[0].strip()
    maybe = text.strip()
    if maybe.startswith('{') or maybe.startswith('['):
        return maybe
    return maybe

def route_shots(storyboard, dry_run=False):
    """为分镜中所有镜头分配路由"""
    episode = storyboard.get('episode', 1)
    shots = storyboard.get('shots', [])
    
    if not shots:
        print('❌ 分镜中没有镜头')
        return None
    
    # 构建提示词
    shots_text = json.dumps([
        {
            "shot_number": s.get("shot_number"),
            "description": s.get("description"),
            "camera": s.get("camera"),
            "type": s.get("type"),
            "characters": s.get("characters", []),
            "scenes": s.get("scenes", []),
            "dialogue": s.get("dialogue"),
        }
        for s in shots
    ], ensure_ascii=False)
    
    prompt = f"""你是一个短剧视频生产管线路由专家。请为以下短剧第{episode}集每个镜头分配最优生产路由。

路由规则：
{ROUTE_DESCRIPTION}

分配原则：
1. 优先使用低成本路由（STILL_HOLD > LOCAL_MOTION > LAYERED_2_5D > LOCAL_LIPSYNC > AI_I2V）
2. 只有确实需要复杂运动的镜头才用 AI_I2V
3. 有对话的镜头用 LOCAL_LIPSYNC
4. 角色有行走/奔跑/打斗等全身运动用 AI_I2V
5. 空镜/过渡/静态用 STILL_HOLD

输出JSON格式（LIST，按镜头顺序）：
[
  {{
    "shot_number": "S01",
    "route": "STILL_HOLD",
    "reason": "简短理由"
  }},
  ...
]

分镜数据：
{shots_text}"""
    
    if dry_run:
        print('🧪 DRY RUN — 提示词:')
        print('=' * 60)
        print(prompt[:500])
        print('...')
        print('=' * 60)
        return None
    
    print('🚀 调豆包AI分配路由...')
    result = chat(prompt, system="你是短剧生产管线路由优化专家，输出纯JSON。", model="pro", temperature=0.2, max_tokens=4096)
    
    if 'error' in result:
        print('❌ API ERROR:', json.dumps(result, ensure_ascii=False, indent=2))
        return None
    
    content = result.get('content', '')
    json_str = extract_json(content)
    
    try:
        routes = json.loads(json_str)
        if isinstance(routes, dict) and 'shots' in routes:
            routes = routes['shots']
        print(f'✅ 路由分配完成: {len(routes)} 镜')
        return routes
    except json.JSONDecodeError as e:
        print(f'⚠️ JSON解析失败: {e}')
        print('原始回复:', content[:500])
        return None

def main():
    parser = argparse.ArgumentParser(description='BC-006 逐镜路由分配器')
    parser.add_argument('storyboard', help='分镜JSON文件路径')
    parser.add_argument('-o', '--output', help='输出JSON路径 (默认: 覆盖原文件添加路由)')
    parser.add_argument('--dry-run', action='store_true', help='只打印提示词，不调API')
    parser.add_argument('--summary', action='store_true', help='只统计各路由数量，不调API')
    
    args = parser.parse_args()
    
    with open(args.storyboard, 'r', encoding='utf-8') as f:
        storyboard = json.load(f)
    
    shots = storyboard.get('shots', [])
    
    if args.summary:
        print(f'📖 分镜: {args.storyboard}')
        print(f'🎬 共 {len(shots)} 镜')
        print()
        print('🔍 各镜头类型分布:')
        type_count = {}
        for s in shots:
            t = s.get('type', 'unknown')
            type_count[t] = type_count.get(t, 0) + 1
        for t, c in sorted(type_count.items()):
            print(f'  {t}: {c}镜')
        return
    
    routes = route_shots(shots, dry_run=args.dry_run)
    if routes is None:
        sys.exit(1)
    
    # 合并路由到分镜
    route_map = {}
    for r in routes:
        if isinstance(r, dict):
            sn = r.get('shot_number', '')
            route_map[sn] = r.get('route', 'STILL_HOLD')
    
    for shot in shots:
        sn = shot.get('shot_number', '')
        if sn in route_map:
            shot['route'] = route_map[sn]
        else:
            shot['route'] = 'STILL_HOLD'  # 默认
    
    # 统计
    route_count = {}
    for shot in shots:
        r = shot.get('route', 'UNKNOWN')
        route_count[r] = route_count.get(r, 0) + 1
    
    print()
    print('📊 路由分布:')
    for r, c in sorted(route_count.items()):
        cost = {'STILL_HOLD': '¥0', 'LOCAL_MOTION': '¥0', 'LAYERED_2_5D': '¥0', 
                'LOCAL_LIPSYNC': '¥0', 'AI_I2V': '¥0.02'}.get(r, '¥?')
        print(f'  {r:20s}: {c:2d}镜 ({cost})')
    
    total_cost = route_count.get('AI_I2V', 0) * 0.02
    print(f'  {"总预估成本":20s}: ¥{total_cost:.2f}')
    
    # 输出
    out_path = args.output or args.storyboard
    with open(out_path, 'w', encoding='utf-8') as f:
        json.dump(storyboard, f, ensure_ascii=False, indent=2)
    
    print(f'✅ 已保存: {out_path}')

if __name__ == '__main__':
    main()