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

Interactive demo to test different prompts and see generated images.
Useful for fine-tuning prompt styles and comparing providers.

Usage:
    python demo_prompts.py
    python demo_prompts.py --provider dalle3
    python demo_prompts.py --interactive
"""

import asyncio
import argparse
import logging
from typing import List, Dict
from image_generation import ImageGenerator, COZY_CARTOON_STYLE

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


# ============================================================
# Demo Prompt Collections
# ============================================================

DEMO_PROMPTS = {
    "Quick Test (3 images)": [
        {
            "name": "High School",
            "prompt": "large contemporary architecture high school building exterior",
            "description": "Testing architectural scene generation"
        },
        {
            "name": "Bedroom",
            "prompt": "cozy teenage bedroom with desk computer and posters on walls",
            "description": "Testing interior scene with details"
        },
        {
            "name": "Family Dinner",
            "prompt": "happy family of four having dinner together at dining table",
            "description": "Testing characters and emotional warmth"
        }
    ],

    "Architecture & Buildings": [
        {
            "name": "Modern High School",
            "prompt": "large contemporary architecture high school building exterior with glass windows"
        },
        {
            "name": "Cozy House",
            "prompt": "modern suburban house exterior with driveway lawn and blue sky"
        },
        {
            "name": "College Campus",
            "prompt": "beautiful college campus quad with brick buildings and students on grass"
        },
        {
            "name": "Shopping Mall",
            "prompt": "modern shopping mall interior with stores and bright lighting"
        }
    ],

    "Interior Rooms": [
        {
            "name": "Teen Bedroom",
            "prompt": "cozy teenage bedroom with desk computer posters and messy bed"
        },
        {
            "name": "Classroom",
            "prompt": "modern high school classroom with desks chalkboard and bright windows"
        },
        {
            "name": "Living Room",
            "prompt": "cozy family living room with couch TV and warm lighting"
        },
        {
            "name": "Office",
            "prompt": "modern office workspace with computer desk and organized supplies"
        }
    ],

    "People & Emotions": [
        {
            "name": "Happy Student",
            "prompt": "teenage student looking happy and excited with big smile"
        },
        {
            "name": "Stressed Studying",
            "prompt": "student looking stressed studying with books and papers everywhere"
        },
        {
            "name": "Friends Laughing",
            "prompt": "group of teenage friends laughing together at park"
        },
        {
            "name": "Romantic Couple",
            "prompt": "young couple walking hand in hand in park at sunset"
        }
    ],

    "Activities": [
        {
            "name": "Studying",
            "prompt": "student studying alone at desk with books laptop and focused expression"
        },
        {
            "name": "Playing Sports",
            "prompt": "teenagers playing basketball on outdoor court"
        },
        {
            "name": "Video Gaming",
            "prompt": "friends playing video games together on couch with controllers"
        },
        {
            "name": "Reading",
            "prompt": "person reading book in comfortable chair by window"
        }
    ],

    "Special Moments": [
        {
            "name": "Birthday Party",
            "prompt": "birthday party with cake candles balloons and happy friends"
        },
        {
            "name": "Graduation",
            "prompt": "graduation ceremony with student in cap and gown celebrating"
        },
        {
            "name": "First Date",
            "prompt": "nervous young couple on first date at cozy restaurant"
        },
        {
            "name": "Wedding",
            "prompt": "beautiful wedding ceremony with bride and groom at altar"
        }
    ]
}


# ============================================================
# Demo Functions
# ============================================================

async def demo_category(category: str, prompts: List[Dict], provider: str = 'imagen4'):
    """Run demo for a category of prompts"""
    print(f"\n{'='*70}")
    print(f"  {category}")
    print(f"{'='*70}\n")

    generator = ImageGenerator(provider=provider)

    for i, prompt_data in enumerate(prompts, 1):
        name = prompt_data['name']
        prompt = prompt_data['prompt']
        description = prompt_data.get('description', '')

        print(f"\n[{i}/{len(prompts)}] {name}")
        print(f"Prompt: {prompt}")
        if description:
            print(f"Testing: {description}")
        print(f"Provider: {provider.upper()}")
        print(f"Style: Cozy Cartoon")
        print("-" * 70)

        # Generate image
        image_url, error = await generator.generate_image(prompt, 'cozy_cartoon')

        if image_url:
            print(f"✓ SUCCESS")
            print(f"Image URL: {image_url}")
            print(f"")
            print(f"🖼️  View in browser: {image_url}")
        else:
            print(f"✗ FAILED")
            print(f"Error: {error}")

        print()

        # Wait between requests
        if i < len(prompts):
            print("Waiting 3 seconds before next generation...")
            await asyncio.sleep(3)

    print(f"\n{'='*70}")
    print(f"  Completed: {category}")
    print(f"{'='*70}\n")


async def demo_comparison(prompt_data: Dict, providers: List[str]):
    """Compare same prompt across different providers"""
    name = prompt_data['name']
    prompt = prompt_data['prompt']

    print(f"\n{'='*70}")
    print(f"  Provider Comparison: {name}")
    print(f"{'='*70}")
    print(f"Prompt: {prompt}\n")

    results = {}

    for provider in providers:
        print(f"\nGenerating with {provider.upper()}...")
        generator = ImageGenerator(provider=provider)
        image_url, error = await generator.generate_image(prompt, 'cozy_cartoon')

        results[provider] = {
            'success': image_url is not None,
            'url': image_url,
            'error': error
        }

        if image_url:
            print(f"✓ {provider.upper()}: {image_url[:60]}...")
        else:
            print(f"✗ {provider.upper()} failed: {error}")

        await asyncio.sleep(3)

    print(f"\n{'='*70}")
    print(f"  Comparison Results")
    print(f"{'='*70}\n")

    for provider, result in results.items():
        status = "✓ SUCCESS" if result['success'] else "✗ FAILED"
        print(f"{provider.upper():12} {status}")
        if result['url']:
            print(f"             {result['url']}")
        print()


async def interactive_mode():
    """Interactive prompt testing"""
    print(f"\n{'='*70}")
    print(f"  Interactive Prompt Testing")
    print(f"{'='*70}\n")

    print("Enter custom prompts to generate images.")
    print("The cozy cartoon style will be applied automatically.")
    print("Type 'quit' or 'exit' to stop.\n")

    provider = input("Choose provider (imagen4/flux/dalle3) [imagen4]: ").strip().lower() or 'imagen4'
    generator = ImageGenerator(provider=provider)

    while True:
        print("\n" + "-" * 70)
        prompt = input("\nEnter prompt (or 'quit'): ").strip()

        if prompt.lower() in ['quit', 'exit', 'q']:
            print("\nExiting interactive mode.")
            break

        if not prompt:
            print("Please enter a prompt.")
            continue

        print(f"\nGenerating with {provider.upper()}...")
        print(f"Full prompt: {prompt}, {COZY_CARTOON_STYLE}")

        image_url, error = await generator.generate_image(prompt, 'cozy_cartoon')

        if image_url:
            print(f"\n✓ SUCCESS")
            print(f"Image URL: {image_url}")
            print(f"\n🖼️  View: {image_url}")
        else:
            print(f"\n✗ FAILED")
            print(f"Error: {error}")


async def run_full_demo(provider: str = 'imagen4', categories: List[str] = None):
    """Run full demo with all or selected categories"""
    print(f"\n{'#'*70}")
    print(f"#  BaoLife Image Generation Demo")
    print(f"#  Provider: {provider.upper()}")
    print(f"#  Style: Cozy Cartoon")
    print(f"{'#'*70}\n")

    if categories:
        selected = {k: v for k, v in DEMO_PROMPTS.items() if k in categories}
    else:
        selected = DEMO_PROMPTS

    total_prompts = sum(len(prompts) for prompts in selected.values())
    print(f"Generating {total_prompts} images across {len(selected)} categories...\n")

    for category, prompts in selected.items():
        await demo_category(category, prompts, provider)

    print(f"\n{'#'*70}")
    print(f"#  Demo Complete!")
    print(f"#  Total Images: {total_prompts}")
    print(f"#  Estimated Cost: ${total_prompts * 0.04:.2f}")
    print(f"{'#'*70}\n")


# ============================================================
# CLI Interface
# ============================================================

async def main():
    parser = argparse.ArgumentParser(
        description="Demo AI image generation with various prompts",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Quick test (3 images)
  python demo_prompts.py --quick

  # Test specific category
  python demo_prompts.py --category "Interior Rooms"

  # Compare providers
  python demo_prompts.py --compare

  # Interactive mode
  python demo_prompts.py --interactive

  # Full demo with DALL-E 3
  python demo_prompts.py --full --provider dalle3

  # Full demo with Imagen 4 Ultra (recommended)
  python demo_prompts.py --full --provider imagen4
        """
    )

    parser.add_argument(
        '--provider',
        choices=['imagen4', 'flux', 'dalle3'],
        default='imagen4',
        help='Image generation provider (default: imagen4)'
    )

    parser.add_argument(
        '--quick',
        action='store_true',
        help='Run quick test with 3 images'
    )

    parser.add_argument(
        '--full',
        action='store_true',
        help='Run full demo with all categories'
    )

    parser.add_argument(
        '--category',
        choices=list(DEMO_PROMPTS.keys()),
        help='Test specific category'
    )

    parser.add_argument(
        '--compare',
        action='store_true',
        help='Compare FLUX vs DALL-E 3 side-by-side'
    )

    parser.add_argument(
        '--interactive',
        action='store_true',
        help='Interactive prompt testing mode'
    )

    parser.add_argument(
        '--list',
        action='store_true',
        help='List all available categories'
    )

    args = parser.parse_args()

    # List categories
    if args.list:
        print("\nAvailable Categories:")
        for i, (category, prompts) in enumerate(DEMO_PROMPTS.items(), 1):
            print(f"{i}. {category} ({len(prompts)} prompts)")
        print()
        return

    # Interactive mode
    if args.interactive:
        await interactive_mode()
        return

    # Comparison mode
    if args.compare:
        print("\nComparing providers with sample prompt...")
        sample = DEMO_PROMPTS["Quick Test (3 images)"][0]
        await demo_comparison(sample, ['imagen4', 'flux', 'dalle3'])
        return

    # Quick test
    if args.quick:
        await demo_category(
            "Quick Test (3 images)",
            DEMO_PROMPTS["Quick Test (3 images)"],
            args.provider
        )
        return

    # Specific category
    if args.category:
        await demo_category(
            args.category,
            DEMO_PROMPTS[args.category],
            args.provider
        )
        return

    # Full demo
    if args.full:
        await run_full_demo(provider=args.provider)
        return

    # Default: show help
    parser.print_help()


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