#!/usr/bin/env python3
"""
Generate anime opening key frames for Broken Spire
Uses ComfyUI API with Animagine XL 3.1
"""

import requests
import time
import json

COMFYUI_URL = "https://mu1aafo3zo4da1-8188.proxy.runpod.net"

FRAMES = [
    {
        "name": "01_void",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, dark void, fractured clock faces spinning backward, glitch effects, purple energy cracks, blue temporal distortion, cinematic lighting, dark atmosphere, epic, no humans"
    },
    {
        "name": "02_ash_birth",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1boy, solo, long blonde hair, glowing blue eyes, black tribal tattoos on arms, slim athletic body, temporal energy surrounding, silhouette, dramatic backlighting, dark cosmic background, intense, masterpiece"
    },
    {
        "name": "03_far_future_ash_evil",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1boy, solo, mature, long white hair, glowing red eyes, dark aura, terrifying expression, menacing smile, multiple arms with weapons, dark armor with glowing red hexcore patterns, floating, surrounded by dark energy, skull decorations, blood splatter, god of death, apocalyptic background, destroyed city, red sky, horror, scary, dark fantasy, terrifying, evil, brutal"
    },
    {
        "name": "04_post_apocalyptic_city",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, cityscape, post-apocalyptic, ruined buildings, crumbling towers, neon signs flickering, blue hexcore energy pillars in distance, augmented soldiers patrolling, rain, ash covered streets, dark cyberpunk, dramatic lighting, cinematic, wide shot, atmospheric, dystopian"
    },
    {
        "name": "05_everly_soldier",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1woman, solo, mature, short dark hair, intense eyes, military uniform, tactical gear, multiple weapons, dark makeup, confident pose, hand on weapon, dark expression, scars visible, dramatic lighting, cinematic, from below, powerful, dangerous, cyberpunk, dark sci-fi"
    },
    {
        "name": "06_eva_doctor",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1woman, solo, french, shoulder length brown hair, intense eyes, lab coat, medical clothing, surgical gloves, focused expression, holding medical scanner, sterile environment, blue medical lights, dramatic lighting, soft features, intelligent, professional, cyberpunk clinic, detailed background"
    },
    {
        "name": "07_nova_warrior",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1woman, solo, young adult, long flowing red hair, fierce expression, combat gear, multiple weapons, dynamic pose, action pose, battle stance, glowing energy weapons, ruins background, explosion in distance, dramatic lighting, cinematic, anime action, intense, stylish, war scene"
    },
    {
        "name": "08_violet_devil",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1woman, solo, long purple hair, red devil horns, glowing red eyes, devil form, dark wings, elegant pose, hand with claws, beautiful and terrifying, dark aura, crimson energy, night background, dramatic lighting, horror, beautiful, duality, dark fantasy"
    },
    {
        "name": "09_lin_weishan",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1woman, solo, asian, long black hair, determined expression, combat clothing, martial arts pose, detailed face, dramatic lighting, cyberpunk, action, intense"
    },
    {
        "name": "10_tc23_esper",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1woman, solo, mechanical augments, glowing cybernetic eyes, esper powers, glowing energy surrounding, futuristic outfit, robot parts, detailed machinery, dramatic lighting, sci-fi, cyberpunk, intense"
    },
    {
        "name": "11_jonas_ghost",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1man, solo, large muscular build, scars on face, short brown hair, warm smile, ghost effect, ethereal, fading, memories, nostalgic, warm lighting, soft focus, emotional, memorial"
    },
    {
        "name": "12_berserk_battle",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, action scene, intense battle, blood splatter, gore, brutal combat, multiple characters fighting, weapons, explosions, dynamic composition, fast action, dark sci-fi, anime style, cinematic, dramatic lighting, violent, intense, berserk style, dark fantasy, horror"
    },
    {
        "name": "13_ash_everly_tension",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1man and 1woman, couple, close together, intense eye contact, dark expression, emotional tension, sexual tension, dramatic lighting, cinematic, close-up, 18+, mature, dark romance, cyberpunk, beautiful"
    },
    {
        "name": "14_ash_eva_medical",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1man and 1woman, couple, medical setting, gentle touch, her hand on his cheek, tender moment, emotional, warm lighting, soft colors, intimate, 18+, mature, romantic, cyberpunk, beautiful, emotional"
    },
    {
        "name": "15_ash_nova_war",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1man and 1woman, couple, sitting together, exhausted, rare smile, battle aftermath, holding hands, warm moment, intimate, dramatic lighting, cinematic, 18+, mature, romantic, war romance, cyberpunk"
    },
    {
        "name": "16_far_future_ash_watching",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1man, solo, far future, white hair flowing, red glowing eyes, dark silhouette, watching from shadows, menacing presence, dark aura, terrifying, horror, suspense, building tension, dramatic lighting, dark sci-fi, threatening"
    },
    {
        "name": "17_all_characters_triangle",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1man and 3women, group, triangular composition, dramatic pose, emotional, intense, hexcore energy connecting them, different timelines, cosmic background, beautiful composition, 18+, mature, dark romance, cyberpunk, cinematic"
    },
    {
        "name": "18_title_card",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, title card, BROKEN SPIRE text, glitch effects, broken sun in background, ruined city silhouette, characters as shadows, cinematic, dramatic, epic, dark sci-fi, poster style"
    },
    {
        "name": "19_ash_walking_hope",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1man, solo, walking away, blonde hair, back view, broken sun on horizon half black hole half burning, shattered city, rain, atmospheric, cinematic, dramatic, hopeful, beautiful, melancholic, epic"
    },
    {
        "name": "20_final_text",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, text EVERY CYCLE BEGINS WITH AN END, dark background, cinematic, dramatic, epic, minimal"
    }
]

