"""Deprecated Gemini model implementation.

This module is deprecated. Use [`pydantic_ai.models.google.GoogleModel`][pydantic_ai.models.google.GoogleModel] instead.
See the [Google model documentation](https://ai.pydantic.dev/models/google/) for more details.

This module uses a custom HTTP implementation. The recommended `GoogleModel` in
`google.py` uses the official `google-genai` SDK.
"""

from __future__ import annotations as _annotations

from collections.abc import AsyncIterator, Sequence
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from datetime import datetime
from typing import Annotated, Any, Literal, Protocol, cast
from uuid import uuid4

import httpx
import pydantic
from httpx import USE_CLIENT_DEFAULT, Response as HTTPResponse
from typing_extensions import NotRequired, TypedDict, assert_never, deprecated

from .. import ModelHTTPError, UnexpectedModelBehavior, _utils, usage
from .._output import OutputObjectDefinition
from .._run_context import RunContext
from ..exceptions import UserError
from ..messages import (
    BinaryContent,
    BuiltinToolCallPart,
    BuiltinToolReturnPart,
    CachePoint,
    FilePart,
    FileUrl,
    ModelMessage,
    ModelRequest,
    ModelResponse,
    ModelResponsePart,
    ModelResponseStreamEvent,
    RetryPromptPart,
    SystemPromptPart,
    TextPart,
    ThinkingPart,
    ToolCallPart,
    ToolReturnPart,
    UserPromptPart,
    VideoUrl,
)
from ..profiles import ModelProfileSpec
from ..providers import Provider
from ..settings import ModelSettings
from ..tools import ToolDefinition
from . import Model, ModelRequestParameters, StreamedResponse, check_allow_model_requests, download_item, get_user_agent

LatestGeminiModelNames = Literal[
    'gemini-2.0-flash',
    'gemini-2.0-flash-lite',
    'gemini-2.5-flash',
    'gemini-2.5-flash-preview-09-2025',
    'gemini-2.5-flash-lite',
    'gemini-2.5-flash-lite-preview-09-2025',
    'gemini-flash-latest',
    'gemini-flash-lite-latest',
    'gemini-2.5-pro',
]
"""Latest Gemini models."""

GeminiModelName = str | LatestGeminiModelNames
"""Possible Gemini model names.

Since Gemini supports a variety of date-stamped models, we explicitly list the latest models but
allow any name in the type hints.
See [the Gemini API docs](https://ai.google.dev/gemini-api/docs/models/gemini#model-variations) for a full list.
"""


class GeminiModelSettings(ModelSettings, total=False):
    """Settings used for a Gemini model request."""

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

    gemini_safety_settings: list[GeminiSafetySettings]
    """Safety settings options for Gemini model request."""

    gemini_thinking_config: ThinkingConfig
    """Thinking is "on" by default in both the API and AI Studio.

    Being on by default doesn't mean the model will send back thoughts. For that, you would need to set `include_thoughts`
    to `True`. If you want to avoid the model spending any tokens on thinking, you can set `thinking_budget` to `0`.

    See more about it on <https://ai.google.dev/gemini-api/docs/thinking>.
    """

    gemini_labels: dict[str, str]
    """User-defined metadata to break down billed charges. Only supported by the Vertex AI provider.

    See the [Gemini API docs](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/add-labels-to-api-calls) for use cases and limitations.
    """


