utils/image_utils.py

import os
import logging
import requests
import base64
from io import BytesIO
from PIL import Image, ImageDraw, ImageFont
import openai
from typing import Optional, Dict, Any
import time
import random

logger = logging.getLogger(__name__)

class ImageGenerationError(Exception):
    """Custom exception for image generation failures"""
    pass

def generate_image(prompt: str, content_type: str, content_info: str = "") -> Optional[str]:
    """
    Main function to generate images for any D&D content type
    
    Args:
        prompt: Detailed image generation prompt
        content_type: Type of content (character_portrait, item, etc.)
        content_info: Brief info about content for placeholder
    
    Returns:
        str: Path to generated image or None if all methods fail
    """
    
    # Check for API keys
    openai_key = os.getenv('OPENAI_API_KEY')
    stability_key = os.getenv('STABILITY_API_KEY')
    
    if openai_key:
        try:
            return generate_with_dalle(prompt)
        except Exception as e:
            logger.warning(f"DALL-E failed: {e}")
    
    if stability_key:
        try:
            return generate_with_stability(prompt)
        except Exception as e:
            logger.warning(f"Stability AI failed: {e}")
    
    # Generate placeholder if no APIs work
    return generate_placeholder(content_type, content_info)

def generate_with_dalle(prompt: str, size: str = "1024x1024") -> Optional[str]:
    """Generate image using OpenAI's DALL-E"""
    try:
        logger.info(f"Generating image with DALL-E: {prompt[:50]}...")
        
        openai.api_key = os.getenv('OPENAI_API_KEY')
        
        response = openai.images.generate(
            model="dall-e-3",
            prompt=prompt,
            size=size,
            quality="standard",
            n=1,
        )
        
        image_url = response.data[0].url
        logger.info("DALL-E image generated successfully")
        return image_url
        
    except Exception as e:
        logger.error(f"DALL-E generation failed: {e}")
        raise ImageGenerationError(f"DALL-E failed: {str(e)}")

def generate_with_stability(prompt: str) -> Optional[str]:
    """Generate image using Stability AI"""
    try:
        logger.info(f"Generating image with Stability AI: {prompt[:50]}...")
        
        url = "https://api.stability.ai/v1/generation/stable-diffusion-xl-1024-v1-0/text-to-image"
        
        headers = {
            "Accept": "application/json",
            "Content-Type": "application/json",
            "Authorization": f"Bearer {os.getenv('STABILITY_API_KEY')}",
        }
        
        body = {
            "text_prompts": [{"text": prompt, "weight": 1}],
            "cfg_scale": 7,
            "height": 1024,
            "width": 1024,
            "samples": 1,
            "steps": 30,
        }
        
        response = requests.post(url, headers=headers, json=body)
        
        if response.status_code != 200:
            raise ImageGenerationError(f"Stability API error: {response.status_code}")
        
        data = response.json()
        
        # Convert base64 to image
        image_data = base64.b64decode(data["artifacts"][0]["base64"])
        img = Image.open(BytesIO(image_data))
        
        # Save temporarily and return path
        temp_path = f"temp_{int(time.time())}_{random.randint(1000,9999)}.png"
        img.save(temp_path)
        logger.info("Stability AI image generated successfully")
        return temp_path
        
    except Exception as e:
        logger.error(f"Stability AI generation failed: {e}")
        raise ImageGenerationError(f"Stability AI failed: {str(e)}")

def generate_placeholder(content_type: str, content_info: str) -> str:
    """Generate a themed placeholder image when AI services fail"""
    try:
        # Create different colored placeholders for different content types
        colors = {
            "character_portrait": (70, 130, 180),      # Steel Blue
            "npc_portrait": (139, 69, 19),             # Saddle Brown
            "item": (255, 215, 0),                     # Gold
            "location": (34, 139, 34),                 # Forest Green
            "faction_symbol": (128, 0, 128),           # Purple
            "deity": (255, 255, 255),                  # White
            "scenario": (220, 20, 60)                  # Crimson
        }
        
        color = colors.get(content_type, (128, 128, 128))
        
        # Create image
        img = Image.new('RGB', (512, 512), color=color)
        draw = ImageDraw.Draw(img)
        
        # Try to load a font, fall back to default if not available
        try:
            font = ImageFont.truetype("arial.ttf", 24)
        except:
            font = ImageFont.load_default()
        
        # Add text
        text_lines = [
            f"{content_type.replace('_', ' ').title()}",
            "Placeholder Image",
            content_info[:30] + "..." if len(content_info) > 30 else content_info
        ]
        
        y_offset = 200
        for line in text_lines:
            # Get text bounding box
            bbox = draw.textbbox((0, 0), line, font=font)
            text_width = bbox[2] - bbox[0]
            text_height = bbox[3] - bbox[1]
            
            # Center text
            x = (512 - text_width) // 2
            draw.text((x, y_offset), line, fill=(255, 255, 255), font=font)
            y_offset += text_height + 10
        
        # Save placeholder
        temp_path = f"placeholder_{content_type}_{int(time.time())}.png"
        img.save(temp_path)
        
        logger.info(f"Generated {content_type} placeholder image")
        return temp_path
        
    except Exception as e:
        logger.error(f"Failed to create placeholder: {e}")
        return None

def cleanup_temp_files():
    """Clean up temporary image files"""
    import glob
    temp_files = glob.glob("temp_*.png") + glob.glob("placeholder_*.png")
    for file in temp_files:
        try:
            os.remove(file)
            logger.info(f"Cleaned up temporary file: {file}")
        except Exception as e:
            logger.warning(f"Failed to clean up {file}: {e}")