from __future__ import annotations

import base64
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Literal, cast

from typing_extensions import assert_never

from pydantic_ai.exceptions import (
    ContentFilterError,
    UnexpectedModelBehavior,
    UserError,
)
from pydantic_ai.messages import BinaryImage, ImageUrl, UploadedFile
from pydantic_ai.models import check_allow_model_requests, download_item
from pydantic_ai.providers import Provider, infer_provider
from pydantic_ai.usage import RequestUsage

from ._media_type import output_format_from_media_type
from ._validation import validate_image_count, warn_image_generation_settings
from .base import ImageGenerationInput, ImageGenerationModel
from .result import GeneratedImage, ImageGenerationResult
from .settings import ImageGenerationSettings

try:
    import grpc
    from xai_sdk import AsyncClient
    from xai_sdk.aio.image import ImageResponse
    from xai_sdk.proto import usage_pb2
    from xai_sdk.types import (
        ImageAspectRatio as XaiImageAspectRatio,
        ImageGenerationModel as LatestXaiImageGenerationModelNames,
        ImageResolution,
    )

    from pydantic_ai.models.xai import (
        _GRPC_STATUS_TO_HTTP as _CHAT_GRPC_STATUS_TO_HTTP,  # pyright: ignore[reportPrivateUsage]
        _map_api_errors,  # pyright: ignore[reportPrivateUsage]
    )

    from ._xai_geometry import resolve_xai_geometry
except ImportError as _import_error:
    raise ImportError(
        'Please install `xai-sdk` to use the xAI image generation model, '
        'you can use the `xai` optional group — `pip install "pydantic-ai-slim[xai]"`'
    ) from _import_error


XaiImageGenerationModelName = str | LatestXaiImageGenerationModelNames
"""Possible xAI image generation model names."""


class XaiImageGenerationSettings(ImageGenerationSettings, total=False):
    """Settings used for an xAI image generation request.

    All fields from [`ImageGenerationSettings`][pydantic_ai.images.ImageGenerationSettings]
    are supported, plus xAI-specific settings prefixed with `xai_`.
    """

    # ALL FIELDS MUST BE `xai_` PREFIXED SO YOU CAN MERGE THEM WITH OTHER MODELS.

    xai_n: int
    """The number of images to generate."""

    xai_user: str
    """A unique identifier representing your end-user."""

    xai_aspect_ratio: XaiImageAspectRatio
    """The aspect ratio of the generated image."""

    xai_resolution: ImageResolution
    """The resolution tier of the generated image."""


@dataclass
class _XaiInputImages:
    """The reference-image arguments `image.sample` and `image.sample_batch` take.

    Named rather than a tuple because the singular and plural forms share their types pairwise, so a
    transposed pair would typecheck and send reference URLs as file IDs.
    """

    image_url: str | None = None
    image_file_id: str | None = None
    image_urls: list[str] | None = None
    image_file_ids: list[str] | None = None


