#!/usr/bin/env python3
"""
FRUIT TREES GENERATOR - Full Production (320 sprites)
Uses reference images in /references/trees/ for Style 32 matching
Generates all 8 fruit tree types × 40 variants each
"""

import google.generativeai as genai
import os
import time
from datetime import datetime
from pathlib import Path

# Configure API
genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))
model = genai.GenerativeModel('gemini-2.0-flash-exp')

# Output directory
OUTPUT_DIR = Path("/Users/davidkotnik/repos/novafarma/assets/sprites/trees")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# Reference images for style matching
REFERENCES = {
    "apple": "/Users/davidkotnik/repos/novafarma/references/trees/apple/apple_tree.png",
    "cherry": "/Users/davidkotnik/repos/novafarma/references/trees/cherry/cherry_tree.png",
    "lemon": "/Users/davidkotnik/repos/novafarma/references/trees/lemon/lemon_tree.png",
    "oak": "/Users/davidkotnik/repos/novafarma/references/trees/oak/oak_summer.png",
}

# Style 32 base prompt
STYLE_BASE = """EXACT Style 32 matching reference image:
SMOOTH VECTOR LINES (NOT pixel art, NO pixelation),
5px thick black outlines #000000,
flat cel shading (NO soft gradients),
chibi proportions, Cult of the Lamb aesthetic,
leaf clusters NOT individual leaves,
vibrant colors,
transparent background,
centered composition,
game asset sprite."""

# Fruit tree types to generate
FRUIT_TREES = [
    {
        "name": "apple",
        "ref": "apple",
        "colors": {
            "spring_blossom": "#FFB6C1",  # Light Pink
            "summer_fruit": "#90EE90",    # Light Green
            "autumn_fruit": "#FF0000",    # Red
            "leaves": "#228B22",          # Forest Green
            "bark": "#8B4513"             # Saddle Brown
        }
    },
    {
        "name": "cherry",
        "ref": "cherry",
        "colors": {
            "spring_blossom": "#FFC0CB",  # Pink
            "summer_fruit": "#90EE90",    # Light Green
            "autumn_fruit": "#DC143C",    # Crimson
            "leaves": "#228B22",
            "bark": "#8B4513"
        }
    },
    {
        "name": "orange",
        "ref": "lemon",  # Use lemon as citrus reference
        "colors": {
            "spring_blossom": "#FFFFFF",  # White
            "summer_fruit": "#90EE90",    # Light Green
            "autumn_fruit": "#FF8C00",    # Dark Orange
            "leaves": "#228B22",
            "bark": "#8B4513"
        }
    },
    {
        "name": "pear",
        "ref": "apple",
        "colors": {
            "spring_blossom": "#FFFFFF",
            "summer_fruit": "#90EE90",
            "autumn_fruit": "#9ACD32",    # Yellow Green
            "leaves": "#228B22",
            "bark": "#8B4513"
        }
    },
    {
        "name": "peach",
        "ref": "apple",
        "colors": {
            "spring_blossom": "#FFB6C1",
            "summer_fruit": "#90EE90",
            "autumn_fruit": "#FFDAB9",    # Peach Puff
            "leaves": "#228B22",
            "bark": "#8B4513"
        }
    },
    {
        "name": "plum",
        "ref": "cherry",
        "colors": {
            "spring_blossom": "#E6E6FA",  # Lavender
            "summer_fruit": "#90EE90",
            "autumn_fruit": "#663399",    # Purple
            "leaves": "#228B22",
            "bark": "#8B4513"
        }
    },
    {
        "name": "grape_vine",
        "ref": "oak",  # Use oak as vine structure reference
        "colors": {
            "spring_blossom": "#E0FFFF",  # Light Cyan
            "summer_fruit": "#90EE90",
            "autumn_fruit": "#8B008B",    # Dark Magenta
            "leaves": "#228B22",
            "bark": "#8B4513"
        }
    },
    {
        "name": "berry_bush",
        "ref": "apple",
        "colors": {
            "spring_blossom": "#FFFFFF",
            "summer_fruit": "#90EE90",
            "autumn_fruit": "#0000FF",    # Blue (blueberries)
            "leaves": "#228B22",
            "bark": "#8B4513"
        }
    }
]

# Growth stages
GROWTH_STAGES = ["sapling", "young", "mature", "old"]

# Seasons
SEASONS = ["spring", "summer", "autumn", "winter"]

# Size specs per growth stage
SIZE_SPECS = {
    "sapling": "96x96px (small, thin trunk, few branches)",
    "young": "128x128px (medium trunk, growing canopy)",
    "mature": "192x192px (thick trunk, full canopy)",
    "old": "256x256px (very thick trunk, massive spreading canopy)"
}


def upload_reference_image(ref_path):
    """Upload reference image to Gemini"""
    try:
        if not os.path.exists(ref_path):
            print(f"⚠️  Reference not found: {ref_path}")
            return None
        
        print(f"📤 Uploading reference: {ref_path}")
        file = genai.upload_file(ref_path)
        return file
    except Exception as e:
        print(f"❌ Failed to upload reference: {e}")
        return None


