"""OpenAI image generation provider.

Supports gpt-image-1 (current) and dall-e-3 (legacy, deprecated May 2026).
Ported from questfoundry — prompt distillation removed.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, NoReturn

from image_generation_mcp.providers.types import (
    ImageContentPolicyError,
    ImageProviderConnectionError,
    ImageProviderError,
    ImageResult,
)

if TYPE_CHECKING:
    from openai import AsyncOpenAI

logger = logging.getLogger(__name__)

_GPT_IMAGE_SIZES: dict[str, str] = {
    "1:1": "1024x1024",
    "16:9": "1536x1024",
    "9:16": "1024x1536",
    "3:2": "1536x1024",
    "2:3": "1024x1536",
}

_DALLE3_SIZES: dict[str, str] = {
    "1:1": "1024x1024",
    "16:9": "1792x1024",
    "9:16": "1024x1792",
    "3:2": "1792x1024",
    "2:3": "1024x1792",
}

_FORMAT_TO_CONTENT_TYPE: dict[str, str] = {
    "png": "image/png",
    "jpeg": "image/jpeg",
    "webp": "image/webp",
}

_DALLE3_FORMATS: frozenset[str] = frozenset({"png"})


def _is_gpt_image_model(model: str) -> bool:
    """Return True for gpt-image-* models (not dall-e)."""
    return model.startswith("gpt-image")


class OpenAIImageProvider:
    """Image generation via OpenAI's Images API.

    Args:
        model: Model name (``gpt-image-1`` or ``dall-e-3``).
        api_key: OpenAI API key.
        output_format: Image format (``png``, ``jpeg``, ``webp``).
    """

    def __init__(
        self,
        api_key: str,
        model: str = "gpt-image-1",
        output_format: str = "png",
    ) -> None:
        self._model = model
        self._output_format = output_format

        if output_format not in _FORMAT_TO_CONTENT_TYPE:
            supported = ", ".join(sorted(_FORMAT_TO_CONTENT_TYPE))
            raise ImageProviderError(
                "openai",
                f"Unsupported output_format '{output_format}'. Supported: {supported}",
            )

        self._is_gpt_image = _is_gpt_image_model(model)

        if not self._is_gpt_image and output_format not in _DALLE3_FORMATS:
            supported_dalle3 = ", ".join(sorted(_DALLE3_FORMATS))
            raise ImageProviderError(
                "openai",
                f"dall-e-3 does not support output_format '{output_format}'. "
                f"Supported: {supported_dalle3}",
            )

        self._content_type = _FORMAT_TO_CONTENT_TYPE[output_format]
        self._sizes = _GPT_IMAGE_SIZES if self._is_gpt_image else _DALLE3_SIZES
        self._client: AsyncOpenAI = self._create_client(api_key)

    def _create_client(self, api_key: str) -> AsyncOpenAI:
        """Create the AsyncOpenAI client (deferred import)."""
        try:
            from openai import AsyncOpenAI as _AsyncOpenAI
        except ImportError as e:
            raise ImageProviderError(
                "openai", "openai package not installed. Run: uv add openai"
            ) from e
        return _AsyncOpenAI(api_key=api_key)

    async def generate(
        self,
        prompt: str,
        *,
        negative_prompt: str | None = None,
        aspect_ratio: str = "1:1",
        quality: str = "standard",
    ) -> ImageResult:
        """Generate an image via OpenAI Images API.

        Args:
            prompt: Positive text prompt.
            negative_prompt: Appended as ``"Avoid: ..."`` clause.
            aspect_ratio: Maps to OpenAI size parameter.
            quality: Quality level.

        Returns:
            ImageResult with generated image.

        Raises:
            ImageProviderError: On API errors.
            ImageContentPolicyError: On content policy rejection.
            ImageProviderConnectionError: On network errors.
        """
        effective_prompt = prompt
        if negative_prompt:
            effective_prompt = f"{prompt}\n\nAvoid: {negative_prompt}"

        api_quality = quality
        if self._is_gpt_image:
            api_quality = {"standard": "high", "hd": "high"}.get(quality, quality)

        size = self._sizes.get(aspect_ratio)
        if size is None:
            supported = ", ".join(sorted(self._sizes))
            raise ImageProviderError(
                "openai",
                f"Unsupported aspect_ratio '{aspect_ratio}'. Supported: {supported}",
            )

        logger.debug(
            "OpenAI image generation: model=%s size=%s quality=%s",
            self._model,
            size,
            api_quality,
        )

        try:
            api_kwargs: dict[str, Any] = {
                "model": self._model,
                "prompt": effective_prompt,
                "n": 1,
                "size": size,
                "quality": api_quality,
            }
            if self._is_gpt_image:
                api_kwargs["output_format"] = self._output_format
            else:
                api_kwargs["response_format"] = "b64_json"
            response = await self._client.images.generate(**api_kwargs)
        except ImageProviderError:
            raise
        except Exception as e:
            self._handle_error(e)

        if not response.data:
            raise ImageProviderError("openai", "Empty response from image API")

        image_item = response.data[0]
        b64_data = image_item.b64_json
        if not b64_data:
            raise ImageProviderError("openai", "No image data in response")

        metadata: dict[str, Any] = {
            "model": self._model,
            "size": size,
            "quality": quality,
            "api_quality": api_quality,
        }
        revised_prompt = getattr(image_item, "revised_prompt", None)
        if revised_prompt:
            metadata["revised_prompt"] = revised_prompt

        logger.info("OpenAI image generated: model=%s size=%s", self._model, size)

        return ImageResult.from_base64(
            b64_data,
            content_type=self._content_type,
            **metadata,
        )

    def _handle_error(self, error: Exception) -> NoReturn:
        """Convert OpenAI exceptions to ImageProvider exceptions."""
        from openai import APIConnectionError, APIStatusError

        if isinstance(error, APIConnectionError):
            raise ImageProviderConnectionError(
                "openai", f"Connection error: {error}"
            ) from error

        if isinstance(error, APIStatusError):
            status = error.status_code
            body = getattr(error, "body", None) or {}
            code = (
                body.get("error", {}).get("code", "") if isinstance(body, dict) else ""
            )
            if status == 400 and (
                code == "content_policy_violation"
                or "content_policy" in str(error).lower()
            ):
                raise ImageContentPolicyError(
                    "openai", f"Content policy rejection: {error}"
                ) from error
            raise ImageProviderError(
                "openai", f"API error (HTTP {status}): {error}"
            ) from error

        raise ImageProviderError(
            "openai", f"Image generation failed: {error}"
        ) from error
