# AI Image Creator: Enhanced UI and UX
# Part 1: Core Setup, Model Classes and API Configuration

import gradio as gr
import logging
import sys
import random
import time
from huggingface_hub import InferenceClient
from PIL import Image
import io
import base64

os.getenv("HF_API_KEY")

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger("ai_image_creator")

# =============== MODEL CLIENTS SETUP ===============
def setup_client(api_key, provider=None):
    """Initialize and return API client"""
    try:
        if provider:
            client = InferenceClient(provider=provider, api_key=api_key)
            logger.info(f"{provider} client initialized successfully")
        else:
            client = InferenceClient(api_key=api_key)
            logger.info("Hugging Face client initialized successfully")
        return client
    except Exception as e:
        logger.error(f"Error initializing client: {str(e)}")
        return None

# Initialize clients
try:
    # Replace with your actual HF API key
    hf_api_key = os.getenv("HF_API_KEY")
    hf_client = setup_client(hf_api_key)
    logger.info("Hugging Face client created successfully")
    
    # Set up Llama client if API key is provided
    llama_api_key = os.getenv("HF_API_KEY")  # Replace with actual key if available
    try:
        llama_client = setup_client(llama_api_key, "sambanova")
        use_llama = True
        logger.info("Llama client created successfully")
    except Exception as e:
        logger.warning(f"Llama client not available: {str(e)}. Will use fallback enhancement.")
        llama_client = None
        use_llama = False
except Exception as e:
    logger.error(f"Failed to create Hugging Face client: {str(e)}")
    hf_client = None
    llama_client = None
    use_llama = False

# =============== DATA MODELS ===============

# Image Models with friendly names, descriptions and icons
IMAGE_MODELS = {
    "stabilityai/stable-diffusion-xl-base-1.0": {
        "display_name": "SDXL 1.0", 
        "description": "Best overall quality, slower generation",
        "icon": "⭐",
        "speed": "slow",
        "quality": "excellent"
    },
    "runwayml/stable-diffusion-v1-5": {
        "display_name": "SD 1.5", 
        "description": "Good for general purpose, faster generation",
        "icon": "🚀",
        "speed": "fast",
        "quality": "good"
    },
    "stabilityai/stable-diffusion-2-1": {
        "display_name": "SD 2.1", 
        "description": "Improved details, balanced speed and quality",
        "icon": "✨",
        "speed": "medium",
        "quality": "very good"
    },
    "prompthero/openjourney": {
        "display_name": "OpenJourney", 
        "description": "Midjourney-like stylized results",
        "icon": "🎨",
        "speed": "medium",
        "quality": "stylized"
    },
    "dreamlike-art/dreamlike-diffusion-1.0": {
        "display_name": "Dreamlike", 
        "description": "Artistic style with dreamy aesthetics",
        "icon": "💫",
        "speed": "medium",
        "quality": "artistic"
    }
}

# Creation types with icons and detailed descriptions
CREATION_TYPES = {
    "Realistic Photo": {
        "description": "Create a photorealistic image with natural details and lighting",
        "icon": "📷",
        "prompt_hint": "Try to include details about lighting, time of day, and environment"
    },
    "Digital Art": {
        "description": "Create colorful digital artwork with clean lines and vibrant colors",
        "icon": "🖌️",
        "prompt_hint": "Consider specifying color palette and mood for better results"
    },
    "Fantasy Illustration": {
        "description": "Create magical and fantastical scenes with otherworldly elements",
        "icon": "🧙",
        "prompt_hint": "Describe magical elements, creatures, and environments in detail"
    },
    "Concept Art": {
        "description": "Create professional concept art for characters, environments or objects",
        "icon": "🎮",
        "prompt_hint": "Include details about perspective, purpose, and design influences"
    },
    "Anime/Manga": {
        "description": "Create Japanese anime or manga style illustration",
        "icon": "🍙",
        "prompt_hint": "Specify anime aesthetics like shading style and character features"
    },
    "Oil Painting": {
        "description": "Create an image with oil painting textures and artistic brushstrokes",
        "icon": "🖼️",
        "prompt_hint": "Consider describing texture, brushwork style, and canvas feel"
    },
    "Watercolor": {
        "description": "Create a soft watercolor illustration with subtle color blending",
        "icon": "💧",
        "prompt_hint": "Mention color blending, paper texture, and watercolor-specific effects"
    },
    "Sketch": {
        "description": "Create a detailed sketch or drawing with line art focus",
        "icon": "✏️",
        "prompt_hint": "Describe line weight, hatching style, and sketch medium (pencil, charcoal, etc.)"
    },
    "3D Rendering": {
        "description": "Create an image that looks like a 3D rendered scene with realistic lighting",
        "icon": "💻",
        "prompt_hint": "Include details about lighting setup, materials, and camera perspective"
    },
    "Pixel Art": {
        "description": "Create retro-style pixel art with limited color palette",
        "icon": "👾",
        "prompt_hint": "Specify resolution, color limitations, and pixel art style (e.g., 16-bit, 8-bit)"
    }
}

