#!/usr/bin/env python3
"""
Auto-Regenerate Failing Frames with IP-Adapter
Uses ComfyUI API to regenerate inconsistent frames
"""

import requests
import time
import os

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

# Frames that need regeneration (below 70%)
FAILING_FRAMES = [
    {
        "name": "regen_everly",
        "frame_id": "05_everly_soldier_00002_",
        "reference": "Everly.png",
        "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": "regen_violet",
        "frame_id": "08_violet_devil_00002_",
        "reference": "Violet Devil.png",
        "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": "regen_nova",
        "frame_id": "07_nova_warrior_00002_",
        "reference": "Nova Human.png",
        "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": "regen_eva",
        "frame_id": "06_eva_doctor_00002_",
        "reference": "Éva Moreau.png",
        "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": "regen_tc23",
        "frame_id": "10_tc23_esper_00002_",
        "reference": "TC-23.png",
        "prompt": "masterpiece, best quality, very aesthetic, absurdres, 1woman, solo, mechanical body parts, cybernetic augments, glowing cybernetic eyes, esper powers, glowing energy surrounding, futuristic outfit, robot parts, detailed machinery, dramatic lighting, sci-fi, cyberpunk, esper, mechanical"
    },
]

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"

def check_comfyui_ready():
    """Check if ComfyUI is ready"""
    try:
        resp = requests.get(f"{COMFYUI_URL}/api/system_stats", timeout=10)
        return resp.status_code == 200
    except:
        return False

def queue_ipadapter_prompt(ref_image_name, positive_prompt, filename_prefix):
    """
    Queue a prompt with IP-Adapter
    Note: Full IP-Adapter workflow requires web UI interaction
    This uses the API with enhanced prompting for better consistency
    """
    
    # Enhanced prompt with reference keywords
    enhanced_prompt = f"{positive_prompt}, highly detailed face, detailed eyes, matching reference photo, same character, accurate features"
    
    payload = {
        "prompt": {
            "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "animagine-xl-3.1.safetensors"}},
            "2": {"class_type": "CLIPTextEncode", "inputs": {"text": enhanced_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": 30,  # Higher quality
                "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 wait_for_completion():
    """Wait for queue to clear"""
    while True:
        try:
            resp = requests.get(f"{COMFYUI_URL}/api/queue")
            queue = resp.json()
            if not queue.get('queue_running') and not queue.get('queue_pending'):
                break
            print(f"  Waiting... {len(queue.get('queue_pending', []))} in queue")
            time.sleep(3)
        except:
            break

def main():
    print("=" * 70)
    print("AUTO-REGENERATION WITH IP-ADAPTER")
    print("=" * 70)
    
    if not check_comfyui_ready():
        print("❌ ComfyUI not ready. Start the pod first.")
        return
    
    print("✅ ComfyUI ready")
    
    for frame in FAILING_FRAMES:
        print(f"\n🔄 Regenerating: {frame['name']}")
        print(f"   Reference: {frame['reference']}")
        print(f"   Using enhanced prompt for better consistency...")
        
        result = queue_ipadapter_prompt(
            frame['reference'],
            frame['prompt'],
            frame['name']
        )
        
        if 'prompt_id' in result:
            print(f"   ✅ Queued: {result['prompt_id']}")
            wait_for_completion()
        else:
            print(f"   ❌ Failed: {result}")
        
        time.sleep(2)
    
    print("\n" + "=" * 70)
    print("REGENERATION COMPLETE")
    print("=" * 70)
    print("""
NOTE: Full IP-Adapter with reference images requires:
1. Open ComfyUI: https://bvlnsjffdkgdkt-8188.proxy.runpod.net
2. Load animagine-xl-3.1.safetensors
3. Add IPAdapter Unified Loader → IPAdapter Advanced
4. Load reference image: P:\\...\\Characters images concepts\\{reference}
5. Generate with same prompt
6. Replace in generated folder

The enhanced prompts above may help but IP-Adapter is more reliable.
""")

if __name__ == "__main__":
    main()