#!/usr/bin/env python
"""
BaoLife Image Pre-Generation Script

Batch generates images for common events and scenarios.
Run this script to populate the image cache before deploying new events.

Usage:
    python pregenerate_images.py --provider flux --batch-size 10
    python pregenerate_images.py --category education --priority 10
    python pregenerate_images.py --test  # Test with 3 sample images
"""

import asyncio
import argparse
import logging
from typing import List, Dict
from image_generation import (
    generate_and_cache_image,
    add_to_generation_queue,
    process_generation_queue
)

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


# ============================================================
# Event Image Prompts Library
# ============================================================
# Add your event-specific prompts here

EVENT_PROMPTS = {
    # Education
    "education": [
        {
            "prompt": "large contemporary architecture high school building exterior with students",
            "event_type": "school_exterior",
            "event_category": "education",
            "priority": 10
        },
        {
            "prompt": "modern high school classroom with desks and chalkboard",
            "event_type": "classroom",
            "event_category": "education",
            "priority": 10
        },
        {
            "prompt": "school cafeteria with students eating lunch",
            "event_type": "cafeteria",
            "event_category": "education",
            "priority": 9
        },
        {
            "prompt": "school hallway with lockers and students walking",
            "event_type": "school_hallway",
            "event_category": "education",
            "priority": 9
        },
        {
            "prompt": "school library with bookshelves and study tables",
            "event_type": "library",
            "event_category": "education",
            "priority": 8
        },
        {
            "prompt": "outdoor school playground with swings and slides",
            "event_type": "playground",
            "event_category": "education",
            "priority": 8
        },
        {
            "prompt": "school gymnasium with basketball court",
            "event_type": "gymnasium",
            "event_category": "education",
            "priority": 7
        },
        {
            "prompt": "college campus quad with students on grass",
            "event_type": "college_campus",
            "event_category": "education",
            "priority": 9
        },
        {
            "prompt": "large college lecture hall with rows of seats",
            "event_type": "lecture_hall",
            "event_category": "education",
            "priority": 8
        },
    ],

    # Home & Family
    "home": [
        {
            "prompt": "cozy teenage bedroom with desk computer and posters on walls",
            "event_type": "bedroom_teen",
            "event_category": "home",
            "priority": 10
        },
        {
            "prompt": "cozy child bedroom with toys and stuffed animals",
            "event_type": "bedroom_child",
            "event_category": "home",
            "priority": 9
        },
        {
            "prompt": "modern suburban house exterior with driveway and lawn",
            "event_type": "house_exterior",
            "event_category": "home",
            "priority": 9
        },
        {
            "prompt": "warm family dining room with dinner table set for meal",
            "event_type": "dining_room",
            "event_category": "home",
            "priority": 9
        },
        {
            "prompt": "cozy living room with couch and TV",
            "event_type": "living_room",
            "event_category": "home",
            "priority": 8
        },
        {
            "prompt": "modern kitchen with cooking appliances and counters",
            "event_type": "kitchen",
            "event_category": "home",
            "priority": 8
        },
        {
            "prompt": "backyard with green grass and patio furniture",
            "event_type": "backyard",
            "event_category": "home",
            "priority": 7
        },
    ],

    # Family Events
    "family": [
        {
            "prompt": "happy family of four having dinner together at dining table",
            "event_type": "family_dinner",
            "event_category": "family",
            "priority": 10
        },
        {
            "prompt": "family watching TV together on couch in living room",
            "event_type": "family_tv",
            "event_category": "family",
            "priority": 8
        },
        {
            "prompt": "family playing board games together at table",
            "event_type": "family_games",
            "event_category": "family",
            "priority": 7
        },
        {
            "prompt": "mother hugging child with warm expression",
            "event_type": "parent_hug",
            "event_category": "family",
            "priority": 9
        },
        {
            "prompt": "siblings playing together in bedroom",
            "event_type": "siblings_play",
            "event_category": "family",
            "priority": 7
        },
    ],

    # Social & Activities
    "social": [
        {
            "prompt": "teenagers hanging out together at park",
            "event_type": "friends_park",
            "event_category": "social",
            "priority": 9
        },
        {
            "prompt": "group of friends at birthday party with cake",
            "event_type": "birthday_party",
            "event_category": "social",
            "priority": 9
        },
        {
            "prompt": "teenagers at shopping mall walking together",
            "event_type": "mall_friends",
            "event_category": "social",
            "priority": 8
        },
        {
            "prompt": "friends playing video games together on couch",
            "event_type": "gaming_friends",
            "event_category": "social",
            "priority": 8
        },
        {
            "prompt": "teenagers at movie theater watching film",
            "event_type": "movie_theater",
            "event_category": "social",
            "priority": 7
        },
    ],

    # Romance & Dating
    "romance": [
        {
            "prompt": "young couple on romantic date at nice restaurant",
            "event_type": "restaurant_date",
            "event_category": "romance",
            "priority": 9
        },
        {
            "prompt": "couple walking hand in hand in park at sunset",
            "event_type": "romantic_walk",
            "event_category": "romance",
            "priority": 9
        },
        {
            "prompt": "young couple at coffee shop talking and smiling",
            "event_type": "coffee_date",
            "event_category": "romance",
            "priority": 8
        },
        {
            "prompt": "couple watching movie together at home on couch",
            "event_type": "movie_date_home",
            "event_category": "romance",
            "priority": 7
        },
    ],

    # Activities
    "activities": [
        {
            "prompt": "person studying alone at desk with books and laptop",
            "event_type": "studying",
            "event_category": "activities",
            "priority": 10
        },
        {
            "prompt": "person reading book in comfortable chair",
            "event_type": "reading",
            "event_category": "activities",
            "priority": 8
        },
        {
            "prompt": "person exercising at gym with equipment",
            "event_type": "gym",
            "event_category": "activities",
            "priority": 8
        },
        {
            "prompt": "person jogging in park on sunny day",
            "event_type": "jogging",
            "event_category": "activities",
            "priority": 7
        },
        {
            "prompt": "person playing musical instrument in bedroom",
            "event_type": "music_practice",
            "event_category": "activities",
            "priority": 7
        },
    ],

    # Jobs & Career
    "career": [
        {
            "prompt": "modern office with desk computer and office supplies",
            "event_type": "office",
            "event_category": "career",
            "priority": 9
        },
        {
            "prompt": "person working at computer in office setting",
            "event_type": "office_work",
            "event_category": "career",
            "priority": 9
        },
        {
            "prompt": "business meeting with people around conference table",
            "event_type": "meeting",
            "event_category": "career",
            "priority": 8
        },
        {
            "prompt": "retail store interior with products on shelves",
            "event_type": "retail_store",
            "event_category": "career",
            "priority": 7
        },
    ],

    # Special Events
    "special": [
        {
            "prompt": "graduation ceremony with student in cap and gown",
            "event_type": "graduation",
            "event_category": "special",
            "priority": 10
        },
        {
            "prompt": "wedding ceremony with couple at altar",
            "event_type": "wedding",
            "event_category": "special",
            "priority": 10
        },
        {
            "prompt": "hospital room with patient in bed",
            "event_type": "hospital",
            "event_category": "special",
            "priority": 8
        },
        {
            "prompt": "birthday cake with candles",
            "event_type": "birthday_cake",
            "event_category": "special",
            "priority": 9
        },
    ],

    # Emotions & States
    "emotions": [
        {
            "prompt": "person looking sad and depressed sitting alone",
            "event_type": "sad",
            "event_category": "emotions",
            "priority": 7
        },
        {
            "prompt": "person looking happy and excited with big smile",
            "event_type": "happy",
            "event_category": "emotions",
            "priority": 7
        },
        {
            "prompt": "person looking angry with frustrated expression",
            "event_type": "angry",
            "event_category": "emotions",
            "priority": 6
        },
        {
            "prompt": "person looking stressed studying with books everywhere",
            "event_type": "stressed",
            "event_category": "emotions",
            "priority": 7
        },
    ],
}