# Art styles with icons and detailed descriptions
ART_STYLES = {
    "Photorealistic": {
        "description": "Detailed realistic style that resembles a photograph with accurate lighting and textures",
        "icon": "📸",
        "examples": "Works by Chuck Close, Richard Estes, or modern 3D renderings"
    },
    "Impressionist": {
        "description": "Soft brushstrokes that capture light and atmosphere over precise details, like Monet",
        "icon": "🌈",
        "examples": "Claude Monet, Pierre-Auguste Renoir, Camille Pissarro"
    },
    "Surrealist": {
        "description": "Dreamlike quality with impossible or irrational scenes, like Salvador Dali",
        "icon": "🌀",
        "examples": "Salvador Dali, René Magritte, Frida Kahlo"
    },
    "Pop Art": {
        "description": "Bold colors, sharp lines and popular culture references, like Andy Warhol",
        "icon": "🎭",
        "examples": "Andy Warhol, Roy Lichtenstein, Keith Haring"
    },
    "Minimalist": {
        "description": "Simplified forms, limited color palette, and clean composition with minimal elements",
        "icon": "⬜",
        "examples": "Piet Mondrian, Kazimir Malevich, Agnes Martin"
    },
    "Abstract": {
        "description": "Non-representational style using shapes, colors, and forms to express ideas",
        "icon": "🔶",
        "examples": "Wassily Kandinsky, Jackson Pollock, Mark Rothko"
    },
    "Cubist": {
        "description": "Geometric shapes and multiple perspectives shown simultaneously, like Picasso",
        "icon": "📐",
        "examples": "Pablo Picasso, Georges Braque, Juan Gris"
    },
    "Art Nouveau": {
        "description": "Ornate, flowing lines inspired by natural forms with decorative elegance",
        "icon": "🌿",
        "examples": "Alphonse Mucha, Gustav Klimt, Antoni Gaudí"
    },
    "Gothic": {
        "description": "Dark, medieval-inspired aesthetic with dramatic lighting and architectural elements",
        "icon": "🏰",
        "examples": "Zdzisław Beksiński, H.R. Giger, medieval architecture"
    },
    "Cyberpunk": {
        "description": "Futuristic dystopian style with neon colors, technology, and urban decay",
        "icon": "🤖",
        "examples": "Blade Runner, Ghost in the Shell, Akira"
    },
    "Steampunk": {
        "description": "Victorian-era aesthetic combined with steam-powered technology and brass elements",
        "icon": "⚙️",
        "examples": "Works by James Ng, Keith Thompson, retrofuturistic Jules Verne adaptations"
    },
    "Retro/Vintage": {
        "description": "Nostalgic style reminiscent of past decades with period-appropriate elements",
        "icon": "📺",
        "examples": "1950s advertisements, vintage travel posters, pulp magazine covers"
    },
    "Art Deco": {
        "description": "Geometric patterns, bold colors, and luxurious materials in a symmetrical style",
        "icon": "🏢",
        "examples": "Works from the 1920s-30s, Chrysler Building, Tamara de Lempicka paintings"
    },
    "Baroque": {
        "description": "Dramatic, ornate style with rich details, contrast, and dynamic composition",
        "icon": "👑",
        "examples": "Caravaggio, Rembrandt, Peter Paul Rubens"
    },
    "Ukiyo-e": {
        "description": "Traditional Japanese woodblock print style with flat areas of color and strong outlines",
        "icon": "🌊",
        "examples": "Hokusai's Great Wave, Hiroshige's landscapes, traditional Japanese prints"
    },
    "Comic Book": {
        "description": "Bold outlines, bright colors, and action-oriented composition like classic comics",
        "icon": "💥",
        "examples": "Jack Kirby, Steve Ditko, modern Marvel/DC art styles"
    },
    "Psychedelic": {
        "description": "Vibrant, swirling colors with abstract patterns inspired by 1960s art",
        "icon": "🌈",
        "examples": "1960s concert posters, Peter Max, Alex Grey"
    },
    "Vaporwave": {
        "description": "Glitch aesthetics with pastel colors, 80s/90s nostalgia and digital elements",
        "icon": "📼",
        "examples": "Retrowave aesthetics, 80s digital graphics, glitch art"
    },
    "Studio Ghibli": {
        "description": "Whimsical, detailed animation style inspired by Japanese animated films",
        "icon": "🐉",
        "examples": "Spirited Away, My Neighbor Totoro, Howl's Moving Castle"
    },
    "Hyperrealism": {
        "description": "Extremely detailed realism that exceeds photograph-like precision",
        "icon": "🔍",
        "examples": "Works by Roberto Bernardi, Denis Peterson, Gottfried Helnwein"
    }
}

