"""
Image generation tool supporting both Gemini and Imagen models.
"""

import base64
import json
import logging
from pathlib import Path
from typing import Any

from fastmcp import Image

from ..config import ALL_MODELS, ASPECT_RATIOS, get_settings
from ..core import (
    validate_aspect_ratio,
    validate_image_format,
    validate_model,
    validate_number_of_images,
    validate_person_generation,
    validate_prompt,
    validate_seed,
)
from ..services import ImageService

logger = logging.getLogger(__name__)


async def generate_image_tool(
    prompt: str,
    model: str | None = None,
    enhance_prompt: bool = True,
    number_of_images: int = 1,
    aspect_ratio: str = "1:1",
    output_format: str = "png",
    # Gemini-specific options
    input_image_path: str | None = None,
    maintain_character_consistency: bool = False,
    blend_images: bool = False,
    use_world_knowledge: bool = False,
    # Imagen-specific options
    person_generation: str = "allow_adult",
    negative_prompt: str | None = None,
    seed: int | None = None,
    # Output options
    save_to_disk: bool = True,
    **kwargs: Any,
) -> dict[str, Any]:
    """
    Generate images using Gemini or Imagen models.

    Args:
        prompt: Text description for image generation
        model: Model to use (gemini-2.5-flash-image, imagen-3, imagen-4, imagen-4-ultra)
        enhance_prompt: Automatically enhance prompt for better results
        number_of_images: Number of images to generate (1-4)
        aspect_ratio: Image aspect ratio (1:1, 16:9, 9:16, etc.)
        output_format: Image format (png, jpeg, webp)
        input_image_path: Path to input image for editing (Gemini only)
        maintain_character_consistency: Maintain character features across generations (Gemini)
        blend_images: Enable multi-image blending (Gemini)
        use_world_knowledge: Use real-world knowledge for context (Gemini)
        person_generation: Person generation policy (Imagen: dont_allow, allow_adult, allow_all)
        negative_prompt: What to avoid in the image (Imagen)
        seed: Random seed for reproducibility (Imagen)
        save_to_disk: Save images to output directory

    Returns:
        Dict with generated images and metadata
    """
    # Validate inputs
    validate_prompt(prompt)
    if model:
        validate_model(model)
    validate_number_of_images(number_of_images)
    validate_aspect_ratio(aspect_ratio)
    validate_image_format(output_format)

    if person_generation:
        validate_person_generation(person_generation)
    if seed is not None:
        validate_seed(seed)

    # Get settings
    settings = get_settings()

    # Determine model
    if model is None:
        model = settings.api.default_gemini_model

    # Initialize image service
    image_service = ImageService(
        api_key=settings.api.gemini_api_key,
        enable_enhancement=settings.api.enable_prompt_enhancement,
        timeout=settings.api.request_timeout,
    )

    try:
        # Prepare parameters based on model type
        params: dict[str, Any] = {
            "aspect_ratio": aspect_ratio,
        }

        # Add input image if provided (Gemini)
        if input_image_path:
            image_path = Path(input_image_path)
            if image_path.exists():
                image_data = base64.b64encode(image_path.read_bytes()).decode()
                params["input_image"] = image_data
            else:
                logger.warning(f"Input image not found: {input_image_path}")

        # Add Gemini-specific options
        if maintain_character_consistency:
            params["maintainCharacterConsistency"] = True
        if blend_images:
            params["blendImages"] = True
        if use_world_knowledge:
            params["useWorldKnowledge"] = True

        # Add Imagen-specific options
        if model.startswith("imagen"):
            params["number_of_images"] = number_of_images
            params["output_format"] = f"image/{output_format}"
            params["person_generation"] = person_generation
            if negative_prompt:
                params["negative_prompt"] = negative_prompt
            if seed is not None:
                params["seed"] = seed

        # Generate images
        results = await image_service.generate(
            prompt=prompt,
            model=model,
            enhance_prompt=enhance_prompt and settings.api.enable_prompt_enhancement,
            **params
        )

        # Prepare response
        response = {
            "success": True,
            "model": model,
            "prompt": prompt,
            "images_generated": len(results),
            "images": [],
            "metadata": {
                "enhance_prompt": enhance_prompt,
                "aspect_ratio": aspect_ratio,
            }
        }

        # Save images and prepare for MCP response
        for result in results:
            image_info = {
                "index": result.index,
                "size": result.get_size(),
                "timestamp": result.timestamp.isoformat(),
            }

            if save_to_disk:
                # Save to output directory
                file_path = result.save(settings.output_dir)
                image_info["path"] = str(file_path)
                image_info["filename"] = file_path.name

            # Add enhanced prompt info
            if "enhanced_prompt" in result.metadata:
                image_info["enhanced_prompt"] = result.metadata["enhanced_prompt"]

            response["images"].append(image_info)

        return response

    finally:
        await image_service.close()


def register_generate_image_tool(mcp_server: Any) -> None:
    """Register generate_image tool with MCP server."""

    @mcp_server.tool()
    async def generate_image(
        prompt: str,
        model: str | None = None,
        enhance_prompt: bool = True,
        number_of_images: int = 1,
        aspect_ratio: str = "1:1",
        output_format: str = "png",
        input_image_path: str | None = None,
        maintain_character_consistency: bool = False,
        blend_images: bool = False,
        use_world_knowledge: bool = False,
        person_generation: str = "allow_adult",
        negative_prompt: str | None = None,
        seed: int | None = None,
    ) -> str:
        """
        Generate images using Google's Gemini or Imagen models.

        Supports both:
        - Gemini 2.5 Flash Image: Advanced image generation with editing, prompt enhancement
        - Imagen 4/Ultra: High-quality image generation with advanced controls

        Args:
            prompt: Text description of the image to generate
            model: Model to use (default: gemini-2.5-flash-image)
            enhance_prompt: Automatically enhance prompt using AI (default: True)
            number_of_images: Number of images to generate, 1-4 (default: 1)
            aspect_ratio: Image aspect ratio like 1:1, 16:9, 9:16 (default: 1:1)
            output_format: Image format: png, jpeg, webp (default: png)
            input_image_path: Path to input image for editing (Gemini only)
            maintain_character_consistency: Maintain character features (Gemini only)
            blend_images: Enable multi-image blending (Gemini only)
            use_world_knowledge: Use real-world knowledge (Gemini only)
            person_generation: Person policy: dont_allow, allow_adult, allow_all (Imagen only)
            negative_prompt: What to avoid in the image (Imagen only)
            seed: Random seed for reproducibility (Imagen only)

        Available models:
        - gemini-2.5-flash-image (default)
        - imagen-4
        - imagen-4-fast
        - imagen-4-ultra

        Returns:
            JSON string with generation results and file paths
        """
        try:
            result = await generate_image_tool(
                prompt=prompt,
                model=model,
                enhance_prompt=enhance_prompt,
                number_of_images=number_of_images,
                aspect_ratio=aspect_ratio,
                output_format=output_format,
                input_image_path=input_image_path,
                maintain_character_consistency=maintain_character_consistency,
                blend_images=blend_images,
                use_world_knowledge=use_world_knowledge,
                person_generation=person_generation,
                negative_prompt=negative_prompt,
                seed=seed,
            )

            return json.dumps(result, indent=2)

        except Exception as e:
            logger.error(f"Error generating image: {e}")
            return json.dumps({
                "success": False,
                "error": str(e),
                "error_type": type(e).__name__
            }, indent=2)