@deprecated('Use `GoogleModel` instead. See <https://ai.pydantic.dev/models/google/> for more details.')
@dataclass(init=False)
class GeminiModel(Model):
    """A model that uses Gemini via `generativelanguage.googleapis.com` API.

    This is implemented from scratch rather than using a dedicated SDK, good API documentation is
    available [here](https://ai.google.dev/api).

    Apart from `__init__`, all methods are private or match those of the base class.
    """

    client: httpx.AsyncClient = field(repr=False)

    _model_name: GeminiModelName = field(repr=False)
    _provider: Provider[httpx.AsyncClient] = field(repr=False)
    _auth: AuthProtocol | None = field(repr=False)
    _url: str | None = field(repr=False)

    def __init__(
        self,
        model_name: GeminiModelName,
        *,
        provider: Literal['google-gla', 'google-vertex'] | Provider[httpx.AsyncClient] = 'google-gla',
        profile: ModelProfileSpec | None = None,
        settings: ModelSettings | None = None,
    ):
        """Initialize a Gemini model.

        Args:
            model_name: The name of the model to use.
            provider: The provider to use for authentication and API access. Can be either the string
                'google-gla' or 'google-vertex' or an instance of `Provider[httpx.AsyncClient]`.
                If not provided, a new provider will be created using the other parameters.
            profile: The model profile to use. Defaults to a profile picked by the provider based on the model name.
            settings: Default model settings for this model instance.
        """
        self._model_name = model_name

        if isinstance(provider, str):
            if provider == 'google-gla':
                from ..providers.google_gla import GoogleGLAProvider  # type: ignore[reportDeprecated]

                provider = GoogleGLAProvider()  # type: ignore[reportDeprecated]
            else:
                from ..providers.google_vertex import GoogleVertexProvider  # type: ignore[reportDeprecated]

                provider = GoogleVertexProvider()  # type: ignore[reportDeprecated]
        self._provider = provider
        self.client = provider.client
        self._url = str(self.client.base_url)

        super().__init__(settings=settings, profile=profile or provider.model_profile)

    @property
    def base_url(self) -> str:
        assert self._url is not None, 'URL not initialized'
        return self._url

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

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

    async def request(
        self,
        messages: list[ModelMessage],
        model_settings: ModelSettings | None,
        model_request_parameters: ModelRequestParameters,
    ) -> ModelResponse:
        check_allow_model_requests()
        model_settings, model_request_parameters = self.prepare_request(
            model_settings,
            model_request_parameters,
        )
        async with self._make_request(
            messages, False, cast(GeminiModelSettings, model_settings or {}), model_request_parameters
        ) as http_response:
            data = await http_response.aread()
            response = _gemini_response_ta.validate_json(data)
        return self._process_response(response)

    @asynccontextmanager
    async def request_stream(
        self,
        messages: list[ModelMessage],
        model_settings: ModelSettings | None,
        model_request_parameters: ModelRequestParameters,
        run_context: RunContext[Any] | None = None,
    ) -> AsyncIterator[StreamedResponse]:
        check_allow_model_requests()
        model_settings, model_request_parameters = self.prepare_request(
            model_settings,
            model_request_parameters,
        )
        async with self._make_request(
            messages, True, cast(GeminiModelSettings, model_settings or {}), model_request_parameters
        ) as http_response:
            yield await self._process_streamed_response(http_response, model_request_parameters)

    def _get_tools(self, model_request_parameters: ModelRequestParameters) -> _GeminiTools | None:
        tools = [_function_from_abstract_tool(t) for t in model_request_parameters.tool_defs.values()]
        return _GeminiTools(function_declarations=tools) if tools else None

    def _get_tool_config(
        self, model_request_parameters: ModelRequestParameters, tools: _GeminiTools | None
    ) -> _GeminiToolConfig | None:
        if not model_request_parameters.allow_text_output and tools:
            return _tool_config([t['name'] for t in tools['function_declarations']])
        else:
            return None

    @asynccontextmanager
    async def _make_request(
        self,
        messages: list[ModelMessage],
        streamed: bool,
        model_settings: GeminiModelSettings,
        model_request_parameters: ModelRequestParameters,
    ) -> AsyncIterator[HTTPResponse]:
        tools = self._get_tools(model_request_parameters)
        tool_config = self._get_tool_config(model_request_parameters, tools)
        sys_prompt_parts, contents = await self._message_to_gemini_content(messages, model_request_parameters)

        request_data = _GeminiRequest(contents=contents)
        if sys_prompt_parts:
            request_data['systemInstruction'] = _GeminiTextContent(role='user', parts=sys_prompt_parts)
        if tools is not None:
            request_data['tools'] = tools
        if tool_config is not None:
            request_data['toolConfig'] = tool_config

        generation_config = _settings_to_generation_config(model_settings)
        if model_request_parameters.output_mode == 'native':
            if tools:
                raise UserError(
                    'Gemini does not support `NativeOutput` and tools at the same time. Use `output_type=ToolOutput(...)` instead.'
                )

            generation_config['response_mime_type'] = 'application/json'

            output_object = model_request_parameters.output_object
            assert output_object is not None
            generation_config['response_json_schema'] = self._map_response_schema(output_object)
        elif model_request_parameters.output_mode == 'prompted' and not tools:
            generation_config['response_mime_type'] = 'application/json'

        if generation_config:
            request_data['generationConfig'] = generation_config

        if gemini_safety_settings := model_settings.get('gemini_safety_settings'):
            request_data['safetySettings'] = gemini_safety_settings

        if gemini_labels := model_settings.get('gemini_labels'):
            if self._provider.name == 'google-vertex':
                request_data['labels'] = gemini_labels  # pragma: lax no cover

        headers = {'Content-Type': 'application/json', 'User-Agent': get_user_agent()}
        url = f'/{self._model_name}:{"streamGenerateContent" if streamed else "generateContent"}'

        request_json = _gemini_request_ta.dump_json(request_data, by_alias=True)
        async with self.client.stream(
            'POST',
            url,
            content=request_json,
            headers=headers,
            timeout=model_settings.get('timeout', USE_CLIENT_DEFAULT),
        ) as r:
            if (status_code := r.status_code) != 200:
                await r.aread()
                if status_code >= 400:
                    raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=r.text)
                raise UnexpectedModelBehavior(  # pragma: no cover
                    f'Unexpected response from gemini {status_code}', r.text
                )
            yield r

    def _process_response(self, response: _GeminiResponse) -> ModelResponse:
        vendor_details: dict[str, Any] | None = None

        if len(response['candidates']) != 1:
            raise UnexpectedModelBehavior('Expected exactly one candidate in Gemini response')  # pragma: no cover
        if 'content' not in response['candidates'][0]:
            if response['candidates'][0].get('finish_reason') == 'SAFETY':
                raise UnexpectedModelBehavior('Safety settings triggered', str(response))
            else:
                raise UnexpectedModelBehavior(  # pragma: no cover
                    'Content field missing from Gemini response', str(response)
                )
        parts = response['candidates'][0]['content']['parts']
        vendor_id = response.get('vendor_id', None)
        finish_reason = response['candidates'][0].get('finish_reason')
        if finish_reason:
            vendor_details = {'finish_reason': finish_reason}
        usage = _metadata_as_usage(response)
        return _process_response_from_parts(
            parts,
            response.get('model_version', self._model_name),
            usage,
            vendor_id=vendor_id,
            vendor_details=vendor_details,
            provider_name=self._provider.name,
            provider_url=self.base_url,
        )

    async def _process_streamed_response(
        self, http_response: HTTPResponse, model_request_parameters: ModelRequestParameters
    ) -> StreamedResponse:
        """Process a streamed response, and prepare a streaming response to return."""
        aiter_bytes = http_response.aiter_bytes()
        start_response: _GeminiResponse | None = None
        content = bytearray()

        async for chunk in aiter_bytes:
            content.extend(chunk)
            responses = _gemini_streamed_response_ta.validate_json(
                _ensure_decodeable(content),
                experimental_allow_partial='trailing-strings',
            )
            if responses:  # pragma: no branch
                last = responses[-1]
                if last['candidates'] and last['candidates'][0].get('content', {}).get('parts'):
                    start_response = last
                    break

        if start_response is None:
            raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')

        return GeminiStreamedResponse(
            model_request_parameters=model_request_parameters,
            _model_name=self._model_name,
            _content=content,
            _stream=aiter_bytes,
            _provider_name=self._provider.name,
            _provider_url=self.base_url,
        )

    async def _message_to_gemini_content(
        self, messages: list[ModelMessage], model_request_parameters: ModelRequestParameters
    ) -> tuple[list[_GeminiTextPart], list[_GeminiContent]]:
        sys_prompt_parts: list[_GeminiTextPart] = []
        contents: list[_GeminiContent] = []
        for m in messages:
            if isinstance(m, ModelRequest):
                message_parts: list[_GeminiPartUnion] = []

                for part in m.parts:
                    if isinstance(part, SystemPromptPart):
                        sys_prompt_parts.append(_GeminiTextPart(text=part.content))
                    elif isinstance(part, UserPromptPart):
                        message_parts.extend(await self._map_user_prompt(part))
                    elif isinstance(part, ToolReturnPart):
                        message_parts.append(_response_part_from_response(part.tool_name, part.model_response_object()))
                    elif isinstance(part, RetryPromptPart):
                        if part.tool_name is None:
                            message_parts.append(_GeminiTextPart(text=part.model_response()))  # pragma: no cover
                        else:
                            response = {'call_error': part.model_response()}
                            message_parts.append(_response_part_from_response(part.tool_name, response))
                    else:
                        assert_never(part)

                if message_parts:  # pragma: no branch
                    contents.append(_GeminiContent(role='user', parts=message_parts))
            elif isinstance(m, ModelResponse):
                contents.append(_content_model_response(m))
            else:
                assert_never(m)
        if instructions := self._get_instructions(messages, model_request_parameters):
            sys_prompt_parts.append(_GeminiTextPart(text=instructions))
        return sys_prompt_parts, contents

    async def _map_user_prompt(self, part: UserPromptPart) -> list[_GeminiPartUnion]:
        if isinstance(part.content, str):
            return [{'text': part.content}]
        else:
            content: list[_GeminiPartUnion] = []
            for item in part.content:
                if isinstance(item, str):
                    content.append({'text': item})
                elif isinstance(item, BinaryContent):
                    content.append(
                        _GeminiInlineDataPart(inline_data={'data': item.base64, 'mime_type': item.media_type})
                    )
                elif isinstance(item, VideoUrl) and item.is_youtube:
                    file_data = _GeminiFileDataPart(file_data={'file_uri': item.url, 'mime_type': item.media_type})
                    content.append(file_data)
                elif isinstance(item, FileUrl):
                    if self.system == 'google-gla' or item.force_download:
                        downloaded_item = await download_item(item, data_format='base64')
                        inline_data = _GeminiInlineDataPart(
                            inline_data={'data': downloaded_item['data'], 'mime_type': downloaded_item['data_type']}
                        )
                        content.append(inline_data)
                    else:  # pragma: lax no cover
                        file_data = _GeminiFileDataPart(file_data={'file_uri': item.url, 'mime_type': item.media_type})
                        content.append(file_data)
                elif isinstance(item, CachePoint):
                    # Gemini doesn't support inline CachePoint markers. Google's caching requires
                    # pre-creating cache objects via the API, then referencing them by name using
                    # `GoogleModelSettings.google_cached_content`. See https://ai.google.dev/gemini-api/docs/caching
                    pass
                else:
                    assert_never(item)  # pragma: lax no cover
        return content

    def _map_response_schema(self, o: OutputObjectDefinition) -> dict[str, Any]:
        response_schema = o.json_schema.copy()
        if o.name:
            response_schema['title'] = o.name
        if o.description:
            response_schema['description'] = o.description

        return response_schema


