from __future__ import annotations as _annotations

from collections.abc import AsyncIterable, AsyncIterator, Iterator, Sequence
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Literal, cast

import pydantic_core
from httpx import Timeout
from typing_extensions import assert_never

from .. import ModelHTTPError, UnexpectedModelBehavior, _utils
from .._run_context import RunContext
from .._utils import generate_tool_call_id as _generate_tool_call_id, now_utc as _now_utc, number_to_datetime
from ..exceptions import ModelAPIError
from ..messages import (
    AudioUrl,
    BinaryContent,
    CachePoint,
    CompactionPart,
    DocumentUrl,
    FilePart,
    FinishReason,
    ImageUrl,
    ModelMessage,
    ModelRequest,
    ModelResponse,
    ModelResponsePart,
    ModelResponseStreamEvent,
    NativeToolCallPart,
    NativeToolReturnPart,
    RetryPromptPart,
    SystemPromptPart,
    TextContent,
    TextPart,
    ThinkingPart,
    ToolCallPart,
    ToolReturnPart,
    UploadedFile,
    UserContent,
    UserPromptPart,
    VideoUrl,
)
from ..profiles import ModelProfileSpec
from ..providers import Provider, infer_provider
from ..settings import ModelSettings
from ..tools import ToolDefinition
from ..usage import RequestUsage
from . import (
    Model,
    ModelRequestParameters,
    StreamedResponse,
    check_allow_model_requests,
    download_item,
    get_user_agent,
)
from ._tool_choice import resolve_tool_choice

try:
    from mistralai.client import Mistral
    from mistralai.client.errors import SDKError
    from mistralai.client.models import (
        AudioChunk as MistralAudioChunk,
        ChatCompletionChoiceFinishReason as MistralFinishReason,
        ChatCompletionRequestMessage as MistralMessages,
        ChatCompletionResponse as MistralChatCompletionResponse,
        CompletionChunk as MistralCompletionChunk,
        CompletionEvent as MistralCompletionEvent,
        ContentChunk as MistralContentChunk,
        DocumentURLChunk as MistralDocumentURLChunk,
        FileChunk as MistralFileChunk,
        FunctionCall as MistralFunctionCall,
        ImageURL as MistralImageURL,
        ImageURLChunk as MistralImageURLChunk,
        ReferenceChunk as MistralReferenceChunk,
        TextChunk as MistralTextChunk,
        ThinkChunk as MistralThinkChunk,
        Tool as MistralTool,
        ToolCall as MistralToolCall,
        ToolChoiceEnum as MistralToolChoiceEnum,
        UnknownContentChunk as MistralUnknownContentChunk,
    )
    from mistralai.client.models.assistantmessage import (
        AssistantMessage as MistralAssistantMessage,
        AssistantMessageContent as MistralContent,
    )
    from mistralai.client.models.function import Function as MistralFunction
    from mistralai.client.models.systemmessage import SystemMessage as MistralSystemMessage
    from mistralai.client.models.thinkchunk import Thinking as MistralThinking
    from mistralai.client.models.toolmessage import ToolMessage as MistralToolMessage
    from mistralai.client.models.usermessage import UserMessage as MistralUserMessage
    from mistralai.client.types import UNSET, OptionalNullable as MistralOptionalNullable
    from mistralai.client.types.basemodel import Unset as MistralUnset
    from mistralai.client.utils.eventstreaming import EventStreamAsync as MistralEventStreamAsync
except ImportError as e:  # pragma: lax no cover
    raise ImportError(
        'Please install `mistral` to use the Mistral model, '
        'you can use the `mistral` optional group — `pip install "pydantic-ai-slim[mistral]"`'
    ) from e


@contextmanager
def _map_api_errors(model_name: str) -> Iterator[None]:
    try:
        yield
    except SDKError as e:
        if (status_code := e.status_code) >= 400:
            raise ModelHTTPError(status_code=status_code, model_name=model_name, body=e.body) from e
        raise ModelAPIError(model_name=model_name, message=e.message) from e  # pragma: lax no cover