# `eq=False`: a model bound to a live provider client compares by identity. `_settings` lives on the
# non-dataclass base, so it is outside generated field equality, which would call two differently
# configured models equal, merge them to the first, and leave the class unhashable.
@dataclass(init=False, eq=False)
class XaiImageGenerationModel(ImageGenerationModel):
    """xAI image generation model implementation.

    This model works with the Grok Imagine models, such as `grok-imagine-image` and
    `grok-imagine-image-quality`, through the official xAI SDK, which connects over gRPC.

    xAI moderates silently: a flagged image in a batch comes back empty rather than as an error, so the
    clean images are returned and the flagged positions are reported through
    `provider_details['moderated_image_indices']`. A
    [`ContentFilterError`][pydantic_ai.exceptions.ContentFilterError] is raised only when every image
    was flagged. See the [xAI model page](../models/xai.md#image-generation) for details.

    Example:
    ```python
    from pydantic_ai.images.xai import XaiImageGenerationModel
    from pydantic_ai.providers.xai import XaiProvider

    # Using xAI directly (requires XAI_API_KEY env var)
    model = XaiImageGenerationModel('grok-imagine-image')

    # Or with explicit provider configuration
    model = XaiImageGenerationModel(
        'grok-imagine-image',
        provider=XaiProvider(api_key='your-api-key'),
    )
    ```
    """

    _model_name: XaiImageGenerationModelName = field(repr=False)
    _provider: Provider[AsyncClient] = field(repr=False)

    def __init__(
        self,
        model_name: XaiImageGenerationModelName,
        *,
        provider: Literal['xai'] | Provider[AsyncClient] = 'xai',
        settings: ImageGenerationSettings | None = None,
    ):
        """Initialize an xAI image generation model.

        Args:
            model_name: The name of the Grok Imagine model to use.
                See [xAI's image generation documentation](https://docs.x.ai/developers/model-capabilities/images/generation)
                for available models.
            provider: The provider to use for authentication and API access. Can be:

                - `'xai'` (default): Uses the standard xAI API
                - An [`XaiProvider`][pydantic_ai.providers.xai.XaiProvider] instance for custom
                  configuration, such as a custom `api_host` or `xai_client`
            settings: Model-specific
                [`ImageGenerationSettings`][pydantic_ai.images.ImageGenerationSettings]
                to use as defaults for this model.
        """
        self._model_name = model_name

        if isinstance(provider, str):
            provider = infer_provider(provider)
        self._provider = provider

        super().__init__(settings=settings)

    @property
    def _client(self) -> AsyncClient:
        return self._provider.client

    @property
    def base_url(self) -> str:
        return self._provider.base_url

    @property
    def model_name(self) -> XaiImageGenerationModelName:
        """The image generation model name."""
        return self._model_name

    @property
    def system(self) -> str:
        """The image generation model provider."""
        return self._provider.name

    async def generate(
        self,
        prompt: str,
        *,
        images: Sequence[ImageGenerationInput] | None = None,
        settings: ImageGenerationSettings | None = None,
    ) -> ImageGenerationResult:
        check_allow_model_requests()
        prompt, images, settings = self.prepare_generate(prompt, images=images, settings=settings)
        xai_settings = cast(XaiImageGenerationSettings, settings)
        resolved = _resolve_xai_settings(xai_settings, model_name=self.model_name)
        warn_image_generation_settings(self.system, ignored=resolved.ignored, conflicts=resolved.conflicts)
        input_images = await self._map_input_images(images)
        n = xai_settings.get('xai_n') or 1

        with _map_api_errors(self.model_name, status_map=_GRPC_STATUS_TO_HTTP):
            if n == 1:
                response = await self._client.image.sample(
                    prompt,
                    self.model_name,
                    image_url=input_images.image_url,
                    image_file_id=input_images.image_file_id,
                    image_urls=input_images.image_urls,
                    image_file_ids=input_images.image_file_ids,
                    user=xai_settings.get('xai_user'),
                    image_format='base64',
                    aspect_ratio=resolved.aspect_ratio,
                    resolution=resolved.resolution,
                )
                responses = [response]
            else:
                responses = list(
                    await self._client.image.sample_batch(
                        prompt,
                        self.model_name,
                        n,
                        image_url=input_images.image_url,
                        image_file_id=input_images.image_file_id,
                        image_urls=input_images.image_urls,
                        image_file_ids=input_images.image_file_ids,
                        user=xai_settings.get('xai_user'),
                        image_format='base64',
                        aspect_ratio=resolved.aspect_ratio,
                        resolution=resolved.resolution,
                    )
                )

        return self._map_response(prompt, responses)

    async def _map_input_images(self, images: Sequence[ImageGenerationInput]) -> _XaiInputImages:
        image_references: list[str] = []
        file_ids: list[str] = []
        seen_reference = False
        order_violated = False

        for image in images:
            if isinstance(image, UploadedFile):
                self._validate_uploaded_file_provider(image)
                if seen_reference:
                    order_violated = True
                file_ids.append(image.file_id)
            elif isinstance(image, BinaryImage):
                image_references.append(image.data_uri)
                seen_reference = True
            elif isinstance(image, ImageUrl):
                if image.force_download:
                    downloaded_image = await download_item(image, data_format='base64_uri')
                    image_references.append(downloaded_image['data'])
                else:
                    image_references.append(image.url)
                seen_reference = True
            else:
                assert_never(image)

        if len(images) == 1:
            if file_ids:
                return _XaiInputImages(image_file_id=file_ids[0])
            return _XaiInputImages(image_url=image_references[0])

        # Reported after the loop so a per-image validation error takes precedence over the ordering.
        if order_violated:
            raise UserError(
                'xAI sends file-ID image inputs before URL or binary inputs. '
                'Place all `UploadedFile` inputs first to preserve reference-image order.'
            )

        return _XaiInputImages(image_urls=image_references or None, image_file_ids=file_ids or None)

    def _map_response(
        self,
        prompt: str,
        responses: Sequence[ImageResponse],
    ) -> ImageGenerationResult:
        if not responses:
            raise UnexpectedModelBehavior('xAI image generation response did not contain any images')

        images: list[GeneratedImage] = []
        moderated_indices: list[int] = []
        for index, response in enumerate(responses):
            # xAI moderation is silent: a flagged slot comes back with `respect_moderation=False` and an
            # empty payload, and reading its `.base64` raises client-side. Skip it so one flagged slot
            # doesn't discard the rest of a paid batch.
            if not response.respect_moderation:
                moderated_indices.append(index)
                continue
            try:
                content = _decode_data_url(response.base64)
            except (ValueError, TypeError) as e:
                raise UnexpectedModelBehavior(
                    'xAI image generation response did not contain valid base64 image data'
                ) from e
            images.append(
                GeneratedImage(
                    content=content,
                    output_format=output_format_from_media_type(content.media_type),
                    provider_details={'respect_moderation': response.respect_moderation},
                )
            )

        if not images:
            raise ContentFilterError('xAI flagged all generated images for content moderation')

        # A batch is one `GenerateImage` RPC that answers with one `ImageResponse` proto holding every
        # image, which `sample_batch` wraps as n views over that single proto — so `usage` and
        # `cost_usd` are the same batch-wide object on each element, not a per-image share. Reading the
        # first is therefore the whole batch; summing would multiply it by n. Recorded across
        # `test_xai_image_generation_vcr` and `test_xai_image_generation_batch_vcr`: one image costs
        # $0.02 against the $0.04 a two-image batch reports.
        first_response = responses[0]
        provider_details = _response_provider_details(first_response)
        if moderated_indices:
            provider_details['moderated_image_indices'] = moderated_indices
        return ImageGenerationResult(
            images=images,
            prompt=prompt,
            usage=_map_usage(first_response.usage, self.system, self.base_url, self.model_name),
            model_name=first_response.model or self.model_name,
            provider_name=self.system,
            provider_url=self.base_url,
            provider_details=provider_details,
        )