def _settings_to_generation_config(model_settings: GeminiModelSettings) -> _GeminiGenerationConfig:
    config: _GeminiGenerationConfig = {}
    if (max_tokens := model_settings.get('max_tokens')) is not None:
        config['max_output_tokens'] = max_tokens
    if (stop_sequences := model_settings.get('stop_sequences')) is not None:
        config['stop_sequences'] = stop_sequences  # pragma: no cover
    if (temperature := model_settings.get('temperature')) is not None:
        config['temperature'] = temperature
    if (top_p := model_settings.get('top_p')) is not None:
        config['top_p'] = top_p
    if (presence_penalty := model_settings.get('presence_penalty')) is not None:
        config['presence_penalty'] = presence_penalty
    if (frequency_penalty := model_settings.get('frequency_penalty')) is not None:
        config['frequency_penalty'] = frequency_penalty
    if (thinkingConfig := model_settings.get('gemini_thinking_config')) is not None:
        config['thinking_config'] = thinkingConfig
    return config


class AuthProtocol(Protocol):
    """Abstract definition for Gemini authentication."""

    async def headers(self) -> dict[str, str]: ...


@dataclass
class ApiKeyAuth:
    """Authentication using an API key for the `X-Goog-Api-Key` header."""

    api_key: str

    async def headers(self) -> dict[str, str]:
        # https://cloud.google.com/docs/authentication/api-keys-use#using-with-rest
        return {'X-Goog-Api-Key': self.api_key}  # pragma: no cover


