#!/usr/bin/env python3
"""
storyboard_to_workflow.py — 分镜JSON → ComfyUI批量出图桥接
读取分镜JSON，为每个镜头生成ComfyUI API调用脚本
用法:
  python tools/storyboard_to_workflow.py <分镜JSON> [-o 输出目录] [--workflow 基础工作流]
  python tools/storyboard_to_workflow.py <分镜JSON> --generate-sh  # 生成批量调用shell脚本
"""
import sys, os, json, argparse, shutil

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 默认的工作流模板路径（ComfyUI蓝图）
DEFAULT_WORKFLOW = os.path.expanduser("~/comfy/ComfyUI/blueprints/Text to Image.json")

def load_storyboard(json_path):
    """读取分镜JSON"""
    with open(json_path, 'r', encoding='utf-8') as f:
        return json.load(f)

def shot_to_prompt(shot, project_style=""):
    """将单镜转换为图像生成prompt"""
    desc = shot.get('description', '')
    camera = shot.get('camera', '中景')
    characters = shot.get('characters', [])
    scenes = shot.get('scenes', [])
    props = shot.get('props', [])
    
    # 构建英文prompt
    parts = []
    if camera:
        parts.append(camera)
    if scenes:
        parts.append(f"in {scenes[0]}")
    parts.append(desc[:80])
    if project_style:
        parts.append(project_style)
    
    prompt = ", ".join(parts)
    return prompt

def generate_comfyui_api_json(shot, base_workflow_path, prompt, output_dir, shot_num):
    """为单镜生成ComfyUI API调用JSON"""
    # 读取基础工作流
    try:
        with open(base_workflow_path, 'r') as f:
            workflow = json.load(f)
    except:
        # 如果文件不存在，创建一个最小工作流
        workflow = {
            "3": {
                "class_type": "KSampler",
                "inputs": {
                    "seed": shot_num * 100 + 42,
                    "steps": 20,
                    "cfg": 7,
                    "sampler_name": "euler",
                    "scheduler": "normal",
                    "denoise": 1,
                    "model": ["4", 0],
                    "positive": ["6", 0],
                    "negative": ["7", 0],
                    "latent_image": ["5", 0]
                }
            },
            "4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "realvisxlV40_v40BDFPonNoobVae.safetensors"}},
            "5": {"class_type": "EmptyLatentImage", "inputs": {"width": 1024, "height": 768, "batch_size": 1}},
            "6": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["4", 1]}},
            "7": {"class_type": "CLIPTextEncode", "inputs": {"text": "worst quality, low quality, blurry", "clip": ["4", 1]}},
            "8": {"class_type": "VAEDecode", "inputs": {"samples": ["3", 0], "vae": ["4", 2]}},
            "9": {"class_type": "SaveImage", "inputs": {"filename_prefix": f"shot_{shot_num:04d}", "images": ["8", 0]}}
        }
        return workflow
    
    # 如果用了现有工作流，替换prompt
    for node_id, node in workflow.items():
        if isinstance(node, dict):
            ct = node.get('class_type', '')
            if 'CLIPTextEncode' in ct or 'Prompt' in ct:
                if 'text' in node.get('inputs', {}):
                    workflow[node_id]['inputs']['text'] = prompt
            if 'KSampler' in ct or 'Sampler' in ct:
                if 'seed' in node.get('inputs', {}):
                    workflow[node_id]['inputs']['seed'] = shot_num * 100 + 42
    
    return workflow

def generate_shell_script(shots, output_dir, comfy_api_url="http://127.0.0.1:8188"):
    """生成批量调用ComfyUI API的shell脚本"""
    script_path = os.path.join(output_dir, 'batch_render.sh')
    
    lines = [
        '#!/bin/bash',
        f'# 批量渲染分镜 - 自动生成',
        f'# ComfyUI API: {comfy_api_url}',
        f'# 分镜数: {len(shots)}',
        f'# 生成时间: {__import__("datetime").datetime.now().isoformat()}',
        '',
        'set -e',
        '',
        'TOTAL=' + str(len(shots)),
        'SUCCESS=0',
        'FAIL=0',
        '',
    ]
    
    for i, shot in enumerate(shots):
        sn = shot.get('shot_number', f'S{i+1:02d}')
        desc = shot.get('description', '')[:40]
        wf_file = os.path.join(output_dir, f'workflow_{i+1:04d}.json')
        out_file = os.path.join(output_dir, f'shot_{i+1:04d}.png')
        
        lines.extend([
            f'',
            f'echo "🎬 [{i+1}/$TOTAL] {sn}: {desc}"',
            f'echo "  → 发送ComfyUI API..."',
            f'RESP=$(curl -s -X POST "{comfy_api_url}/prompt" \\',
            f'  -H "Content-Type: application/json" \\',
            f'  -d @{wf_file})',
            f'if echo "$RESP" | grep -q "error"; then',
            f'  echo "  ❌ 失败: $RESP"',
            f'  FAIL=$((FAIL + 1))',
            f'else',
            f'  echo "  ✅ 已提交"',
            f'  SUCCESS=$((SUCCESS + 1))',
            f'fi',
        ])
    
    lines.extend([
        '',
        'echo "=========================="',
        'echo "📊 渲染完成: $SUCCESS 成功, $FAIL 失败 / $TOTAL 总镜"',
    ])
    
    with open(script_path, 'w') as f:
        f.write('\n'.join(lines))
    os.chmod(script_path, 0o755)
    return script_path