# ============================================================
# Generation Functions
# ============================================================

async def generate_all_images(provider: str = 'imagen4', use_queue: bool = False):
    """Generate all images from EVENT_PROMPTS library"""
    total = sum(len(prompts) for prompts in EVENT_PROMPTS.values())
    logger.info(f"Starting generation of {total} images using {provider}")

    count = 0
    for category, prompts in EVENT_PROMPTS.items():
        logger.info(f"\n{'='*60}")
        logger.info(f"Category: {category.upper()} ({len(prompts)} images)")
        logger.info(f"{'='*60}")

        for prompt_data in prompts:
            count += 1
            logger.info(f"\n[{count}/{total}] Generating: {prompt_data['event_type']}")

            if use_queue:
                # Add to queue for later processing
                await add_to_generation_queue(
                    prompt=prompt_data['prompt'],
                    event_type=prompt_data['event_type'],
                    event_category=prompt_data['event_category'],
                    priority=prompt_data.get('priority', 5),
                    provider=provider
                )
            else:
                # Generate immediately
                image_url = await generate_and_cache_image(
                    prompt=prompt_data['prompt'],
                    event_type=prompt_data['event_type'],
                    event_category=prompt_data['event_category'],
                    provider=provider,
                    use_cache=False  # Force new generation
                )

                if image_url:
                    logger.info(f"✓ Success: {image_url[:80]}...")
                else:
                    logger.error(f"✗ Failed: {prompt_data['event_type']}")

            # Rate limiting: wait between requests
            await asyncio.sleep(3)

    logger.info(f"\n{'='*60}")
    logger.info(f"Completed: {count} images processed")
    logger.info(f"{'='*60}")