@dataclass
class GeminiStreamedResponse(StreamedResponse):
    """Implementation of `StreamedResponse` for the Gemini model."""

    _model_name: GeminiModelName
    _content: bytearray
    _stream: AsyncIterator[bytes]
    _provider_name: str
    _provider_url: str
    _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)

    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
        async for gemini_response in self._get_gemini_responses():
            candidate = gemini_response['candidates'][0]
            if 'content' not in candidate:
                raise UnexpectedModelBehavior('Streamed response has no content field')  # pragma: no cover
            gemini_part: _GeminiPartUnion
            for gemini_part in candidate['content']['parts']:
                if 'text' in gemini_part:
                    # Using vendor_part_id=None means we can produce multiple text parts if their deltas are sprinkled
                    # amongst the tool call deltas
                    for event in self._parts_manager.handle_text_delta(
                        vendor_part_id=None, content=gemini_part['text']
                    ):
                        yield event

                elif 'function_call' in gemini_part:
                    # Here, we assume all function_call parts are complete and don't have deltas.
                    # We do this by assigning a unique randomly generated "vendor_part_id".
                    # We need to confirm whether this is actually true, but if it isn't, we can still handle it properly
                    # it would just be a bit more complicated. And we'd need to confirm the intended semantics.
                    maybe_event = self._parts_manager.handle_tool_call_delta(
                        vendor_part_id=uuid4(),
                        tool_name=gemini_part['function_call']['name'],
                        args=gemini_part['function_call']['args'],
                        tool_call_id=None,
                    )
                    if maybe_event is not None:  # pragma: no branch
                        yield maybe_event
                else:
                    if not any([key in gemini_part for key in ['function_response', 'thought']]):
                        raise AssertionError(f'Unexpected part: {gemini_part}')  # pragma: no cover

    async def _get_gemini_responses(self) -> AsyncIterator[_GeminiResponse]:
        # This method exists to ensure we only yield completed items, so we don't need to worry about
        # partial gemini responses, which would make everything more complicated

        gemini_responses: list[_GeminiResponse] = []
        current_gemini_response_index = 0
        # Right now, there are some circumstances where we will have information that could be yielded sooner than it is
        # But changing that would make things a lot more complicated.
        async for chunk in self._stream:
            self._content.extend(chunk)

            gemini_responses = _gemini_streamed_response_ta.validate_json(
                _ensure_decodeable(self._content),
                experimental_allow_partial='trailing-strings',
            )

            # The idea: yield only up to the latest response, which might still be partial.
            # Note that if the latest response is complete, we could yield it immediately, but there's not a good
            # allow_partial API to determine if the last item in the list is complete.
            responses_to_yield = gemini_responses[:-1]
            for r in responses_to_yield[current_gemini_response_index:]:
                current_gemini_response_index += 1
                yield r

        # Now yield the final response, which should be complete
        if gemini_responses:  # pragma: no branch
            r = gemini_responses[-1]
            self._usage = _metadata_as_usage(r)
            yield r

    @property
    def model_name(self) -> GeminiModelName:
        """Get the model name of the response."""
        return self._model_name

    @property
    def provider_name(self) -> str:
        """Get the provider name."""
        return self._provider_name

    @property
    def provider_url(self) -> str:
        """Get the provider base URL."""
        return self._provider_url

    @property
    def timestamp(self) -> datetime:
        """Get the timestamp of the response."""
        return self._timestamp