def _decode_data_url(value: str) -> BinaryImage:
    """Decode a base64 image data URL into a `BinaryImage`.

    Stricter than `BinaryContent.from_data_uri`, which decodes without `validate=True`: a malformed
    payload must raise here rather than silently decode to truncated image bytes.

    The header's media type is the provider's own claim and is trusted, as Google's `mime_type` is.
    Only the OpenAI adapter overrides a provider's claim by sniffing the bytes, because gpt-image
    echoes a requested `output_format` it did not return (openai-node#1850).
    """
    header, encoded = value.split(',', maxsplit=1)
    if not header.startswith('data:image/') or not header.endswith(';base64'):
        raise ValueError('Not a base64 image data URL')
    media_type = header.removeprefix('data:').removesuffix(';base64')
    return BinaryImage(data=base64.b64decode(encoded, validate=True), media_type=media_type)


@dataclass
class _XaiResolvedSettings:
    aspect_ratio: XaiImageAspectRatio | None
    resolution: ImageResolution | None
    ignored: list[str]
    conflicts: list[str]


def _resolve_xai_settings(
    settings: XaiImageGenerationSettings, *, model_name: XaiImageGenerationModelName
) -> _XaiResolvedSettings:
    validate_image_count('xAI', settings.get('xai_n'))
    geometry = resolve_xai_geometry(
        model_name,
        settings,
        provider_aspect_ratio=settings.get('xai_aspect_ratio'),
        provider_resolution=settings.get('xai_resolution'),
    )

    # xAI is reached over gRPC, which has no per-request body or header escape hatch, so these
    # portable settings cannot be honored here as they are on the HTTP-based providers.
    ignored: list[str] = []
    if settings.get('extra_headers'):
        ignored.append('extra_headers')
    if settings.get('extra_body'):
        ignored.append('extra_body')

    return _XaiResolvedSettings(
        aspect_ratio=geometry.aspect_ratio,
        resolution=geometry.resolution,
        ignored=ignored,
        conflicts=geometry.conflicts,
    )


def _map_usage(
    usage: usage_pb2.SamplingUsage,
    provider: str,
    provider_url: str,
    model: str,
) -> RequestUsage:
    details: dict[str, int] = {}
    if reasoning_tokens := usage.reasoning_tokens:
        details['reasoning_tokens'] = reasoning_tokens
    if prompt_text_tokens := usage.prompt_text_tokens:
        details['input_text_tokens'] = prompt_text_tokens
    if prompt_image_tokens := usage.prompt_image_tokens:
        details['input_image_tokens'] = prompt_image_tokens

    # `cached_prompt_text_tokens` is fed to `extract` rather than `details` so genai-prices maps it
    # onto the typed `cache_read_tokens` and prices it at the cached rate, matching `models/xai.py`.
    usage_data: dict[str, int] = {
        'prompt_tokens': usage.prompt_tokens,
        'completion_tokens': usage.completion_tokens,
    }
    if cached_tokens := usage.cached_prompt_text_tokens:
        usage_data['cached_prompt_text_tokens'] = cached_tokens

    extracted_usage = RequestUsage.extract(
        {'model': model, 'usage': usage_data},
        provider=provider,
        provider_url=provider_url,
        provider_fallback='x-ai',
        details=details,
    )
    # Backfill only the counts genai-prices failed to derive, so the typed fields it did populate
    # (notably `cache_read_tokens`) survive.
    if extracted_usage.input_tokens == 0 and usage.prompt_tokens:
        extracted_usage.input_tokens = usage.prompt_tokens
    if extracted_usage.output_tokens == 0 and usage.completion_tokens:
        extracted_usage.output_tokens = usage.completion_tokens

    return extracted_usage


def _response_provider_details(response: ImageResponse) -> dict[str, object]:
    provider_details: dict[str, object] = {}
    usage = response.usage
    if usage.HasField('cost_in_usd_ticks'):
        provider_details['cost_in_usd_ticks'] = usage.cost_in_usd_ticks
    if (cost_usd := response.cost_usd) is not None:
        provider_details['cost_usd'] = cost_usd
    return provider_details


# The image path additionally maps `INVALID_ARGUMENT` to 400, which the chat table leaves unmapped.
_GRPC_STATUS_TO_HTTP: dict[grpc.StatusCode, int] = {
    **_CHAT_GRPC_STATUS_TO_HTTP,
    grpc.StatusCode.INVALID_ARGUMENT: 400,
}