"""Types for Nano Banana (Gemini) image generation."""
from __future__ import annotations

from enum import Enum
from pydantic import BaseModel, Field


# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------

class ImageModel(str, Enum):
    """Available image generation models."""
    STANDARD = "standard"  # Gemini 2.5 Flash — fast generation
    PRO = "pro"            # Gemini 3.0 Pro — high quality
    V2 = "v2"              # Gemini 3.1 Flash — next-gen (Nano Banana 2)


class ImageQuality(str, Enum):
    """Image quality options (pro and v2 models only)."""
    STANDARD = "standard"
    HD_2K = "2k"
    HD_4K = "4k"


class ImageAspectRatio(str, Enum):
    """Supported aspect ratios."""
    AUTO = "auto"
    # Landscape
    ULTRA_WIDE = "21:9"
    WIDE = "16:9"
    STANDARD_L = "4:3"
    PHOTO_L = "3:2"
    # Square
    SQUARE = "1:1"
    # Portrait
    PORTRAIT = "9:16"
    PORTRAIT_34 = "3:4"
    PORTRAIT_23 = "2:3"
    # Flexible
    FLEX_54 = "5:4"
    FLEX_45 = "4:5"


class ImageStatus(str, Enum):
    """Generation task status."""
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"


class ImageErrorCode(str, Enum):
    """Known error codes from the API."""
    INSUFFICIENT_CREDITS = "INSUFFICIENT_CREDITS"
    MODERATION_FAILED = "MODERATION_FAILED"
    API_ERROR = "API_ERROR"
    TASK_FAILED = "TASK_FAILED"
    TIMEOUT = "TIMEOUT"
    INVALID_IMAGE = "INVALID_IMAGE"
    IMAGE_TOO_LARGE = "IMAGE_TOO_LARGE"
    UNSUPPORTED_IMAGE_FORMAT = "UNSUPPORTED_IMAGE_FORMAT"


# ---------------------------------------------------------------------------
# Request models
# ---------------------------------------------------------------------------

class ImageGenerateParams(BaseModel):
    """Parameters for POST /v1/gemini/generate."""

    prompt: str = Field(
        ...,
        min_length=1,
        max_length=4000,
        description="Text description of the image to generate",
    )
    model: ImageModel = Field(
        default=ImageModel.STANDARD,
        description="Model version: standard, pro, or v2",
    )
    quality: ImageQuality = Field(
        default=ImageQuality.STANDARD,
        description="Image quality (2k/4k only for pro and v2)",
    )
    aspect_ratio: ImageAspectRatio = Field(
        default=ImageAspectRatio.AUTO,
        description="Image aspect ratio",
    )
    temperature: float = Field(
        default=0.7,
        ge=0.0,
        le=1.0,
        description="Creativity level: 0 (deterministic) to 1 (creative)",
    )
    images: list[str] | None = Field(
        default=None,
        description="Input images (base64 or URLs) for editing. Max 4 (standard) or 8 (pro/v2).",
    )
    callback_url: str | None = Field(
        default=None,
        description="Webhook URL for completion notifications",
    )

    def to_api_body(self) -> dict:
        """Convert to API request body, excluding None fields."""
        data: dict = {
            "prompt": self.prompt,
            "model": self.model.value,
            "quality": self.quality.value,
            "aspect_ratio": self.aspect_ratio.value,
            "temperature": self.temperature,
        }
        if self.images:
            data["images"] = self.images
        if self.callback_url:
            data["callback_url"] = self.callback_url
        return data


# ---------------------------------------------------------------------------
# Response models
# ---------------------------------------------------------------------------

class ImageGenerateResponse(BaseModel):
    """Response from POST /v1/gemini/generate."""

    task_uuid: str = Field(description="UUID for tracking the generation")
    status: ImageStatus = Field(description="Initial status (always 'pending')")
    price: str = Field(description="Amount charged in USD")


class ImageStatusResponse(BaseModel):
    """Response from GET /v1/gemini/generations/:task_uuid."""

    task_uuid: str
    status: ImageStatus
    prompt: str
    result_image_url: str | None = None
    message: str | None = Field(default=None, description="Error code if failed")

    @property
    def is_terminal(self) -> bool:
        """Whether the generation has reached a final state."""
        return self.status in (ImageStatus.COMPLETED, ImageStatus.FAILED)

    @property
    def is_completed(self) -> bool:
        return self.status == ImageStatus.COMPLETED

    @property
    def is_failed(self) -> bool:
        return self.status == ImageStatus.FAILED

    @property
    def error_code(self) -> ImageErrorCode | None:
        """Parsed error code, or None if not failed."""
        if self.message is None:
            return None
        try:
            return ImageErrorCode(self.message)
        except ValueError:
            return None