LatestMistralModelNames = Literal[
    'mistral-large-latest', 'mistral-small-latest', 'codestral-latest', 'mistral-moderation-latest'
]
"""Latest  Mistral models."""

MistralModelName = str | LatestMistralModelNames
"""Possible Mistral model names.

Since Mistral supports a variety of date-stamped models, we explicitly list the most popular models but
allow any name in the type hints.
Since [the Mistral docs](https://docs.mistral.ai/getting-started/models/models_overview/) for a full list.
"""

_FINISH_REASON_MAP: dict[MistralFinishReason, FinishReason] = {
    'stop': 'stop',
    'length': 'length',
    'model_length': 'length',
    'error': 'error',
    'tool_calls': 'tool_call',
}


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

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

    # This class is a placeholder for any future mistral-specific settings


@dataclass(init=False)
class MistralModel(Model[Mistral]):
    """A model that uses Mistral.

    Internally, this uses the [Mistral Python client](https://github.com/mistralai/client-python) to interact with the API.

    [API Documentation](https://docs.mistral.ai/)
    """

    json_mode_schema_prompt: str

    _model_name: MistralModelName = field(repr=False)
    _provider: Provider[Mistral] = field(repr=False)

    def __init__(
        self,
        model_name: MistralModelName,
        *,
        provider: Literal['mistral'] | Provider[Mistral] = 'mistral',
        profile: ModelProfileSpec | None = None,
        json_mode_schema_prompt: str = """Answer in JSON Object, respect the format:\n```\n{schema}\n```\n""",
        settings: ModelSettings | None = None,
    ):
        """Initialize a Mistral 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
                'mistral' or an instance of `Provider[Mistral]`. 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.
            json_mode_schema_prompt: The prompt to show when the model expects a JSON object as input.
            settings: Model-specific settings that will be used as defaults for this model.
        """
        self._model_name = model_name
        self.json_mode_schema_prompt = json_mode_schema_prompt

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

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

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

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

    @property
    def model_name(self) -> MistralModelName:
        """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:
        """Make a non-streaming request to the model from Pydantic AI call."""
        check_allow_model_requests()
        model_settings, model_request_parameters = self.prepare_request(
            model_settings,
            model_request_parameters,
        )
        response = await self._completions_create(
            messages, cast(MistralModelSettings, model_settings or {}), model_request_parameters
        )
        model_response = self._process_response(response)
        return model_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]:
        """Make a streaming request to the model from Pydantic AI call."""
        check_allow_model_requests()
        model_settings, model_request_parameters = self.prepare_request(
            model_settings,
            model_request_parameters,
        )
        with _map_api_errors(self.model_name):
            response = await self._stream_completions_create(
                messages, cast(MistralModelSettings, model_settings or {}), model_request_parameters
            )
        async with response:
            yield await self._process_streamed_response(response, model_request_parameters)

    async def _completions_create(
        self,
        messages: list[ModelMessage],
        model_settings: MistralModelSettings,
        model_request_parameters: ModelRequestParameters,
    ) -> MistralChatCompletionResponse:
        """Make a non-streaming request to the model."""
        # TODO(Marcelo): We need to replace the current MistralAI client to use the beta client.
        # See https://docs.mistral.ai/agents/connectors/websearch/ to support web search.
        tools, tool_choice = self._get_tool_choice(model_request_parameters, model_settings)

        with _map_api_errors(self.model_name):
            response = await self.client.chat.complete_async(
                model=str(self._model_name),
                messages=await self._map_messages(messages, model_request_parameters),
                n=1,
                tools=tools or UNSET,
                tool_choice=tool_choice,
                stream=False,
                max_tokens=model_settings.get('max_tokens', UNSET),
                temperature=model_settings.get('temperature', UNSET),
                top_p=model_settings.get('top_p', 1),
                timeout_ms=self._get_timeout_ms(model_settings.get('timeout')),
                random_seed=model_settings.get('seed', UNSET),
                stop=model_settings.get('stop_sequences', None),
                http_headers={'User-Agent': get_user_agent()},
            )

        assert response, 'An unexpected empty response from Mistral.'
        return response

    async def _stream_completions_create(
        self,
        messages: list[ModelMessage],
        model_settings: MistralModelSettings,
        model_request_parameters: ModelRequestParameters,
    ) -> MistralEventStreamAsync[MistralCompletionEvent]:
        """Create a streaming completion request to the Mistral model."""
        response: MistralEventStreamAsync[MistralCompletionEvent] | None
        mistral_messages = await self._map_messages(messages, model_request_parameters)

        # TODO(Marcelo): We need to replace the current MistralAI client to use the beta client.
        # See https://docs.mistral.ai/agents/connectors/websearch/ to support web search.
        tools, tool_choice = self._get_tool_choice(model_request_parameters, model_settings)

        if tools:
            # Function Calling mode (with filtered tools)
            response = await self.client.chat.stream_async(
                model=str(self._model_name),
                messages=mistral_messages,
                n=1,
                tools=tools,
                tool_choice=tool_choice,
                temperature=model_settings.get('temperature', UNSET),
                top_p=model_settings.get('top_p', 1),
                max_tokens=model_settings.get('max_tokens', UNSET),
                timeout_ms=self._get_timeout_ms(model_settings.get('timeout')),
                presence_penalty=model_settings.get('presence_penalty'),
                frequency_penalty=model_settings.get('frequency_penalty'),
                stop=model_settings.get('stop_sequences', None),
                http_headers={'User-Agent': get_user_agent()},
            )

        elif model_request_parameters.output_tools:  # pragma: no cover
            # this branch is dead code (output tool is being handled above)
            # leaving it in for the TODO (support NativeOutput properly)
            # TODO: Port to native "manual JSON" mode
            # Json Mode (only output tools, no function tools filtered in)
            parameters_json_schemas = [tool.parameters_json_schema for tool in model_request_parameters.output_tools]
            user_output_format_message = self._generate_user_output_format(parameters_json_schemas)
            mistral_messages.append(user_output_format_message)

            response = await self.client.chat.stream_async(
                model=str(self._model_name),
                messages=mistral_messages,
                response_format={
                    'type': 'json_object'
                },  # TODO: Should be able to use json_schema now: https://docs.mistral.ai/capabilities/structured-output/custom_structured_output/, https://github.com/mistralai/client-python/blob/bc4adf335968c8a272e1ab7da8461c9943d8e701/src/mistralai/extra/utils/response_format.py#L9
                stream=True,
                temperature=model_settings.get('temperature', UNSET),
                top_p=model_settings.get('top_p', 1),
                max_tokens=model_settings.get('max_tokens', UNSET),
                timeout_ms=self._get_timeout_ms(model_settings.get('timeout')),
                presence_penalty=model_settings.get('presence_penalty'),
                frequency_penalty=model_settings.get('frequency_penalty'),
                stop=model_settings.get('stop_sequences', None),
                http_headers={'User-Agent': get_user_agent()},
            )

        else:
            # Stream Mode (no tools at all)
            response = await self.client.chat.stream_async(
                model=str(self._model_name),
                messages=mistral_messages,
                stream=True,
                http_headers={'User-Agent': get_user_agent()},
            )
        assert response, 'An unexpected empty response from Mistral.'
        return response

    def _get_tool_choice(
        self,
        model_request_parameters: ModelRequestParameters,
        model_settings: MistralModelSettings,
    ) -> tuple[list[MistralTool] | None, MistralToolChoiceEnum | None]:
        """Get tools and tool choice for the model.

        Returns a tuple of (tools, tool_choice):
        - tools: List of MistralTool definitions to send, or None if no tools
        - tool_choice: "auto", "any", "none", "required", or None

        Tool choice semantics:
        - "auto": Default mode. Model decides if it uses the tool or not.
        - "any": Select any tool.
        - "none": Prevents tool use.
        - "required": Forces tool use.
        """
        resolved_tool_choice = resolve_tool_choice(model_settings, model_request_parameters)
        tool_defs = model_request_parameters.tool_defs

        tool_choice: MistralToolChoiceEnum
        if resolved_tool_choice == 'auto':
            tool_choice = 'auto'
        elif resolved_tool_choice == 'required':
            tool_choice = 'any'
        elif resolved_tool_choice == 'none':
            # Mistral returns garbled responses when tool_choice='none' with tools present.
            # Don't send tools at all.
            return None, None
        elif isinstance(resolved_tool_choice, tuple):
            tool_choice_mode, tool_names = resolved_tool_choice
            # Breaks caching, but Mistral doesn't support limiting tools via API arg
            tool_defs = {k: v for k, v in tool_defs.items() if k in tool_names}
            tool_choice = 'auto' if tool_choice_mode == 'auto' else 'any'
        else:
            assert_never(resolved_tool_choice)

        if not tool_defs:
            return None, None

        _tool_functions = [
            MistralFunction(name=r.name, parameters=r.parameters_json_schema, description=r.description or '')
            for r in tool_defs.values()
        ]
        tools = [MistralTool(function=f) for f in _tool_functions]

        return tools, tool_choice

    def _process_response(self, response: MistralChatCompletionResponse) -> ModelResponse:
        """Process a non-streamed response, and prepare a message to return."""
        assert response.choices, 'Unexpected empty response choice.'

        choice = response.choices[0]
        content = choice.message.content
        tool_calls = choice.message.tool_calls

        parts: list[ModelResponsePart] = []
        text, thinking = _map_content(content)
        for thought in thinking:
            parts.append(ThinkingPart(content=thought))
        if text:
            parts.append(TextPart(content=text))

        if isinstance(tool_calls, list):
            for tool_call in tool_calls:
                tool = self._map_mistral_to_pydantic_tool_call(tool_call=tool_call)
                parts.append(tool)

        raw_finish_reason = choice.finish_reason
        provider_details: dict[str, Any] = {'finish_reason': raw_finish_reason}
        if response.created:  # pragma: no branch
            provider_details['timestamp'] = number_to_datetime(response.created)
        finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)

        return ModelResponse(
            parts=parts,
            usage=_map_usage(response),
            model_name=response.model,
            provider_response_id=response.id,
            provider_name=self._provider.name,
            provider_url=self._provider.base_url,
            finish_reason=finish_reason,
            provider_details=provider_details,
        )

    async def _process_streamed_response(
        self,
        response: MistralEventStreamAsync[MistralCompletionEvent],
        model_request_parameters: ModelRequestParameters,
    ) -> StreamedResponse:
        """Process a streamed response, and prepare a streaming response to return."""
        peekable_response: _utils.PeekableAsyncStream[
            MistralCompletionEvent, MistralEventStreamAsync[MistralCompletionEvent]
        ] = _utils.PeekableAsyncStream(response)
        with _map_api_errors(self.model_name):
            first_chunk = await peekable_response.peek()
        if isinstance(first_chunk, _utils.Unset):
            raise UnexpectedModelBehavior(  # pragma: no cover
                'Streamed response ended without content or tool calls'
            )

        return MistralStreamedResponse(
            model_request_parameters=model_request_parameters,
            _response=peekable_response,
            _model_name=first_chunk.data.model,
            _provider_name=self._provider.name,
            _provider_url=self._provider.base_url,
            _provider_timestamp=number_to_datetime(first_chunk.data.created) if first_chunk.data.created else None,
        )

    @staticmethod
    def _map_mistral_to_pydantic_tool_call(tool_call: MistralToolCall) -> ToolCallPart:
        """Maps a MistralToolCall to a ToolCall."""
        tool_call_id = tool_call.id or _generate_tool_call_id()
        func_call = tool_call.function

        return ToolCallPart(func_call.name, func_call.arguments, tool_call_id)

    @staticmethod
    def _map_tool_call(t: ToolCallPart) -> MistralToolCall:
        """Maps a pydantic-ai ToolCall to a MistralToolCall."""
        return MistralToolCall(
            id=_utils.guard_tool_call_id(t=t),
            type='function',
            function=MistralFunctionCall(name=t.tool_name, arguments=t.args or {}),
        )

    def _generate_user_output_format(self, schemas: list[dict[str, Any]]) -> MistralUserMessage:
        """Get a message with an example of the expected output format."""
        examples: list[dict[str, Any]] = []
        for schema in schemas:
            typed_dict_definition: dict[str, Any] = {}
            for key, value in schema.get('properties', {}).items():
                typed_dict_definition[key] = self._get_python_type(value)
            examples.append(typed_dict_definition)

        example_schema = examples[0] if len(examples) == 1 else examples
        return MistralUserMessage(content=self.json_mode_schema_prompt.format(schema=example_schema))

    @classmethod
    def _get_python_type(cls, value: dict[str, Any]) -> str:
        """Return a string representation of the Python type for a single JSON schema property.

        This function handles recursion for nested arrays/objects and `anyOf`.
        """
        # 1) Handle anyOf first, because it's a different schema structure
        if any_of := value.get('anyOf'):
            # Simplistic approach: pick the first option in anyOf
            # (In reality, you'd possibly want to merge or union types)
            return f'Optional[{cls._get_python_type(any_of[0])}]'

        # 2) If we have a top-level "type" field
        value_type = value.get('type')
        if not value_type:
            # No explicit type; fallback
            return 'Any'

        # 3) Direct simple type mapping (string, integer, float, bool, None)
        if value_type in SIMPLE_JSON_TYPE_MAPPING and value_type != 'array' and value_type != 'object':
            return SIMPLE_JSON_TYPE_MAPPING[value_type]

        # 4) Array: Recursively get the item type
        if value_type == 'array':
            items = value.get('items', {})
            return f'list[{cls._get_python_type(items)}]'

        # 5) Object: Check for additionalProperties
        if value_type == 'object':
            additional_properties = value.get('additionalProperties', {})
            if isinstance(additional_properties, bool):
                return 'bool'  # pragma: lax no cover
            additional_properties_type = additional_properties.get('type')
            if (
                additional_properties_type in SIMPLE_JSON_TYPE_MAPPING
                and additional_properties_type != 'array'
                and additional_properties_type != 'object'
            ):
                # dict[str, bool/int/float/etc...]
                return f'dict[str, {SIMPLE_JSON_TYPE_MAPPING[additional_properties_type]}]'
            elif additional_properties_type == 'array':
                array_items = additional_properties.get('items', {})
                return f'dict[str, list[{cls._get_python_type(array_items)}]]'
            elif additional_properties_type == 'object':
                # nested dictionary of unknown shape
                return 'dict[str, dict[str, Any]]'
            else:
                # If no additionalProperties type or something else, default to a generic dict
                return 'dict[str, Any]'

        # 6) Fallback
        return 'Any'

    @staticmethod
    def _get_timeout_ms(timeout: Timeout | float | None) -> int | None:
        """Convert a timeout to milliseconds."""
        if timeout is None:
            return None
        if isinstance(timeout, float):  # pragma: no cover
            return int(1000 * timeout)
        raise NotImplementedError('Timeout object is not yet supported for MistralModel.')

    async def _map_user_message(self, message: ModelRequest) -> AsyncIterable[MistralMessages]:
        file_content: list[UserContent] = []
        for part in message.parts:
            if isinstance(part, SystemPromptPart):
                yield MistralSystemMessage(content=part.content)
            elif isinstance(part, UserPromptPart):
                yield await self._map_user_prompt(part)
            elif isinstance(part, ToolReturnPart):
                tool_text, files = part.model_response_str_and_user_content()
                file_content.extend(files)
                yield MistralToolMessage(
                    tool_call_id=part.tool_call_id,
                    content=tool_text,
                )
            elif isinstance(part, RetryPromptPart):
                if part.tool_name is None:
                    yield MistralUserMessage(content=part.model_response())  # pragma: no cover
                else:
                    yield MistralToolMessage(
                        tool_call_id=part.tool_call_id,
                        content=part.model_response(),
                    )
            else:
                assert_never(part)
        if file_content:
            yield await self._map_user_prompt(UserPromptPart(content=file_content))

    async def _map_messages(  # noqa: C901
        self, messages: Sequence[ModelMessage], model_request_parameters: ModelRequestParameters
    ) -> list[MistralMessages]:
        """Just maps a `pydantic_ai.Message` to a `MistralMessage`."""
        mistral_messages: list[MistralMessages] = []
        for message in messages:
            if isinstance(message, ModelRequest):
                async for msg in self._map_user_message(message):
                    mistral_messages.append(msg)
            elif isinstance(message, ModelResponse):
                content_chunks: list[MistralContentChunk] = []
                thinking_chunks: list[MistralThinking] = []
                tool_calls: list[MistralToolCall] = []

                for part in message.parts:
                    if isinstance(part, TextPart):
                        content_chunks.append(MistralTextChunk(text=part.content))
                    elif isinstance(part, ThinkingPart):
                        thinking_chunks.append(MistralTextChunk(text=part.content))
                    elif isinstance(part, ToolCallPart):
                        tool_calls.append(self._map_tool_call(part))
                    elif isinstance(part, NativeToolCallPart | NativeToolReturnPart):  # pragma: no cover
                        # This is currently never returned from mistral
                        pass
                    elif isinstance(part, FilePart):  # pragma: no cover
                        # Files generated by models are not sent back to models that don't themselves generate files.
                        pass
                    elif isinstance(part, CompactionPart):  # pragma: no cover
                        # Compaction parts are not sent back to models that don't support compaction.
                        pass
                    else:
                        assert_never(part)
                if thinking_chunks:
                    content_chunks.insert(0, MistralThinkChunk(thinking=thinking_chunks))
                mistral_messages.append(MistralAssistantMessage(content=content_chunks, tool_calls=tool_calls))
            else:
                assert_never(message)
        if instruction_parts := self._get_instruction_parts(messages, model_request_parameters):
            system_prompt_count = next(
                (i for i, m in enumerate(mistral_messages) if not isinstance(m, MistralSystemMessage)),
                len(mistral_messages),
            )
            mistral_messages[system_prompt_count:system_prompt_count] = [
                MistralSystemMessage(content=part.content) for part in instruction_parts
            ]

        # Post-process messages to insert fake assistant message after tool message if followed by user message
        # to work around `Unexpected role 'user' after role 'tool'` error.
        processed_messages: list[MistralMessages] = []
        for i, current_message in enumerate(mistral_messages):
            processed_messages.append(current_message)

            if isinstance(current_message, MistralToolMessage) and i + 1 < len(mistral_messages):
                next_message = mistral_messages[i + 1]
                if isinstance(next_message, MistralUserMessage):
                    # Insert a dummy assistant message
                    processed_messages.append(MistralAssistantMessage(content=[MistralTextChunk(text='OK')]))

        return processed_messages

    async def _map_user_prompt(self, part: UserPromptPart) -> MistralUserMessage:  # noqa: C901
        content: str | list[MistralContentChunk]
        if isinstance(part.content, str):
            content = part.content
        else:
            content = []
            for item in part.content:
                if isinstance(item, str | TextContent):
                    text = item if isinstance(item, str) else item.content
                    content.append(MistralTextChunk(text=text))
                elif isinstance(item, ImageUrl):
                    if item.force_download:
                        downloaded = await download_item(item, data_format='base64_uri')
                        image_url = MistralImageURL(url=downloaded['data'])
                        content.append(MistralImageURLChunk(image_url=image_url, type='image_url'))
                    else:
                        content.append(MistralImageURLChunk(image_url=MistralImageURL(url=item.url)))
                elif isinstance(item, BinaryContent):
                    if item.is_image:
                        image_url = MistralImageURL(url=item.data_uri)
                        content.append(MistralImageURLChunk(image_url=image_url, type='image_url'))
                    elif item.media_type == 'application/pdf':
                        content.append(MistralDocumentURLChunk(document_url=item.data_uri, type='document_url'))
                    else:
                        raise NotImplementedError(
                            'BinaryContent other than image or PDF is not supported in Mistral user prompts'
                        )
                elif isinstance(item, DocumentUrl):
                    if item.media_type == 'application/pdf':
                        if item.force_download:
                            downloaded = await download_item(item, data_format='base64_uri')
                            content.append(
                                MistralDocumentURLChunk(document_url=downloaded['data'], type='document_url')
                            )
                        else:
                            content.append(MistralDocumentURLChunk(document_url=item.url, type='document_url'))
                    else:
                        raise NotImplementedError('DocumentUrl other than PDF is not supported in Mistral user prompts')
                elif isinstance(item, AudioUrl):
                    raise NotImplementedError('AudioUrl is not supported in Mistral user prompts')
                elif isinstance(item, VideoUrl):
                    raise NotImplementedError('VideoUrl is not supported in Mistral user prompts')
                elif isinstance(item, UploadedFile):
                    raise NotImplementedError('UploadedFile is not supported in Mistral user prompts')
                elif isinstance(item, CachePoint):
                    pass
                else:
                    assert_never(item)
        return MistralUserMessage(content=content)


MistralToolCallId = str | None


@dataclass
class MistralStreamedResponse(StreamedResponse):
    """Implementation of `StreamedResponse` for Mistral models."""

    _model_name: MistralModelName
    _response: _utils.PeekableAsyncStream[MistralCompletionEvent, MistralEventStreamAsync[MistralCompletionEvent]]
    _provider_name: str
    _provider_url: str
    _provider_timestamp: datetime | None = None
    _timestamp: datetime = field(default_factory=_now_utc)

    _delta_content: str = field(default='', init=False)

    async def close_stream(self) -> None:
        await self._response.source.response.aclose()

    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
        with _map_api_errors(self._model_name):
            if self._provider_timestamp is not None:  # pragma: no branch
                self.provider_details = {'timestamp': self._provider_timestamp}
            chunk: MistralCompletionEvent
            async for chunk in self._response:
                self._usage += _map_usage(chunk.data)

                if chunk.data.id:  # pragma: no branch
                    self.provider_response_id = chunk.data.id

                try:
                    choice = chunk.data.choices[0]
                except IndexError:
                    continue

                if raw_finish_reason := choice.finish_reason:
                    self.provider_details = {**(self.provider_details or {}), 'finish_reason': raw_finish_reason}
                    self.finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)

                # Handle the text part of the response
                content = choice.delta.content
                text, thinking = _map_content(content)
                for thought in thinking:
                    for event in self._parts_manager.handle_thinking_delta(vendor_part_id='thinking', content=thought):
                        yield event
                if text:
                    # Attempt to produce an output tool call from the received text
                    output_tools = {c.name: c for c in self.model_request_parameters.output_tools}
                    if output_tools:
                        self._delta_content += text
                        # TODO: Port to native "manual JSON" mode
                        maybe_tool_call_part = self._try_get_output_tool_from_text(self._delta_content, output_tools)
                        if maybe_tool_call_part:
                            yield self._parts_manager.handle_tool_call_part(
                                vendor_part_id='output',
                                tool_name=maybe_tool_call_part.tool_name,
                                args=maybe_tool_call_part.args_as_dict(),
                                tool_call_id=maybe_tool_call_part.tool_call_id,
                            )
                    else:
                        for event in self._parts_manager.handle_text_delta(vendor_part_id='content', content=text):
                            yield event

                # Handle the explicit tool calls
                for index, dtc in enumerate(choice.delta.tool_calls or []):
                    # It seems that mistral just sends full tool calls, so we just use them directly, rather than building
                    yield self._parts_manager.handle_tool_call_part(
                        vendor_part_id=index,
                        tool_name=dtc.function.name,
                        args=dtc.function.arguments,
                        tool_call_id=dtc.id,
                    )

    @property
    def model_name(self) -> MistralModelName:
        """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

    @staticmethod
    def _try_get_output_tool_from_text(text: str, output_tools: dict[str, ToolDefinition]) -> ToolCallPart | None:
        output_json: dict[str, Any] | None = pydantic_core.from_json(text, allow_partial='trailing-strings')
        if output_json:
            for output_tool in output_tools.values():
                # NOTE: Additional verification to prevent JSON validation to crash
                # Ensures required parameters in the JSON schema are respected, especially for stream-based return types.
                # Example with BaseModel and required fields.
                if not MistralStreamedResponse._validate_required_json_schema(
                    output_json, output_tool.parameters_json_schema
                ):
                    continue

                # The following part_id will be thrown away
                return ToolCallPart(tool_name=output_tool.name, args=output_json)

    @staticmethod
    def _validate_required_json_schema(json_dict: dict[str, Any], json_schema: dict[str, Any]) -> bool:
        """Validate that all required parameters in the JSON schema are present in the JSON dictionary."""
        required_params = json_schema.get('required', [])
        properties = json_schema.get('properties', {})

        for param in required_params:
            if param not in json_dict:
                return False

            param_schema = properties.get(param, {})
            param_type = param_schema.get('type')
            param_items_type = param_schema.get('items', {}).get('type')

            if param_type == 'array' and param_items_type:
                if not isinstance(json_dict[param], list):
                    return False
                for item in json_dict[param]:
                    if not isinstance(item, VALID_JSON_TYPE_MAPPING[param_items_type]):
                        return False
            elif param_type and not isinstance(json_dict[param], VALID_JSON_TYPE_MAPPING[param_type]):
                return False

            if isinstance(json_dict[param], dict) and 'properties' in param_schema:
                nested_schema = param_schema
                if not MistralStreamedResponse._validate_required_json_schema(json_dict[param], nested_schema):
                    return False

        return True