NEGATIVE_PROMPT = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name, nsfw"

def queue_prompt(positive_prompt, filename_prefix):
    payload = {
        "prompt": {
            "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "animagine-xl-3.1.safetensors"}},
            "2": {"class_type": "CLIPTextEncode", "inputs": {"text": positive_prompt, "clip": ["1", 1]}},
            "3": {"class_type": "CLIPTextEncode", "inputs": {"text": NEGATIVE_PROMPT, "clip": ["1", 1]}},
            "4": {"class_type": "EmptyLatentImage", "inputs": {"width": 832, "height": 1216, "batch_size": 1}},
            "5": {"class_type": "KSampler", "inputs": {"seed": 42, "steps": 28, "cfg": 7, "sampler_name": "euler_ancestral", "scheduler": "normal", "denoise": 1.0, "model": ["1", 0], "positive": ["2", 0], "negative": ["3", 0], "latent_image": ["4", 0]}},
            "6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}},
            "7": {"class_type": "SaveImage", "inputs": {"filename_prefix": filename_prefix, "images": ["6", 0]}}
        }
    }
    
    response = requests.post(f"{COMFYUI_URL}/api/prompt", json=payload)
    return response.json()

def check_queue():
    response = requests.get(f"{COMFYUI_URL}/api/queue")
    return response.json()

def main():
    print("Starting Broken Spire anime opening frame generation...")
    print(f"Total frames: {len(FRAMES)}")
    
    for i, frame in enumerate(FRAMES):
        print(f"\nGenerating frame {i+1}/20: {frame['name']}")
        print(f"Prompt: {frame['prompt'][:80]}...")
        
        result = queue_prompt(frame['prompt'], frame['name'])
        print(f"Result: {result}")
        
        # Wait for generation to complete before next prompt
        while True:
            queue_status = check_queue()
            if not queue_status.get('queue_running') and not queue_status.get('queue_pending'):
                break
            print(f"  Waiting... queue: {queue_status.get('queue_pending', [])}")
            time.sleep(3)
        
        time.sleep(2)  # Brief pause between frames
    
    print("\nAll frames generated!")
    print(f"View at: {COMFYUI_URL}")

if __name__ == "__main__":
    main()