src/image_processor.py

"""Image processing utilities."""

import json
import zipfile
from pathlib import Path
from typing import Optional, List, Dict, Any
from datetime import datetime
from PIL import Image
import logging

from config.settings import OUTPUT_DIR

logger = logging.getLogger(__name__)


def save_image(
    image: Image.Image,
    prompt: str,
    params: Dict[str, Any],
    metadata: Optional[Dict[str, Any]] = None
) -> Path:
    """Save an image with metadata.
    
    Args:
        image: PIL Image object to save.
        prompt: Prompt used to generate the image.
        params: Generation parameters.
        metadata: Additional metadata to include.
        
    Returns:
        Path to the saved image file.
    """
    timestamp = datetime.now().timestamp()
    filename = f"texture_{int(timestamp)}.png"
    filepath = OUTPUT_DIR / filename
    
    # Save image
    image.save(filepath, "PNG")
    logger.info(f"Saved image to {filepath}")
    
    # Save metadata
    metadata_path = filepath.with_suffix(".json")
    metadata_dict = {
        "timestamp": timestamp,
        "prompt": prompt,
        "params": params,
        "image_path": str(filepath),
        "filename": filename,
    }
    
    if metadata:
        metadata_dict.update(metadata)
    
    with open(metadata_path, "w", encoding="utf-8") as f:
        json.dump(metadata_dict, f, indent=2)
    
    return filepath


def create_thumbnail(image: Image.Image, size: tuple[int, int] = (256, 256)) -> Image.Image:
    """Create a thumbnail from an image.
    
    Args:
        image: PIL Image object.
        size: Thumbnail size (width, height).
        
    Returns:
        Thumbnail Image object.
    """
    thumb = image.copy()
    thumb.thumbnail(size, Image.Resampling.LANCZOS)
    return thumb


def validate_image_format(filepath: Path) -> bool:
    """Validate if a file is a valid image format.
    
    Args:
        filepath: Path to the image file.
        
    Returns:
        True if valid, False otherwise.
    """
    valid_formats = {".png", ".jpg", ".jpeg", ".webp"}
    return filepath.suffix.lower() in valid_formats


def create_zip(files: List[Path], output_path: Path) -> Path:
    """Create a ZIP archive from multiple files.
    
    Args:
        files: List of file paths to include.
        output_path: Path for the output ZIP file.
        
    Returns:
        Path to the created ZIP file.
    """
    with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zipf:
        for file in files:
            if file.exists():
                zipf.write(file, file.name)
                logger.info(f"Added {file.name} to ZIP")
    
    logger.info(f"Created ZIP archive at {output_path}")
    return output_path


def optimize_image_for_web(image: Image.Image, max_size: int = 2048) -> Image.Image:
    """Optimize an image for web display.
    
    Args:
        image: PIL Image object.
        max_size: Maximum dimension size.
        
    Returns:
        Optimized Image object.
    """
    width, height = image.size
    
    if width <= max_size and height <= max_size:
        return image
    
    # Calculate new dimensions maintaining aspect ratio
    if width > height:
        new_width = max_size
        new_height = int(height * (max_size / width))
    else:
        new_height = max_size
        new_width = int(width * (max_size / height))
    
    return image.resize((new_width, new_height), Image.Resampling.LANCZOS)


def load_metadata(image_path: Path) -> Optional[Dict[str, Any]]:
    """Load metadata for an image.
    
    Args:
        image_path: Path to the image file.
        
    Returns:
        Metadata dictionary or None if not found.
    """
    metadata_path = image_path.with_suffix(".json")
    
    if not metadata_path.exists():
        return None
    
    try:
        with open(metadata_path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception as e:
        logger.error(f"Error loading metadata: {e}")
        return None