VALID_JSON_TYPE_MAPPING: dict[str, Any] = {
    'string': str,
    'integer': int,
    'number': float,
    'boolean': bool,
    'array': list,
    'object': dict,
    'null': type(None),
}

SIMPLE_JSON_TYPE_MAPPING = {
    'string': 'str',
    'integer': 'int',
    'number': 'float',
    'boolean': 'bool',
    'array': 'list',
    'null': 'None',
}


def _map_usage(response: MistralChatCompletionResponse | MistralCompletionChunk) -> RequestUsage:
    """Maps a Mistral Completion Chunk or Chat Completion Response to a Usage."""
    if response.usage:
        return RequestUsage(
            input_tokens=response.usage.prompt_tokens or 0,
            output_tokens=response.usage.completion_tokens or 0,
        )
    else:
        return RequestUsage()


def _map_content(content: MistralOptionalNullable[MistralContent]) -> tuple[str | None, list[str]]:
    """Maps the delta content from a Mistral Completion Chunk to a string or None."""
    text: str | None = None
    thinking: list[str] = []

    if isinstance(content, MistralUnset) or not content:
        return None, []
    elif isinstance(content, list):
        for chunk in content:
            if isinstance(chunk, MistralTextChunk):
                text = (text or '') + chunk.text
            elif isinstance(chunk, MistralThinkChunk):
                for thought in chunk.thinking:
                    if thought.type == 'text':  # pragma: no branch
                        thinking.append(thought.text)
            elif isinstance(chunk, MistralReferenceChunk):
                pass
            elif isinstance(
                chunk,
                MistralImageURLChunk | MistralDocumentURLChunk | MistralFileChunk | MistralAudioChunk,
            ):  # pragma: no cover
                pass
            elif isinstance(chunk, MistralUnknownContentChunk):  # pragma: no cover
                pass
            else:
                assert_never(chunk)
    elif isinstance(content, str):
        text = content

    # Note: Check len to handle potential mismatch between function calls and responses from the API. (`msg: not the same number of function class and responses`)
    if text and len(text) == 0:  # pragma: no cover
        text = None

    return text, thinking