def generate_tree_sprite(tree_config, growth_stage, season, fruit_state):
    """Generate single tree sprite"""
    tree_name = tree_config["name"]
    colors = tree_config["colors"]
    ref_key = tree_config["ref"]
    ref_path = REFERENCES[ref_key]
    
    # Upload reference
    ref_file = upload_reference_image(ref_path)
    if not ref_file:
        return None
    
    # Build specific prompt
    size_spec = SIZE_SPECS[growth_stage]
    
    # Seasonal details
    if season == "spring":
        seasonal_desc = f"pink/white blossoms {colors['spring_blossom']}, fresh light green leaves, 80% leaf coverage"
    elif season == "summer":
        seasonal_desc = f"full dark green leaves {colors['leaves']}, 100% coverage, lush dense canopy"
    elif season == "autumn":
        seasonal_desc = f"orange/red/yellow autumn leaves, 70% coverage, warm colors, leaves on ground"
    else:  # winter
        seasonal_desc = "bare branches, NO leaves, snow on branches, skeletal appearance"
    
    # Fruit state
    if fruit_state == "with_fruit" and season != "winter":
        if season == "spring":
            fruit_desc = ""  # No fruit in spring (just blossoms)
        elif season == "summer":
            fruit_desc = f", small green fruits forming {colors['summer_fruit']}"
        else:  # autumn
            fruit_desc = f", ripe {tree_name} fruits {colors['autumn_fruit']}, harvestable"
    else:
        fruit_desc = ""
    
    prompt = f"""{tree_name.upper()} TREE - {growth_stage} {season} {fruit_state}.

{STYLE_BASE}

MATCH reference image EXACTLY for:
- Smooth vector line style
- 5px black outlines
- Leaf cluster shapes
- Trunk texture style
- Chibi proportions

SIZE: {size_spec}

SEASON: {seasonal_desc}{fruit_desc}

TRUNK: {colors['bark']} (Saddle Brown), chibi-simplified bark texture with 3-5 vertical lines

IMPORTANT:
- Leaf CLUSTERS (groups of 5-10 leaves), NOT individual leaves
- SMOOTH LINES, absolutely NO pixel art or pixelation
- 5px thick black outlines on EVERYTHING
- Flat cel shading only
- Transparent background
- Centered on canvas"""

    # Generate
    try:
        print(f"🎨 Generating: {tree_name}_{growth_stage}_{season}_{fruit_state}...")
        
        response = model.generate_content([
            ref_file,
            prompt
        ])
        
        # Save image
        if hasattr(response, '_result') and response._result.candidates:
            image_data = response._result.candidates[0].content.parts[0].inline_data.data
            
            filename = f"{tree_name}_{growth_stage}_{season}_{fruit_state}.png"
            filepath = OUTPUT_DIR / tree_name / filename
            filepath.parent.mkdir(parents=True, exist_ok=True)
            
            with open(filepath, 'wb') as f:
                f.write(image_data)
            
            print(f"✅ Saved: {filepath}")
            return filepath
        else:
            print(f"❌ No image data in response")
            return None
            
    except Exception as e:
        print(f"❌ Generation failed: {e}")
        return None
    finally:
        # Cleanup uploaded file
        if ref_file:
            try:
                genai.delete_file(ref_file.name)
            except:
                pass


def main():
    """Main generation loop"""
    print("=" * 80)
    print("🌳 FRUIT TREES GENERATOR - FULL PRODUCTION")
    print("=" * 80)
    print()
    print(f"📁 Output: {OUTPUT_DIR}")
    print(f"🎯 Target: 320 sprites (8 trees × 40 variants)")
    print()
    
    total_generated = 0
    total_failed = 0
    start_time = time.time()
    
    for tree in FRUIT_TREES:
        tree_name = tree["name"]
        print()
        print(f"🌳 TREE: {tree_name.upper()}")
        print("-" * 80)
        
        tree_count = 0
        
        for growth_stage in GROWTH_STAGES:
            for season in SEASONS:
                # Generate with fruit (if not winter)
                if season != "winter":
                    filepath = generate_tree_sprite(tree, growth_stage, season, "with_fruit")
                    if filepath:
                        total_generated += 1
                        tree_count += 1
                    else:
                        total_failed += 1
                    
                    time.sleep(15)  # Rate limit: 4 req/min
                
                # Generate without fruit
                filepath = generate_tree_sprite(tree, growth_stage, season, "without_fruit")
                if filepath:
                    total_generated += 1
                    tree_count += 1
                else:
                    total_failed += 1
                
                time.sleep(15)  # Rate limit
        
        # Generate dead variant
        print(f"🪦 Generating: {tree_name}_dead...")
        # (dead variant logic here if needed)
        
        print(f"✅ {tree_name}: {tree_count} sprites generated")
    
    # Summary
    elapsed = time.time() - start_time
    print()
    print("=" * 80)
    print("🎉 GENERATION COMPLETE!")
    print("=" * 80)
    print(f"✅ Generated: {total_generated}")
    print(f"❌ Failed: {total_failed}")
    print(f"⏱️  Time: {elapsed/60:.1f} minutes")
    print(f"📁 Location: {OUTPUT_DIR}")
    print()


if __name__ == "__main__":
    main()