GPT Image MCP Server

"""OpenAI provider implementation."""

import base64
import logging
from typing import Any

import httpx
from openai import AsyncOpenAI

from .base import (
    ImageResponse,
    LLMProvider,
    ModelCapability,
    ProviderConfig,
    ProviderError,
)

logger = logging.getLogger(__name__)


class OpenAIProvider(LLMProvider):
    """OpenAI provider for image generation using gpt-image-1 and DALL-E models."""

    # Supported models and their capabilities
    SUPPORTED_MODELS = {
        "gpt-image-1": ModelCapability(
            model_id="gpt-image-1",
            supported_sizes=["auto", "1024x1024", "1536x1024", "1024x1536"],
            supported_qualities=["auto", "high", "medium", "low"],
            supported_formats=["png", "jpeg", "webp"],
            max_images_per_request=1,
            supports_style=False,
            supports_background=True,
            supports_compression=True,
            custom_parameters={
                "moderation": ["auto", "low"],
                "background": ["auto", "transparent", "opaque"],
            },
        )
    }

    def __init__(self, config: ProviderConfig):
        super().__init__(config)
        self.client = AsyncOpenAI(
            api_key=config.api_key,
            organization=config.organization,
            base_url=config.base_url or "https://api.openai.com/v1",
            timeout=config.timeout,
            max_retries=config.max_retries,
        )

    def get_supported_models(self) -> set[str]:
        """Return set of supported OpenAI model IDs."""
        return set(self.SUPPORTED_MODELS.keys())

    def get_model_capabilities(self, model_id: str) -> ModelCapability | None:
        """Get capabilities for a specific OpenAI model."""
        return self.SUPPORTED_MODELS.get(model_id)

    async def generate_image(
        self,
        model: str,
        prompt: str,
        quality: str = "auto",
        size: str = "auto",
        style: str = "vivid",
        moderation: str = "auto",
        output_format: str = "png",
        compression: int = 100,
        background: str = "auto",
        n: int = 1,
        **kwargs,
    ) -> ImageResponse:
        """Generate image using OpenAI's Images API."""

        # Validate model
        if model not in self.SUPPORTED_MODELS:
            raise ProviderError(
                f"Model '{model}' is not supported by OpenAI provider",
                provider_name=self.name,
                error_code="UNSUPPORTED_MODEL",
            )

        capabilities = self.SUPPORTED_MODELS[model]

        # Build request parameters
        request_params = {
            "model": model,
            "prompt": prompt,
            "n": min(n, capabilities.max_images_per_request),
        }

        # Map quality parameter based on model
        if model == "dall-e-3":
            request_params["quality"] = "hd" if quality == "high" else "standard"
        elif model == "gpt-image-1":
            request_params["quality"] = quality
            request_params["moderation"] = moderation

        # Add size parameter
        if size in capabilities.supported_sizes:
            request_params["size"] = size
        else:
            request_params["size"] = capabilities.supported_sizes[0]
            self._logger.warning(
                f"Size '{size}' not supported, using {request_params['size']}"
            )

        # Add style parameter if supported
        if capabilities.supports_style:
            request_params["style"] = style

        # Add gpt-image-1 specific parameters
        if model == "gpt-image-1":
            request_params.update(
                {
                    "output_format": output_format,
                    "background": background,
                }
            )

            # Add compression for JPEG/WebP
            if output_format in ["jpeg", "webp"] and compression < 100:
                request_params["output_compression"] = compression

        try:
            self._logger.info(f"Generating image with OpenAI model {model}")
            self._logger.debug(f"Request parameters: {request_params}")

            response = await self.client.images.generate(**request_params)

            # Process response - OpenAI returns base64 for gpt-image-1, URLs for DALL-E
            if hasattr(response.data[0], "b64_json") and response.data[0].b64_json:
                # Base64 response (gpt-image-1)
                image_bytes = base64.b64decode(response.data[0].b64_json)
            elif hasattr(response.data[0], "url") and response.data[0].url:
                # URL response (DALL-E models)
                image_bytes = await self._download_image(response.data[0].url)
            else:
                raise ProviderError(
                    "OpenAI response contains neither base64 data nor URL",
                    provider_name=self.name,
                    error_code="INVALID_RESPONSE",
                )

            # Build metadata
            metadata = {
                "model": model,
                "prompt": prompt,
                "size": request_params["size"],
                "quality": request_params.get("quality", "auto"),
                "style": request_params.get("style", "vivid"),
                "output_format": request_params.get("output_format", "png"),
                "provider": self.name,
                "created_at": getattr(response, "created", None),
            }

            # Add usage information if available
            if hasattr(response, "usage") and response.usage:
                metadata["usage"] = {
                    "total_tokens": response.usage.total_tokens,
                    "input_tokens": response.usage.input_tokens,
                    "output_tokens": response.usage.output_tokens,
                }

            return ImageResponse(
                image_data=image_bytes,
                metadata=metadata,
                provider_response=(
                    response.model_dump() if hasattr(response, "model_dump") else None
                ),
            )

        except Exception as e:
            self._logger.error(f"Error generating image with OpenAI: {e}")
            raise ProviderError(
                f"OpenAI image generation failed: {str(e)}",
                provider_name=self.name,
                error_code="GENERATION_FAILED",
            )

    async def edit_image(
        self,
        model: str,
        image_data: str | bytes,
        prompt: str,
        mask_data: str | bytes | None = None,
        quality: str = "auto",
        size: str = "1536x1024",
        output_format: str = "png",
        compression: int = 100,
        background: str = "auto",
        n: int = 1,
        **kwargs,
    ) -> ImageResponse:
        """Edit image using OpenAI's Images API."""

        # Validate model
        if model not in self.SUPPORTED_MODELS:
            raise ProviderError(
                f"Model '{model}' is not supported by OpenAI provider",
                provider_name=self.name,
                error_code="UNSUPPORTED_MODEL",
            )

        capabilities = self.SUPPORTED_MODELS[model]

        # Convert base64 strings to bytes if needed
        if isinstance(image_data, str):
            if image_data.startswith("data:"):
                image_data = image_data.split(",", 1)[1]
            image_bytes = base64.b64decode(image_data)
        else:
            image_bytes = image_data

        mask_bytes = None
        if mask_data:
            if isinstance(mask_data, str):
                if mask_data.startswith("data:"):
                    mask_data = mask_data.split(",", 1)[1]
                mask_bytes = base64.b64decode(mask_data)
            else:
                mask_bytes = mask_data

        # Build request parameters
        request_params = {
            "model": model,
            "image": image_bytes,
            "prompt": prompt,
            "n": min(n, capabilities.max_images_per_request),
        }

        if mask_bytes:
            request_params["mask"] = mask_bytes

        # Add size parameter
        if size in capabilities.supported_sizes:
            request_params["size"] = size
        else:
            request_params["size"] = capabilities.supported_sizes[0]

        # Add gpt-image-1 specific parameters
        if model == "gpt-image-1":
            request_params.update(
                {
                    "quality": quality,
                    "output_format": output_format,
                    "background": background,
                }
            )

            # Add compression for JPEG/WebP
            if output_format in ["jpeg", "webp"] and compression < 100:
                request_params["output_compression"] = compression

        try:
            self._logger.info(f"Editing image with OpenAI model {model}")
            self._logger.debug(f"Request parameters: {list(request_params.keys())}")

            response = await self.client.images.edit(**request_params)

            # Process response (similar to generate_image)
            if hasattr(response.data[0], "b64_json") and response.data[0].b64_json:
                image_bytes = base64.b64decode(response.data[0].b64_json)
            elif hasattr(response.data[0], "url") and response.data[0].url:
                image_bytes = await self._download_image(response.data[0].url)
            else:
                raise ProviderError(
                    "OpenAI response contains neither base64 data nor URL",
                    provider_name=self.name,
                    error_code="INVALID_RESPONSE",
                )

            metadata = {
                "model": model,
                "prompt": prompt,
                "size": request_params["size"],
                "output_format": request_params.get("output_format", "png"),
                "provider": self.name,
                "operation": "edit",
                "created_at": getattr(response, "created", None),
            }

            return ImageResponse(
                image_data=image_bytes,
                metadata=metadata,
                provider_response=(
                    response.model_dump() if hasattr(response, "model_dump") else None
                ),
            )

        except Exception as e:
            self._logger.error(f"Error editing image with OpenAI: {e}")
            raise ProviderError(
                f"OpenAI image editing failed: {str(e)}",
                provider_name=self.name,
                error_code="EDITING_FAILED",
            )

    async def _download_image(self, image_url: str) -> bytes:
        """Download image from URL (for DALL-E models that return URLs)."""
        try:
            async with httpx.AsyncClient() as client:
                response = await client.get(image_url)
                response.raise_for_status()
                return response.content
        except Exception as e:
            raise ProviderError(
                f"Failed to download image from URL: {str(e)}",
                provider_name=self.name,
                error_code="DOWNLOAD_FAILED",
            )

    def estimate_cost(
        self, model: str, prompt: str, image_count: int = 1
    ) -> dict[str, Any]:
        """Estimate cost for OpenAI image generation."""

        # OpenAI pricing (as of 2024)
        pricing = {
            "gpt-image-1": {
                "text_input_per_1m_tokens": 5.0,
                "image_output_per_1m_tokens": 40.0,
                "tokens_per_image": 1750,
            },
            "dall-e-3": {
                "cost_per_image": 0.04,  # $0.04 per image
            },
            "dall-e-2": {
                "cost_per_image": 0.02,  # $0.02 per image
            },
        }

        if model not in pricing:
            return super().estimate_cost(model, prompt, image_count)

        model_pricing = pricing[model]

        if model == "gpt-image-1":
            # Token-based pricing
            text_tokens = len(prompt.split()) * 1.3  # Rough approximation
            text_cost = (text_tokens / 1_000_000) * model_pricing[
                "text_input_per_1m_tokens"
            ]
            image_cost = (
                image_count * model_pricing["tokens_per_image"] / 1_000_000
            ) * model_pricing["image_output_per_1m_tokens"]
            total_cost = text_cost + image_cost

            return {
                "provider": self.name,
                "model": model,
                "estimated_cost_usd": round(total_cost, 4),
                "currency": "USD",
                "breakdown": {
                    "text_input_cost": round(text_cost, 4),
                    "image_output_cost": round(image_cost, 4),
                    "text_tokens": int(text_tokens),
                    "image_tokens": model_pricing["tokens_per_image"] * image_count,
                    "total_images": image_count,
                },
            }
        else:
            # Fixed price per image
            total_cost = model_pricing["cost_per_image"] * image_count

            return {
                "provider": self.name,
                "model": model,
                "estimated_cost_usd": round(total_cost, 4),
                "currency": "USD",
                "breakdown": {
                    "per_image": model_pricing["cost_per_image"],
                    "total_images": image_count,
                    "base_cost": total_cost,
                },
            }