def generate_prompt_file(shots, output_dir, project_style=""):
    """生成每个镜头的prompt文本文件"""
    prompt_path = os.path.join(output_dir, 'prompts.txt')
    with open(prompt_path, 'w', encoding='utf-8') as f:
        for i, shot in enumerate(shots):
            sn = shot.get('shot_number', f'S{i+1:02d}')
            prompt = shot_to_prompt(shot, project_style)
            f.write(f"=== {sn} ===\n")
            f.write(f"景别: {shot.get('camera', 'N/A')}\n")
            f.write(f"时长: {shot.get('duration', 'N/A')}s\n")
            f.write(f"角色: {', '.join(shot.get('characters', []))}\n")
            f.write(f"场景: {', '.join(shot.get('scenes', []))}\n")
            f.write(f"Prompt: {prompt}\n\n")
    return prompt_path

def main():
    parser = argparse.ArgumentParser(description='分镜JSON → ComfyUI批量出图桥接')
    parser.add_argument('storyboard', help='分镜JSON文件路径')
    parser.add_argument('-o', '--output-dir', help='输出目录 (默认: 分镜JSON同目录下的renders/)')
    parser.add_argument('--workflow', default=DEFAULT_WORKFLOW, help=f'基础工作流JSON (默认: {DEFAULT_WORKFLOW})')
    parser.add_argument('--style', default='', help='全局风格描述 (如: cinematic, anime)')
    parser.add_argument('--generate-sh', action='store_true', help='生成批量调用shell脚本')
    parser.add_argument('--comfy-url', default='http://127.0.0.1:8188', help='ComfyUI API地址')
    
    args = parser.parse_args()
    
    # 读取分镜
    storyboard = load_storyboard(args.storyboard)
    shots = storyboard.get('shots', [])
    if not shots:
        print(f'❌ 分镜JSON中没有shots')
        sys.exit(1)
    
    # 输出目录
    if args.output_dir:
        out_dir = args.output_dir
    else:
        sb_dir = os.path.dirname(os.path.abspath(args.storyboard))
        out_dir = os.path.join(sb_dir, 'renders')
    os.makedirs(out_dir, exist_ok=True)
    
    print(f'📖 分镜: {args.storyboard}')
    print(f'🎬 共 {len(shots)} 镜')
    print(f'📁 输出: {out_dir}')
    print()
    
    # 1. 生成prompt文件
    prompt_path = generate_prompt_file(shots, out_dir, args.style)
    print(f'✅ Prompt文件: {prompt_path}')
    
    # 2. 为每镜生成工作流JSON
    wf_dir = os.path.join(out_dir, 'workflows')
    os.makedirs(wf_dir, exist_ok=True)
    
    for i, shot in enumerate(shots):
        sn = shot.get('shot_number', f'S{i+1:02d}')
        prompt = shot_to_prompt(shot, args.style)
        workflow = generate_comfyui_api_json(shot, args.workflow, prompt, wf_dir, i+1)
        
        wf_path = os.path.join(wf_dir, f'workflow_{i+1:04d}.json')
        with open(wf_path, 'w', encoding='utf-8') as f:
            json.dump(workflow, f, ensure_ascii=False, indent=2)
    
    print(f'✅ 工作流JSON: {len(shots)}个 → {wf_dir}/')
    
    # 3. 可选：生成批量shell脚本
    if args.generate_sh:
        sh_path = generate_shell_script(shots, out_dir, args.comfy_url)
        print(f'✅ 批量脚本: {sh_path}')
        print(f'   运行: bash {sh_path}')
    
    # 4. 总结
    print()
    print('📊 各镜一览:')
    for i, shot in enumerate(shots):
        sn = shot.get('shot_number', f'S{i+1:02d}')
        prompt = shot_to_prompt(shot, args.style)
        print(f'  {sn:>4} | {shot.get("camera","?"):　<4} | {shot.get("duration","?"):>2}s | {prompt[:50]}...')

if __name__ == '__main__':
    main()