# We use typed dicts to define the Gemini API response schema
# once Pydantic partial validation supports, dataclasses, we could revert to using them
# TypeAdapters take care of validation and serialization


@pydantic.with_config(pydantic.ConfigDict(defer_build=True))
class _GeminiRequest(TypedDict):
    """Schema for an API request to the Gemini API.

    See <https://ai.google.dev/api/generate-content#request-body> for API docs.
    """

    # Note: Even though Google supposedly supports camelCase and snake_case, we've had user report misbehavior
    # when using snake_case, which is why this typeddict now uses camelCase. The recommended `GoogleModel` in
    # `google.py` uses the official `google-genai` SDK and should be used instead of this deprecated module.
    contents: list[_GeminiContent]
    tools: NotRequired[_GeminiTools]
    toolConfig: NotRequired[_GeminiToolConfig]
    safetySettings: NotRequired[list[GeminiSafetySettings]]
    systemInstruction: NotRequired[_GeminiTextContent]
    """
    Developer generated system instructions, see
    <https://ai.google.dev/gemini-api/docs/system-instructions?lang=rest>
    """
    generationConfig: NotRequired[_GeminiGenerationConfig]
    labels: NotRequired[dict[str, str]]


class GeminiSafetySettings(TypedDict):
    """Safety settings options for Gemini model request.

    See [Gemini API docs](https://ai.google.dev/gemini-api/docs/safety-settings) for safety category and threshold descriptions.
    For an example on how to use `GeminiSafetySettings`, see [here](../../agents.md#model-specific-settings).
    """

    category: Literal[
        'HARM_CATEGORY_UNSPECIFIED',
        'HARM_CATEGORY_HARASSMENT',
        'HARM_CATEGORY_HATE_SPEECH',
        'HARM_CATEGORY_SEXUALLY_EXPLICIT',
        'HARM_CATEGORY_DANGEROUS_CONTENT',
        'HARM_CATEGORY_CIVIC_INTEGRITY',
    ]
    """
    Safety settings category.
    """

    threshold: Literal[
        'HARM_BLOCK_THRESHOLD_UNSPECIFIED',
        'BLOCK_LOW_AND_ABOVE',
        'BLOCK_MEDIUM_AND_ABOVE',
        'BLOCK_ONLY_HIGH',
        'BLOCK_NONE',
        'OFF',
    ]
    """
    Safety settings threshold.
    """


class ThinkingConfig(TypedDict, total=False):
    """The thinking features configuration."""

    include_thoughts: Annotated[bool, pydantic.Field(alias='includeThoughts')]
    """Indicates whether to include thoughts in the response. If true, thoughts are returned only if the model supports thought and thoughts are available."""

    thinking_budget: Annotated[int, pydantic.Field(alias='thinkingBudget')]
    """Indicates the thinking budget in tokens."""


class _GeminiGenerationConfig(TypedDict, total=False):
    """Schema for an API request to the Gemini API.

    Note there are many additional fields available that have not been added yet.

    See <https://ai.google.dev/api/generate-content#generationconfig> for API docs.
    """

    max_output_tokens: int
    temperature: float
    top_p: float
    presence_penalty: float
    frequency_penalty: float
    stop_sequences: list[str]
    thinking_config: ThinkingConfig
    response_mime_type: str
    response_json_schema: dict[str, Any]


class _GeminiContent(TypedDict):
    role: Literal['user', 'model']
    parts: list[_GeminiPartUnion]


