character_forge_image/services/generation_service.py

"""
Generation Service
==================

High-level service for orchestrating image generation workflows.
Coordinates backend routing, file saving, metadata management, and history tracking.
"""

from pathlib import Path
from typing import Optional, Dict, Any
from PIL import Image

from core import BackendRouter
from models.generation_request import GenerationRequest
from models.generation_result import GenerationResult
from utils.file_utils import (
    save_image,
    create_generation_metadata,
    ensure_directory_exists
)
from utils.logging_utils import get_logger
from config.settings import Settings


logger = get_logger(__name__)


class GenerationService:
    """
    High-level service for image generation operations.

    Orchestrates the complete generation workflow:
    1. Validate request
    2. Route to backend
    3. Save results
    4. Update history
    5. Return result
    """

    def __init__(self, api_key: Optional[str] = None):
        """
        Initialize generation service.

        Args:
            api_key: Optional Gemini API key (defaults to Settings)
        """
        self.router = BackendRouter(api_key=api_key)
        logger.info("GenerationService initialized")

    def generate_and_save(
        self,
        request: GenerationRequest,
        output_dir: Path,
        base_filename: str,
        save_metadata: bool = True
    ) -> GenerationResult:
        """
        Generate image and save to disk.

        Complete workflow:
        1. Generate image via backend
        2. Save image to output directory
        3. Save metadata JSON (optional)
        4. Update result with saved path

        Args:
            request: GenerationRequest object
            output_dir: Directory to save image
            base_filename: Base name for output file
            save_metadata: Whether to save metadata JSON

        Returns:
            GenerationResult with saved_path populated
        """
        try:
            logger.info(f"Starting generation: {request.prompt[:50]}...")

            # Ensure output directory exists
            ensure_directory_exists(output_dir)

            # Generate image
            result = self.router.generate(request)

            if not result.success:
                logger.warning(f"Generation failed: {result.message}")
                return result

            # Save image and metadata
            metadata = None
            if save_metadata:
                metadata = create_generation_metadata(
                    prompt=request.prompt,
                    backend=request.backend,
                    aspect_ratio=request.aspect_ratio,
                    temperature=request.temperature,
                    generation_time=result.generation_time,
                    **request.metadata
                )

            image_path, metadata_path = save_image(
                image=result.image,
                directory=output_dir,
                base_name=base_filename,
                metadata=metadata
            )

            # Update result with saved paths
            result.saved_path = image_path
            if metadata_path:
                result.metadata['metadata_path'] = metadata_path

            logger.info(f"Image saved: {image_path}")
            return result

        except Exception as e:
            logger.error(f"Generation and save failed: {e}", exc_info=True)
            return GenerationResult.error_result(
                message=f"Generation service error: {str(e)}"
            )

    def generate_only(self, request: GenerationRequest) -> GenerationResult:
        """
        Generate image without saving to disk.

        Useful for previews or temporary generations.

        Args:
            request: GenerationRequest object

        Returns:
            GenerationResult
        """
        logger.info(f"Generating (no save): {request.prompt[:50]}...")
        return self.router.generate(request)

    def batch_generate(
        self,
        prompts: list[str],
        backend: str,
        aspect_ratio: str,
        temperature: float,
        output_dir: Path,
        base_filename_template: str = "batch_{index}"
    ) -> list[GenerationResult]:
        """
        Generate multiple images from prompt list.

        Args:
            prompts: List of prompts to generate
            backend: Backend to use
            aspect_ratio: Aspect ratio for all images
            temperature: Temperature for all images
            output_dir: Output directory
            base_filename_template: Template for filenames (use {index} placeholder)

        Returns:
            List of GenerationResult objects
        """
        logger.info(f"Starting batch generation: {len(prompts)} prompts")
        results = []

        for i, prompt in enumerate(prompts):
            request = GenerationRequest(
                prompt=prompt,
                backend=backend,
                aspect_ratio=aspect_ratio,
                temperature=temperature
            )

            base_filename = base_filename_template.format(index=i+1)
            result = self.generate_and_save(
                request=request,
                output_dir=output_dir,
                base_filename=base_filename
            )
            results.append(result)

            logger.info(f"Batch {i+1}/{len(prompts)}: {'Success' if result.success else 'Failed'}")

        logger.info(f"Batch complete: {sum(r.success for r in results)}/{len(results)} successful")
        return results

    def check_backend_availability(self, backend: str) -> tuple[bool, str]:
        """
        Check if backend is available.

        Args:
            backend: Backend name to check

        Returns:
            Tuple of (is_available, status_message)
        """
        return self.router.check_backend_health(backend)

    def get_all_backend_status(self) -> Dict[str, Dict[str, Any]]:
        """
        Get status of all configured backends.

        Returns:
            Dictionary mapping backend name to status info
        """
        return self.router.get_all_backend_status()

    def validate_and_generate(
        self,
        prompt: str,
        backend: str,
        aspect_ratio: str,
        temperature: float,
        input_images: Optional[list[Image.Image]] = None,
        output_dir: Optional[Path] = None,
        base_filename: Optional[str] = None
    ) -> GenerationResult:
        """
        Convenience method with built-in validation.

        Args:
            prompt: Generation prompt
            backend: Backend name
            aspect_ratio: Aspect ratio
            temperature: Temperature
            input_images: Optional input images
            output_dir: Optional output directory (if None, no save)
            base_filename: Optional base filename (if None, use sanitized prompt)

        Returns:
            GenerationResult
        """
        from utils.validation import validate_generation_request
        from utils.file_utils import sanitize_filename

        # Validate request
        is_valid, error = validate_generation_request(
            prompt=prompt,
            backend=backend,
            aspect_ratio=aspect_ratio,
            temperature=temperature,
            input_images=input_images
        )

        if not is_valid:
            logger.error(f"Validation failed: {error}")
            return GenerationResult.error_result(message=f"Validation error: {error}")

        # Build request
        request = GenerationRequest(
            prompt=prompt,
            backend=backend,
            aspect_ratio=aspect_ratio,
            temperature=temperature,
            input_images=input_images or []
        )

        # Generate with or without save
        if output_dir is None:
            return self.generate_only(request)
        else:
            if base_filename is None:
                base_filename = sanitize_filename(prompt[:50])
            return self.generate_and_save(
                request=request,
                output_dir=output_dir,
                base_filename=base_filename
            )