# Moods with icons and descriptions
MOODS = {
    "Happy": {
        "description": "Bright, cheerful atmosphere with warm colors",
        "icon": "😊",
        "color_palette": "Warm and vibrant colors: yellows, bright oranges, light blues"
    },
    "Sad": {
        "description": "Melancholic atmosphere with muted colors",
        "icon": "😢",
        "color_palette": "Muted blues, grays, desaturated colors, cool tones"
    },
    "Mysterious": {
        "description": "Enigmatic atmosphere with shadows and hidden elements",
        "icon": "🔮",
        "color_palette": "Deep purples, dark blues, hints of teal, selective lighting"
    },
    "Peaceful": {
        "description": "Serene, calm atmosphere with gentle lighting",
        "icon": "🕊️",
        "color_palette": "Soft blues, gentle greens, pastel colors, balanced light"
    },
    "Tense": {
        "description": "Suspenseful atmosphere with dramatic lighting",
        "icon": "😰",
        "color_palette": "High contrast, dark shadows, selective reds, strong highlights"
    },
    "Whimsical": {
        "description": "Playful, imaginative atmosphere with fanciful elements",
        "icon": "🦄",
        "color_palette": "Pastels, candy colors, unexpected color combinations"
    },
    "Dark": {
        "description": "Gloomy atmosphere with deep shadows and low lighting",
        "icon": "🌑",
        "color_palette": "Dark blues, blacks, deep greens, minimal highlights"
    },
    "Energetic": {
        "description": "Dynamic, vibrant atmosphere with strong colors and movement",
        "icon": "⚡",
        "color_palette": "Saturated primary colors, bold contrasts, vibrant hues"
    },
    "Romantic": {
        "description": "Soft, dreamy atmosphere with warm, gentle lighting",
        "icon": "❤️",
        "color_palette": "Soft pinks, gentle reds, golden highlights, warm tones"
    },
    "Epic": {
        "description": "Grand, impressive atmosphere with dramatic scale and lighting",
        "icon": "🏔️",
        "color_palette": "Bold colors, dramatic contrast, atmospheric lighting, expansive scale"
    }
}