def _content_model_response(m: ModelResponse) -> _GeminiContent:
    parts: list[_GeminiPartUnion] = []
    function_call_requires_signature = True
    for item in m.parts:
        if isinstance(item, ToolCallPart):
            part = _function_call_part_from_call(item)
            if function_call_requires_signature and not part.get('thought_signature'):
                # Per https://ai.google.dev/gemini-api/docs/thought-signatures#faqs:
                # > You can set the following dummy signatures of either "context_engineering_is_the_way_to_go"
                # > or "skip_thought_signature_validator"
                # Per https://cloud.google.com/vertex-ai/generative-ai/docs/thought-signatures#using-rest-or-manual-handling:
                # > You can set thought_signature to skip_thought_signature_validator
                # We use "skip_thought_signature_validator" as it works for both Gemini API and Vertex AI.
                part['thought_signature'] = b'skip_thought_signature_validator'
            # Only the first function call requires a signature
            function_call_requires_signature = False
            parts.append(part)
        elif isinstance(item, ThinkingPart):
            # NOTE: We don't send ThinkingPart to the providers yet. If you are unsatisfied with this,
            # please open an issue. The below code is the code to send thinking to the provider.
            # parts.append(_GeminiTextPart(text=item.content, thought=True))
            pass
        elif isinstance(item, TextPart):
            if item.content:
                parts.append(_GeminiTextPart(text=item.content))
        elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
            # This is currently never returned from gemini
            pass
        elif isinstance(item, FilePart):  # pragma: no cover
            # Files generated by models are not sent back to models that don't themselves generate files.
            pass
        else:
            assert_never(item)
    return _GeminiContent(role='model', parts=parts)


class _BasePart(TypedDict):
    thought: NotRequired[bool]
    """Indicates if the part is thought from the model."""


class _GeminiTextPart(_BasePart):
    text: str


class _GeminiInlineData(_BasePart):
    data: str
    mime_type: Annotated[str, pydantic.Field(alias='mimeType')]


class _GeminiInlineDataPart(_BasePart):
    """See <https://ai.google.dev/api/caching#Blob>."""

    inline_data: Annotated[_GeminiInlineData, pydantic.Field(alias='inlineData')]


class _GeminiFileData(_BasePart):
    """See <https://ai.google.dev/api/caching#FileData>."""

    file_uri: Annotated[str, pydantic.Field(alias='fileUri')]
    mime_type: Annotated[str, pydantic.Field(alias='mimeType')]


class _GeminiFileDataPart(_BasePart):
    file_data: Annotated[_GeminiFileData, pydantic.Field(alias='fileData')]


class _GeminiThoughtPart(TypedDict):
    thought: bool
    thought_signature: Annotated[str, pydantic.Field(alias='thoughtSignature')]


class _GeminiFunctionCallPart(_BasePart):
    function_call: Annotated[_GeminiFunctionCall, pydantic.Field(alias='functionCall')]

    thought_signature: NotRequired[Annotated[bytes, pydantic.Field(alias='thoughtSignature')]]


def _function_call_part_from_call(tool: ToolCallPart) -> _GeminiFunctionCallPart:
    return _GeminiFunctionCallPart(function_call=_GeminiFunctionCall(name=tool.tool_name, args=tool.args_as_dict()))


def _process_response_from_parts(
    parts: Sequence[_GeminiPartUnion],
    model_name: GeminiModelName,
    usage: usage.RequestUsage,
    vendor_id: str | None,
    provider_name: str,
    provider_url: str,
    vendor_details: dict[str, Any] | None = None,
) -> ModelResponse:
    items: list[ModelResponsePart] = []
    for part in parts:
        if 'text' in part:
            if part.get('thought'):
                items.append(ThinkingPart(content=part['text']))
            else:
                items.append(TextPart(content=part['text']))
        elif 'function_call' in part:
            items.append(ToolCallPart(tool_name=part['function_call']['name'], args=part['function_call']['args']))
        elif 'function_response' in part:  # pragma: no cover
            raise UnexpectedModelBehavior(
                f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}'
            )
    return ModelResponse(
        parts=items,
        usage=usage,
        model_name=model_name,
        provider_name=provider_name,
        provider_response_id=vendor_id,
        provider_details=vendor_details,
        provider_url=provider_url,
    )


class _GeminiFunctionCall(TypedDict):
    """See <https://ai.google.dev/api/caching#FunctionCall>."""

    name: str
    args: dict[str, Any]


class _GeminiFunctionResponsePart(TypedDict):
    function_response: Annotated[_GeminiFunctionResponse, pydantic.Field(alias='functionResponse')]


def _response_part_from_response(name: str, response: dict[str, Any]) -> _GeminiFunctionResponsePart:
    return _GeminiFunctionResponsePart(function_response=_GeminiFunctionResponse(name=name, response=response))


class _GeminiFunctionResponse(TypedDict):
    """See <https://ai.google.dev/api/caching#FunctionResponse>."""

    name: str
    response: dict[str, Any]


def _part_discriminator(v: Any) -> str:
    if isinstance(v, dict):  # pragma: no branch
        if 'text' in v:
            return 'text'
        elif 'inlineData' in v:
            return 'inline_data'  # pragma: no cover
        elif 'fileData' in v:
            return 'file_data'  # pragma: no cover
        elif 'thought' in v:
            return 'thought'
        elif 'functionCall' in v or 'function_call' in v:
            return 'function_call'
        elif 'functionResponse' in v or 'function_response' in v:
            return 'function_response'
    return 'text'


