#!/usr/bin/env python3
"""
Vertex AI Imagen Bulk Generator - Dino Valley
Generates all missing Dino Valley assets using Google Vertex AI Imagen 3
"""

import os
import json
import time
from pathlib import Path
from vertexai.preview.vision_models import ImageGenerationModel

# Initialize Vertex AI
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.expanduser(
    "~/.config/gcloud/application_default_credentials.json"
)

PROJECT_ID = "gen-lang-client-0428644398"
LOCATION = "us-central1"

# Style prompts
STYLE_A_SUFFIX = """
2D game asset, SMOOTH CHIBI-STYLE VECTOR ART (NOT pixel art), clean anti-aliased edges, 
bold thick black outlines, flat vibrant bright colors, cute playful proportions, 
simple clean shapes, friendly cartoon aesthetic. Reference style: cute anime/manga chibi characters,
smooth polished vector art. Solid bright chroma key green background (#00FF00), 
no other elements, centered subject, game-ready asset.
"""

STYLE_B_SUFFIX = """
2D game asset, dark hand-drawn gritty noir style with dramatic shadows, high contrast, 
sketchy atmospheric lines, moody lighting, mature post-apocalyptic aesthetic. 
Solid bright chroma key green background (#00FF00), no other elements, centered subject, game-ready asset.
"""

def load_missing_assets():
    """Load missing assets JSON"""
    json_path = Path("/Users/davidkotnik/repos/novafarma/scripts/dino_valley_missing.json")
    with open(json_path, 'r') as f:
        return json.load(f)

def get_category_description(category):
    """Get detailed description for each category"""
    descriptions = {
        "fauna": "prehistoric dinosaur creature",
        "teren": "tileable terrain ground tile",
        "vegetacija": "prehistoric plant or vegetation",
        "rekviziti": "prop object or environmental decoration",
        "zgradbe": "building or structure",
        "hrana": "food item",
        "materiali": "crafting material or resource",
        "oblacila": "clothing or armor piece",
        "orodja": "tool or weapon",
        "npcs": "non-player character, full body sprite, clearly visible eyes with pupils and highlights, detailed facial expression"
    }
    return descriptions.get(category, "game asset")

def generate_prompt(item_name, category, style):
    """Generate detailed prompt for Vertex AI"""
    category_desc = get_category_description(category)
    
    # Clean up item name for prompt
    display_name = item_name.replace("_", " ").replace("dino ", "")
    
    # Build base prompt
    base_prompt = f"{category_desc}: {display_name}, "
    
    # Add style
    if style == "A":
        prompt = base_prompt + STYLE_A_SUFFIX
    else:
        prompt = base_prompt + STYLE_B_SUFFIX
    
    return prompt

def generate_image_vertex(prompt, output_path):
    """Generate single image using Vertex AI Imagen"""
    try:
        model = ImageGenerationModel.from_pretrained("imagegeneration@006")
        
        response = model.generate_images(
            prompt=prompt,
            number_of_images=1,
            aspect_ratio="1:1",
            safety_filter_level="block_some",
            person_generation="allow_adult"
        )
        
        if response.images:
            # Save image
            response.images[0].save(output_path)
            return True
        else:
            print(f"  ❌ No images generated for: {output_path.name}")
            return False
            
    except Exception as e:
        print(f"  ⚠️ Error generating {output_path.name}: {str(e)}")
        return False

def generate_dino_valley():
    """Main generation loop"""
    missing = load_missing_assets()
    base_path = Path("/Users/davidkotnik/repos/novafarma/assets/slike/biomi/dino_valley")
    
    total_items = sum(len(cat["items"]) for cat in missing.values())
    total_png = total_items * 2  # Both styles
    
    print("=" * 60)
    print("🦖 DINO VALLEY BULK GENERATION STARTING")
    print("=" * 60)
    print(f"Total items: {total_items}")
    print(f"Total PNG to generate: {total_png}")
    print(f"Estimated time: {total_png / 6:.1f} minutes")
    print("=" * 60)
    print()
    
    generated_count = 0
    failed_count = 0
    
    for category, data in missing.items():
        if not data["items"]:
            continue
            
        category_path = base_path / category
        category_path.mkdir(parents=True, exist_ok=True)
        
        print(f"\n📁 CATEGORY: {category.upper()}")
        print(f"Items to generate: {len(data['items'])}")
        print("-" * 60)
        
        for item in data["items"]:
            # Generate Style A
            prompt_a = generate_prompt(item, category, "A")
            output_a = category_path / f"{item}_stylea.png"
            
            print(f"  🎨 Generating: {item}_stylea.png")
            if generate_image_vertex(prompt_a, output_a):
                generated_count += 1
                print(f"    ✅ Success! ({generated_count}/{total_png})")
            else:
                failed_count += 1
                
            time.sleep(1)  # Small delay between requests
            
            # Generate Style B
            prompt_b = generate_prompt(item, category, "B")
            output_b = category_path / f"{item}_styleb.png"
            
            print(f"  🎨 Generating: {item}_styleb.png")
            if generate_image_vertex(prompt_b, output_b):
                generated_count += 1
                print(f"    ✅ Success! ({generated_count}/{total_png})")
            else:
                failed_count += 1
                
            time.sleep(1)
            
            # Progress update every 10 items
            if (generated_count + failed_count) % 20 == 0:
                progress_pct = (generated_count / total_png) * 100
                print(f"\n  📊 PROGRESS: {generated_count}/{total_png} ({progress_pct:.1f}%)")
                print(f"  ⏱️  Estimated remaining: {(total_png - generated_count) / 6:.1f} minutes\n")
    
    print("\n" + "=" * 60)
    print("🎉 DINO VALLEY GENERATION COMPLETE!")
    print("=" * 60)
    print(f"✅ Generated: {generated_count}/{total_png}")
    print(f"❌ Failed: {failed_count}")
    print("=" * 60)

if __name__ == "__main__":
    import vertexai
    vertexai.init(project=PROJECT_ID, location=LOCATION)
    generate_dino_valley()