#!/usr/bin/env python3
"""
Test Character Consistency - Generate Kai walking animation frames
"""

import vertexai
from vertexai.preview.vision_models import ImageGenerationModel, Image
from pathlib import Path

# Initialize
vertexai.init(project="gen-lang-client-0428644398", location="us-central1")

# Load Kai reference (Style A - Cartoon)
kai_reference_path = Path("reference_images/kai_master_stylea.png")
kai_reference = Image.load_from_file(str(kai_reference_path))

print("="*60)
print("🎮 KAI CHARACTER CONSISTENCY TEST")
print("="*60)
print(f"📸 Reference: {kai_reference_path}")
print()

# Animation frames to generate
frames = [
    {
        "name": "kai_walk_frame1",
        "prompt": "Same character Kai exactly as shown, walking animation, left leg forward step, right arm forward. Maintain exact same face, green hair, backpack, blue jeans, and brown boots. Cartoon vector style with bold outlines. Solid bright green background (#00FF00)."
    },
    {
        "name": "kai_walk_frame2", 
        "prompt": "Same character Kai exactly as shown, walking animation, standing upright neutral pose, both feet together. Maintain exact same face, green hair, backpack, blue jeans, and brown boots. Cartoon vector style with bold outlines. Solid bright green background (#00FF00)."
    },
    {
        "name": "kai_walk_frame3",
        "prompt": "Same character Kai exactly as shown, walking animation, right leg forward step, left arm forward. Maintain exact same face, green hair, backpack, blue jeans, and brown boots. Cartoon vector style with bold outlines. Solid bright green background (#00FF00)."
    }
]

model = ImageGenerationModel.from_pretrained("imagegeneration@006")

# Generate frames
for i, frame in enumerate(frames, 1):
    print(f"\n🎨 Generating Frame {i}/3: {frame['name']}")
    print(f"📝 Prompt: {frame['prompt'][:80]}...")
    
    try:
        # Use edit_image with reference for consistency
        response = model.edit_image(
            base_image=kai_reference,
            prompt=frame['prompt'],
            edit_mode="inpainting-insert",
            negative_prompt="different character, different face, different hair color, different clothing, different proportions, watermark, text"
        )
        
        output_path = Path(f"test_character/{frame['name']}.png")
        output_path.parent.mkdir(exist_ok=True)
        
        response.images[0].save(location=str(output_path))
        
        print(f"✅ Saved: {output_path}")
        
    except Exception as e:
        print(f"❌ Error: {e}")

print("\n" + "="*60)
print("✅ TEST COMPLETE!")
print("📁 Check: test_character/ folder")
print("="*60)