# See <https://ai.google.dev/api/caching#Part>
# we don't currently support other part types
_GeminiPartUnion = Annotated[
    Annotated[_GeminiTextPart, pydantic.Tag('text')]
    | Annotated[_GeminiFunctionCallPart, pydantic.Tag('function_call')]
    | Annotated[_GeminiFunctionResponsePart, pydantic.Tag('function_response')]
    | Annotated[_GeminiInlineDataPart, pydantic.Tag('inline_data')]
    | Annotated[_GeminiFileDataPart, pydantic.Tag('file_data')]
    | Annotated[_GeminiThoughtPart, pydantic.Tag('thought')],
    pydantic.Discriminator(_part_discriminator),
]


class _GeminiTextContent(TypedDict):
    role: Literal['user', 'model']
    parts: list[_GeminiTextPart]


class _GeminiTools(TypedDict):
    function_declarations: Annotated[list[_GeminiFunction], pydantic.Field(alias='functionDeclarations')]


class _GeminiFunction(TypedDict):
    name: str
    description: str
    parameters_json_schema: NotRequired[dict[str, Any]]


def _function_from_abstract_tool(tool: ToolDefinition) -> _GeminiFunction:
    json_schema = tool.parameters_json_schema
    f = _GeminiFunction(name=tool.name, description=tool.description or '', parameters_json_schema=json_schema)
    return f


class _GeminiToolConfig(TypedDict):
    function_calling_config: _GeminiFunctionCallingConfig


def _tool_config(function_names: list[str]) -> _GeminiToolConfig:
    return _GeminiToolConfig(
        function_calling_config=_GeminiFunctionCallingConfig(mode='ANY', allowed_function_names=function_names)
    )


class _GeminiFunctionCallingConfig(TypedDict):
    mode: Literal['ANY', 'AUTO']
    allowed_function_names: list[str]


@pydantic.with_config(pydantic.ConfigDict(defer_build=True))
class _GeminiResponse(TypedDict):
    """Schema for the response from the Gemini API.

    See <https://ai.google.dev/api/generate-content#v1beta.GenerateContentResponse>
    and <https://cloud.google.com/vertex-ai/docs/reference/rest/v1/GenerateContentResponse>
    """

    candidates: list[_GeminiCandidates]
    # usageMetadata appears to be required by both APIs but is omitted when streaming responses until the last response
    usage_metadata: NotRequired[Annotated[_GeminiUsageMetaData, pydantic.Field(alias='usageMetadata')]]
    prompt_feedback: NotRequired[Annotated[_GeminiPromptFeedback, pydantic.Field(alias='promptFeedback')]]
    model_version: NotRequired[Annotated[str, pydantic.Field(alias='modelVersion')]]
    vendor_id: NotRequired[Annotated[str, pydantic.Field(alias='responseId')]]


class _GeminiCandidates(TypedDict):
    """See <https://ai.google.dev/api/generate-content#v1beta.Candidate>."""

    content: NotRequired[_GeminiContent]
    finish_reason: NotRequired[Annotated[Literal['STOP', 'MAX_TOKENS', 'SAFETY'], pydantic.Field(alias='finishReason')]]
    """
    See <https://ai.google.dev/api/generate-content#FinishReason>, lots of other values are possible,
    but let's wait until we see them and know what they mean to add them here.
    """
    avg_log_probs: NotRequired[Annotated[float, pydantic.Field(alias='avgLogProbs')]]
    index: NotRequired[int]
    safety_ratings: NotRequired[Annotated[list[_GeminiSafetyRating], pydantic.Field(alias='safetyRatings')]]


class _GeminiModalityTokenCount(TypedDict):
    """See <https://ai.google.dev/api/generate-content#modalitytokencount>."""

    modality: Annotated[
        Literal['MODALITY_UNSPECIFIED', 'TEXT', 'IMAGE', 'VIDEO', 'AUDIO', 'DOCUMENT'], pydantic.Field(alias='modality')
    ]
    token_count: Annotated[int, pydantic.Field(alias='tokenCount', default=0)]