# Example prompts with rich metadata
EXAMPLE_PROMPTS = [
    {
        "text": "A serene lake at sunset with mountains in the background and a small wooden boat floating nearby",
        "thumbnail_desc": "Peaceful lake scene at sunset",
        "creation_type": "Realistic Photo",
        "art_style": "Photorealistic", 
        "mood": "Peaceful",
        "tags": ["nature", "landscape", "water", "sunset"]
    },
    {
        "text": "A futuristic cityscape with flying cars, neon lights, and tall skyscrapers under a night sky with two moons",
        "thumbnail_desc": "Futuristic city with flying cars",
        "creation_type": "Concept Art",
        "art_style": "Cyberpunk",
        "mood": "Mysterious",
        "tags": ["scifi", "future", "urban", "night"]
    },
    {
        "text": "A close-up portrait of an elderly craftsman with weathered hands working on an intricate wooden carving in his workshop",
        "thumbnail_desc": "Elderly craftsman working with wood",
        "creation_type": "Oil Painting",
        "art_style": "Hyperrealism",
        "mood": "Peaceful",
        "tags": ["portrait", "craftsmanship", "elderly", "detail"]
    },
    {
        "text": "A magical forest with glowing mushrooms, fairy lights, and a small cottage with smoke coming from the chimney",
        "thumbnail_desc": "Magical forest with glowing elements",
        "creation_type": "Fantasy Illustration",
        "art_style": "Studio Ghibli",
        "mood": "Whimsical",
        "tags": ["fantasy", "forest", "magic", "cottage"]
    },
    {
        "text": "A cybernetic samurai with glowing blue circuits standing in a rainy Tokyo street at night",
        "thumbnail_desc": "Cybernetic samurai in rainy Tokyo",
        "creation_type": "Digital Art",
        "art_style": "Cyberpunk",
        "mood": "Dark",
        "tags": ["character", "cyberpunk", "samurai", "rain"]
    },
    {
        "text": "A cute cat with dragon wings and tiny horns sleeping on a pile of gold coins",
        "thumbnail_desc": "Cat with dragon features on gold",
        "creation_type": "Fantasy Illustration",
        "art_style": "Comic Book",
        "mood": "Whimsical",
        "tags": ["creature", "fantasy", "cute", "treasure"]
    },
    {
        "text": "An ancient temple covered in vines and moss, partially sunken into a crystal-clear cenote in the jungle",
        "thumbnail_desc": "Ancient temple in jungle cenote",
        "creation_type": "Concept Art",
        "art_style": "Photorealistic",
        "mood": "Mysterious",
        "tags": ["architecture", "ruins", "jungle", "water"]
    },
    {
        "text": "A cozy coffee shop interior with rain falling outside the windows, soft lighting, and a few people reading books",
        "thumbnail_desc": "Cozy rainy day in coffee shop",
        "creation_type": "Digital Art",
        "art_style": "Impressionist",
        "mood": "Peaceful",
        "tags": ["interior", "rainy", "cozy", "urban"]
    }
]

# CSS for enhanced UI styling - Will be included in part 2

# =============== HELPER FUNCTIONS ===============

# Helper function to format dropdown choices with icons
def format_dropdown_choices(options_dict):
    return [f"{info.get('icon', '•')} {key}" for key in options_dict.keys()]

# Helper function to extract the key from formatted choice
def extract_key(formatted_choice):
    # Skip the icon and space at the beginning
    parts = formatted_choice.split(' ', 1)
    if len(parts) > 1:
        return parts[1]
    return formatted_choice

# Function to load example prompt
def load_example(example_index):
    if example_index < 0 or example_index >= len(EXAMPLE_PROMPTS):
        return "", "", "", ""
    
    example = EXAMPLE_PROMPTS[example_index]
    creation = f"{CREATION_TYPES[example['creation_type']]['icon']} {example['creation_type']}"
    art = f"{ART_STYLES[example['art_style']]['icon']} {example['art_style']}"
    mood = f"{MOODS[example['mood']]['icon']} {example['mood']}"
    
    return example["text"], creation, art, mood

# Get model key from formatted display name
def get_model_key(formatted_name):
    # Extract display name without the icon
    if ' ' in formatted_name:
        display_name = formatted_name.split(' ', 1)[1]
        # Find the corresponding key
        for key, info in IMAGE_MODELS.items():
            if info['display_name'] == display_name:
                return key
    return list(IMAGE_MODELS.keys())[0]  # Default to first model if not found