async def generate_category(category: str, provider: str = 'imagen4', use_queue: bool = False):
    """Generate images for a specific category"""
    if category not in EVENT_PROMPTS:
        logger.error(f"Unknown category: {category}")
        logger.info(f"Available categories: {', '.join(EVENT_PROMPTS.keys())}")
        return

    prompts = EVENT_PROMPTS[category]
    logger.info(f"Generating {len(prompts)} images for category: {category}")

    for i, prompt_data in enumerate(prompts, 1):
        logger.info(f"\n[{i}/{len(prompts)}] {prompt_data['event_type']}")

        if use_queue:
            await add_to_generation_queue(
                prompt=prompt_data['prompt'],
                event_type=prompt_data['event_type'],
                event_category=prompt_data['event_category'],
                priority=prompt_data.get('priority', 5),
                provider=provider
            )
        else:
            image_url = await generate_and_cache_image(
                prompt=prompt_data['prompt'],
                event_type=prompt_data['event_type'],
                event_category=prompt_data['event_category'],
                provider=provider,
                use_cache=False
            )

            if image_url:
                logger.info(f"✓ Success: {image_url[:80]}...")
            else:
                logger.error(f"✗ Failed")

        await asyncio.sleep(3)


async def generate_test_batch(provider: str = 'imagen4'):
    """Generate a small test batch of 3 images"""
    test_prompts = [
        {
            "prompt": "large contemporary architecture high school building exterior",
            "event_type": "test_school",
            "event_category": "test",
            "priority": 10
        },
        {
            "prompt": "cozy teenage bedroom with desk and computer",
            "event_type": "test_bedroom",
            "event_category": "test",
            "priority": 10
        },
        {
            "prompt": "happy family having dinner together at dining table",
            "event_type": "test_family",
            "event_category": "test",
            "priority": 10
        }
    ]

    logger.info("Generating test batch of 3 images...")

    for i, prompt_data in enumerate(test_prompts, 1):
        logger.info(f"\n[{i}/3] Testing: {prompt_data['event_type']}")

        image_url = await generate_and_cache_image(
            prompt=prompt_data['prompt'],
            event_type=prompt_data['event_type'],
            event_category=prompt_data['event_category'],
            provider=provider,
            use_cache=False
        )

        if image_url:
            logger.info(f"✓ Success: {image_url}")
        else:
            logger.error(f"✗ Failed")

        await asyncio.sleep(2)

    logger.info("\nTest batch complete!")


# ============================================================
# Main CLI
# ============================================================

async def main():
    parser = argparse.ArgumentParser(
        description="Pre-generate images for BaoLife events"
    )
    parser.add_argument(
        '--provider',
        choices=['imagen4', 'flux', 'dalle3'],
        default='imagen4',
        help='Image generation provider (default: imagen4)'
    )
    parser.add_argument(
        '--category',
        choices=list(EVENT_PROMPTS.keys()) + ['all'],
        default='all',
        help='Generate images for specific category or all'
    )
    parser.add_argument(
        '--queue',
        action='store_true',
        help='Add to queue instead of generating immediately'
    )
    parser.add_argument(
        '--process-queue',
        action='store_true',
        help='Process existing queue items'
    )
    parser.add_argument(
        '--batch-size',
        type=int,
        default=10,
        help='Batch size for queue processing (default: 10)'
    )
    parser.add_argument(
        '--test',
        action='store_true',
        help='Generate test batch of 3 images'
    )

    args = parser.parse_args()

    if args.test:
        await generate_test_batch(provider=args.provider)
    elif args.process_queue:
        logger.info(f"Processing queue (batch size: {args.batch_size})...")
        await process_generation_queue(batch_size=args.batch_size)
    elif args.category == 'all':
        await generate_all_images(provider=args.provider, use_queue=args.queue)
    else:
        await generate_category(
            category=args.category,
            provider=args.provider,
            use_queue=args.queue
        )


if __name__ == "__main__":
    asyncio.run(main())