class _GeminiUsageMetaData(TypedDict, total=False):
    """See <https://ai.google.dev/api/generate-content#UsageMetadata>.

    The docs suggest all fields are required, but some are actually not required, so we assume they are all optional.
    """

    prompt_token_count: Annotated[int, pydantic.Field(alias='promptTokenCount')]
    candidates_token_count: NotRequired[Annotated[int, pydantic.Field(alias='candidatesTokenCount')]]
    total_token_count: Annotated[int, pydantic.Field(alias='totalTokenCount')]
    cached_content_token_count: NotRequired[Annotated[int, pydantic.Field(alias='cachedContentTokenCount')]]
    thoughts_token_count: NotRequired[Annotated[int, pydantic.Field(alias='thoughtsTokenCount')]]
    tool_use_prompt_token_count: NotRequired[Annotated[int, pydantic.Field(alias='toolUsePromptTokenCount')]]
    prompt_tokens_details: NotRequired[
        Annotated[list[_GeminiModalityTokenCount], pydantic.Field(alias='promptTokensDetails')]
    ]
    cache_tokens_details: NotRequired[
        Annotated[list[_GeminiModalityTokenCount], pydantic.Field(alias='cacheTokensDetails')]
    ]
    candidates_tokens_details: NotRequired[
        Annotated[list[_GeminiModalityTokenCount], pydantic.Field(alias='candidatesTokensDetails')]
    ]
    tool_use_prompt_tokens_details: NotRequired[
        Annotated[list[_GeminiModalityTokenCount], pydantic.Field(alias='toolUsePromptTokensDetails')]
    ]


def _metadata_as_usage(response: _GeminiResponse) -> usage.RequestUsage:
    metadata = response.get('usage_metadata')
    if metadata is None:
        return usage.RequestUsage()
    details: dict[str, int] = {}
    if cached_content_token_count := metadata.get('cached_content_token_count', 0):
        details['cached_content_tokens'] = cached_content_token_count

    if thoughts_token_count := metadata.get('thoughts_token_count', 0):
        details['thoughts_tokens'] = thoughts_token_count

    if tool_use_prompt_token_count := metadata.get('tool_use_prompt_token_count', 0):
        details['tool_use_prompt_tokens'] = tool_use_prompt_token_count

    input_audio_tokens = 0
    output_audio_tokens = 0
    cache_audio_read_tokens = 0
    for key, metadata_details in metadata.items():
        if key.endswith('_details') and metadata_details:
            metadata_details = cast(list[_GeminiModalityTokenCount], metadata_details)
            suffix = key.removesuffix('_details')
            for detail in metadata_details:
                modality = detail['modality']
                details[f'{modality.lower()}_{suffix}'] = value = detail.get('token_count', 0)
                if value and modality == 'AUDIO':
                    if key == 'prompt_tokens_details':
                        input_audio_tokens = value
                    elif key == 'candidates_tokens_details':
                        output_audio_tokens = value
                    elif key == 'cache_tokens_details':  # pragma: no branch
                        cache_audio_read_tokens = value

    return usage.RequestUsage(
        input_tokens=metadata.get('prompt_token_count', 0),
        output_tokens=metadata.get('candidates_token_count', 0) + thoughts_token_count,
        cache_read_tokens=cached_content_token_count,
        input_audio_tokens=input_audio_tokens,
        output_audio_tokens=output_audio_tokens,
        cache_audio_read_tokens=cache_audio_read_tokens,
        details=details,
    )


class _GeminiSafetyRating(TypedDict):
    """See <https://ai.google.dev/gemini-api/docs/safety-settings#safety-filters>."""

    category: Literal[
        'HARM_CATEGORY_HARASSMENT',
        'HARM_CATEGORY_HATE_SPEECH',
        'HARM_CATEGORY_SEXUALLY_EXPLICIT',
        'HARM_CATEGORY_DANGEROUS_CONTENT',
        'HARM_CATEGORY_CIVIC_INTEGRITY',
    ]
    probability: Literal['NEGLIGIBLE', 'LOW', 'MEDIUM', 'HIGH']
    blocked: NotRequired[bool]


class _GeminiPromptFeedback(TypedDict):
    """See <https://ai.google.dev/api/generate-content#v1beta.GenerateContentResponse>."""

    block_reason: Annotated[str, pydantic.Field(alias='blockReason')]
    safety_ratings: Annotated[list[_GeminiSafetyRating], pydantic.Field(alias='safetyRatings')]


_gemini_request_ta = pydantic.TypeAdapter(_GeminiRequest)
_gemini_response_ta = pydantic.TypeAdapter(_GeminiResponse)

# steam requests return a list of https://ai.google.dev/api/generate-content#method:-models.streamgeneratecontent
_gemini_streamed_response_ta = pydantic.TypeAdapter(list[_GeminiResponse], config=pydantic.ConfigDict(defer_build=True))


def _ensure_decodeable(content: bytearray) -> bytearray:
    """Trim any invalid unicode point bytes off the end of a bytearray.

    This is necessary before attempting to parse streaming JSON bytes.

    This is a temporary workaround until https://github.com/pydantic/pydantic-core/issues/1633 is resolved
    """
    try:
        content.decode()
    except UnicodeDecodeError as e:
        # e.start marks the start of the invalid decoded bytes, so cut up to before the first invalid byte
        return content[: e.start]
    else:
        return content
