import json
import secrets
from collections.abc import Mapping
from dataclasses import dataclass
from typing import cast

from google import genai
from google.genai import (
    errors,
    types,
)
from mcp import Tool as McpTool
from mcp.types import (
    CallToolRequest,
    CallToolRequestParams,
    ContentBlock,
    TextContent,
)

from fast_agent.core.exceptions import ProviderKeyError
from fast_agent.core.prompt import Prompt
from fast_agent.llm.fastagent_llm import FastAgentLLM
from fast_agent.llm.model_database import ModelDatabase
from fast_agent.llm.provider.google.google_converter import GoogleConverter
from fast_agent.llm.provider_types import Provider
from fast_agent.llm.reasoning_effort import (
    format_reasoning_setting,
    parse_reasoning_setting,
)
from fast_agent.llm.stream_types import StreamChunk
from fast_agent.llm.tool_tracking import ToolCallTracker
from fast_agent.llm.usage_tracking import TurnUsage
from fast_agent.types import PromptMessageExtended, RequestParams
from fast_agent.types.llm_stop_reason import LlmStopReason

# Define default model and potentially other Google-specific defaults
DEFAULT_GOOGLE_MODEL = "gemini3"


# Define Google-specific parameter exclusions if necessary
GOOGLE_EXCLUDE_FIELDS = {
    # Add fields that should not be passed directly from RequestParams to google.genai config
    FastAgentLLM.PARAM_MESSAGES,  # Handled by contents
    FastAgentLLM.PARAM_MODEL,  # Handled during client/call setup
    FastAgentLLM.PARAM_SYSTEM_PROMPT,  # Handled by system_instruction in config
    FastAgentLLM.PARAM_USE_HISTORY,  # Handled by FastAgentLLM base / this class's logic
    FastAgentLLM.PARAM_MAX_ITERATIONS,  # Handled by this class's loop
    FastAgentLLM.PARAM_MCP_METADATA,
}.union(FastAgentLLM.BASE_EXCLUDE_FIELDS)


@dataclass(slots=True)
class _GoogleTextTimelineEntry:
    text: str


@dataclass(slots=True)
class _GoogleToolTimelineEntry:
    tool_use_id: str


@dataclass(slots=True)
class _GoogleToolBuffer:
    tool_use_id: str
    name: str
    buffer: str = ""


GoogleTimelineEntry = _GoogleTextTimelineEntry | _GoogleToolTimelineEntry


class GoogleNativeLLM(FastAgentLLM[types.Content, types.Content]):
    """
    Google LLM provider using the native google.genai library.
    """

    def __init__(self, **kwargs) -> None:
        kwargs.pop("provider", None)
        super().__init__(provider=Provider.GOOGLE, **kwargs)
        # Initialize the converter
        self._converter = GoogleConverter()
        self._init_reasoning(kwargs)

    def _init_reasoning(self, kwargs: dict) -> None:
        """Wire up reasoning/thinking from kwargs or config."""
        raw_setting = kwargs.get("reasoning_effort", None)
        model_name = self.default_request_params.model or DEFAULT_GOOGLE_MODEL

        if raw_setting is None:
            google_cfg = getattr(getattr(self.context, "config", None), "google", None)
            if google_cfg:
                raw_setting = (
                    google_cfg.get("reasoning")
                    if isinstance(google_cfg, Mapping)
                    else getattr(google_cfg, "reasoning", None)
                )

        reasoning_mode = self._get_model_reasoning(model_name)
        spec = self._get_model_reasoning_effort_spec(model_name)

        if raw_setting is not None and reasoning_mode != "google_thinking":
            self.logger.warning(
                "Reasoning setting ignored for model without Google thinking support."
            )
            raw_setting = None

        if raw_setting is None and reasoning_mode == "google_thinking" and spec and spec.default:
            raw_setting = spec.default

        setting = parse_reasoning_setting(raw_setting)
        if setting is not None:
            try:
                self.set_reasoning_effort(setting)
            except ValueError as exc:
                self.logger.warning(f"Invalid reasoning setting: {exc}")
                if spec and spec.default:
                    self.set_reasoning_effort(spec.default)
                else:
                    self.set_reasoning_effort(None)
        else:
            self.set_reasoning_effort(None)

        if reasoning_mode == "google_thinking":
            self.logger.info(
                f"Google reasoning resolved: {format_reasoning_setting(self.reasoning_effort)}"
            )

    def _resolve_thinking_config(self) -> tuple[int | None, str | None]:
        """Resolve thinking config from reasoning_effort setting.

        Returns:
            (thinking_budget, thinking_level) tuple where:
            - thinking_budget: None if not configured, 0 to disable, -1 for auto,
              or a positive token count for explicit budgets.
            - thinking_level: SDK ThinkingLevel name (MINIMAL/LOW/MEDIUM/HIGH)
              when an effort level is selected, None otherwise.
        """
        setting = self.reasoning_effort
        if setting is None:
            return (None, None)
        if setting.kind == "toggle":
            return (-1 if setting.value else 0, None)
        if setting.kind == "budget" and isinstance(setting.value, int):
            return (max(0, setting.value), None)
        if setting.kind == "effort":
            effort = str(setting.value).lower()
            if effort in ("none",):
                return (0, None)
            if effort in ("auto",):
                return (-1, None)
            # Map to SDK ThinkingLevel names
            level_map: dict[str, str] = {
                "minimal": "MINIMAL",
                "low": "LOW",
                "medium": "MEDIUM",
                "high": "HIGH",
            }
            level = level_map.get(effort)
            if level:
                return (None, level)
            return (-1, None)
        return (None, None)

    def _vertex_cfg(self) -> tuple[bool, str | None, str | None]:
        """(enabled, project_id, location) for Vertex config; supports dict/mapping or object."""
        google_cfg = getattr(getattr(self.context, "config", None), "google", None)
        vertex = (
            (google_cfg or {}).get("vertex_ai")
            if isinstance(google_cfg, Mapping)
            else getattr(google_cfg, "vertex_ai", None)
        )
        if not vertex:
            return (False, None, None)
        if isinstance(vertex, Mapping):
            return (bool(vertex.get("enabled")), vertex.get("project_id"), vertex.get("location"))
        return (
            bool(getattr(vertex, "enabled", False)),
            getattr(vertex, "project_id", None),
            getattr(vertex, "location", None),
        )

    def _resolve_model_name(self, model: str) -> str:
        """Resolve model name; for Vertex, apply a generic preview→base fallback.

        * If the caller passes a full publisher resource name, it is respected as-is.
        * If Vertex is not enabled, the short id is returned unchanged (Developer API path).
        * If Vertex is enabled and the id contains '-preview-', the suffix is stripped so that
          e.g. 'gemini-2.5-flash-preview-09-2025' becomes 'gemini-2.5-flash'.
        """
        # Fully-qualified publisher / model resource: do not rewrite.
        if model.startswith(("projects/", "publishers/")) or "/publishers/" in model:
            return model

        enabled, project_id, location = self._vertex_cfg()
        # Developer API path: return the short model id unchanged.
        if not (enabled and project_id and location):
            return model

        # Vertex path: strip any '-preview-…' suffix to fall back to the base model id.
        base_model = model.split("-preview-", 1)[0] if "-preview-" in model else model

        return f"projects/{project_id}/locations/{location}/publishers/google/models/{base_model}"

    def _initialize_google_client(self) -> genai.Client:
        """
        Initializes the google.genai client.

        Reads Google API key or Vertex AI configuration from context config.
        """
        try:
            # Prefer Vertex AI (ADC/IAM) if enabled. This path must NOT require an API key.
            vertex_enabled, project_id, location = self._vertex_cfg()
            if vertex_enabled:
                return genai.Client(
                    vertexai=True,
                    project=project_id,
                    location=location,
                    # http_options=types.HttpOptions(api_version='v1')
                )

            # Otherwise, default to Gemini Developer API (API key required).
            api_key = self._api_key()
            if not api_key:
                raise ProviderKeyError(
                    "Google API key not found.",
                    "Please configure your Google API key.",
                )

            return genai.Client(
                api_key=api_key,
                # http_options=types.HttpOptions(api_version='v1')
            )
        except Exception as e:
            # Catch potential initialization errors and raise ProviderKeyError
            raise ProviderKeyError("Failed to initialize Google GenAI client.", str(e)) from e

    def _initialize_default_params(self, kwargs: dict) -> RequestParams:
        """Initialize Google-specific default parameters."""
        chosen_model = (
            self._resolve_default_model_name(kwargs.get("model"), DEFAULT_GOOGLE_MODEL)
            or DEFAULT_GOOGLE_MODEL
        )
        # Gemini models have different max output token limits; for example,
        # gemini-2.0-flash only supports up to 8192 output tokens.
        resolved_model = self._resolved_model_spec
        if (
            resolved_model is not None
            and chosen_model == resolved_model.wire_model_name
            and resolved_model.max_output_tokens is not None
        ):
            max_tokens = resolved_model.max_output_tokens
        else:
            max_tokens = ModelDatabase.get_max_output_tokens(chosen_model) or 65536

        return RequestParams(
            model=chosen_model,
            systemPrompt=self.instruction,  # System instruction will be mapped in _google_completion
            parallel_tool_calls=True,  # Assume parallel tool calls are supported by default with native API
            max_iterations=20,
            use_history=True,
            # Pick a safe default per model (e.g. gemini-2.0-flash is limited to 8192).
            maxTokens=max_tokens,
            # Include other relevant default parameters
        )

    async def _stream_generate_content(
        self,
        *,
        model: str,
        contents: list[types.Content],
        config: types.GenerateContentConfig,
        client: genai.Client,
    ) -> types.GenerateContentResponse | None:
        """Stream Gemini responses and return the final aggregated completion."""
        try:
            response_stream = await client.aio.models.generate_content_stream(
                model=model,
                contents=cast("types.ContentListUnion", contents),
                config=config,
            )
        except AttributeError:
            # Older SDKs might not expose streaming; fall back to non-streaming.
            return None
        except errors.APIError:
            raise
        except Exception as exc:  # pragma: no cover - defensive fallback
            self.logger.warning(
                "Google streaming failed during setup; falling back to non-streaming",
                exc_info=exc,
            )
            return None

        return await self._consume_google_stream(response_stream, model=model)

    @staticmethod
    def _append_google_text_timeline(
        timeline: list[GoogleTimelineEntry],
        text: str,
    ) -> None:
        if not text:
            return
        if timeline and isinstance(timeline[-1], _GoogleTextTimelineEntry):
            timeline[-1].text += text
            return
        timeline.append(_GoogleTextTimelineEntry(text=text))

    @staticmethod
    def _serialize_google_tool_args(args: object) -> str:
        try:
            return json.dumps(args, separators=(",", ":"))
        except Exception:
            return str(args)

    def _start_google_tool_stream(
        self,
        *,
        tracker: ToolCallTracker,
        tool_buffers: dict[str, _GoogleToolBuffer],
        timeline: list[GoogleTimelineEntry],
        tool_index: int,
        tool_name: str,
    ) -> _GoogleToolBuffer:
        tool_use_id = f"tool_{self.chat_turn()}_{tool_index}"
        state = tracker.register(
            tool_use_id=tool_use_id,
            name=tool_name,
            index=tool_index,
        )
        buffer = _GoogleToolBuffer(tool_use_id=state.tool_use_id, name=state.name)
        tool_buffers[state.tool_use_id] = buffer
        self._notify_tool_stream_listeners(
            "start",
            {
                "tool_name": state.name,
                "tool_use_id": state.tool_use_id,
                "index": tool_index,
            },
        )
        state.start_notified = True
        timeline.append(_GoogleToolTimelineEntry(tool_use_id=state.tool_use_id))
        return buffer

    def _close_google_tool_stream(
        self,
        *,
        tracker: ToolCallTracker,
        tool_index: int,
    ) -> None:
        state = tracker.close(index=tool_index)
        if state is None:
            return
        self._notify_tool_stream_listeners(
            "stop",
            {
                "tool_name": state.name,
                "tool_use_id": state.tool_use_id,
                "index": tool_index,
            },
        )

    def _build_google_final_response(
        self,
        *,
        last_chunk: types.GenerateContentResponse | None,
        usage_metadata: types.GenerateContentResponseUsageMetadata | None,
        timeline: list[GoogleTimelineEntry],
        tool_buffers: dict[str, _GoogleToolBuffer],
    ) -> types.GenerateContentResponse | None:
        if not timeline and last_chunk is None:
            return None

        final_parts: list[types.Part] = []
        for entry in timeline:
            if isinstance(entry, _GoogleTextTimelineEntry):
                final_parts.append(types.Part.from_text(text=entry.text))
                continue

            tool_buffer = tool_buffers.get(entry.tool_use_id)
            if tool_buffer is None:
                continue
            try:
                args_obj = json.loads(tool_buffer.buffer) if tool_buffer.buffer else {}
            except json.JSONDecodeError:
                args_obj = {"__raw": tool_buffer.buffer}
            final_parts.append(
                types.Part.from_function_call(
                    name=str(tool_buffer.name or "tool"),
                    args=args_obj,
                )
            )

        final_content = types.Content(role="model", parts=final_parts)

        if last_chunk is not None:
            final_response = last_chunk.model_copy(deep=True)
            candidates = final_response.candidates or []
            if candidates:
                final_candidate = candidates[0]
                final_candidate.content = final_content
            else:
                final_response.candidates = [types.Candidate(content=final_content)]
        else:
            final_response = types.GenerateContentResponse(
                candidates=[types.Candidate(content=final_content)]
            )

        if usage_metadata:
            final_response.usage_metadata = usage_metadata

        return final_response

    async def _consume_google_stream(
        self,
        response_stream,
        *,
        model: str,
    ) -> types.GenerateContentResponse | None:
        """Consume the async streaming iterator and aggregate the final response."""
        estimated_tokens = 0
        timeline: list[GoogleTimelineEntry] = []
        tracker = ToolCallTracker()
        tool_buffers: dict[str, _GoogleToolBuffer] = {}
        active_tool_index: int | None = None
        tool_counter = 0
        usage_metadata: types.GenerateContentResponseUsageMetadata | None = None
        last_chunk: types.GenerateContentResponse | None = None

        try:
            # Cancellation is handled via asyncio.Task.cancel() which raises CancelledError
            async for chunk in response_stream:
                last_chunk = chunk
                if getattr(chunk, "usage_metadata", None):
                    usage_metadata = chunk.usage_metadata

                if not getattr(chunk, "candidates", None):
                    continue

                candidate = chunk.candidates[0]
                content = getattr(candidate, "content", None)
                if content is None or not getattr(content, "parts", None):
                    continue

                for part in content.parts:
                    if getattr(part, "text", None):
                        text = part.text or ""
                        if text:
                            if getattr(part, "thought", False):
                                self._notify_stream_listeners(
                                    StreamChunk(text=text, is_reasoning=True)
                                )
                            else:
                                self._append_google_text_timeline(timeline, text)
                                estimated_tokens = self._emit_stream_text_delta(
                                    text=text,
                                    model=model,
                                    estimated_tokens=estimated_tokens,
                                )

                    if getattr(part, "function_call", None):
                        function_call = part.function_call
                        name = getattr(function_call, "name", None) or "tool"
                        args = getattr(function_call, "args", None) or {}

                        if active_tool_index is None:
                            active_tool_index = tool_counter
                            tool_counter += 1
                            self._start_google_tool_stream(
                                tracker=tracker,
                                tool_buffers=tool_buffers,
                                timeline=timeline,
                                tool_index=active_tool_index,
                                tool_name=name,
                            )
                        state = tracker.resolve_open(index=active_tool_index)
                        if state is None:
                            continue
                        tool_buffer = tool_buffers.get(state.tool_use_id)
                        if tool_buffer is None:
                            tool_buffer = _GoogleToolBuffer(
                                tool_use_id=state.tool_use_id,
                                name=state.name,
                            )
                            tool_buffers[state.tool_use_id] = tool_buffer

                        serialized_args = self._serialize_google_tool_args(args)
                        previous = tool_buffer.buffer
                        delta = (
                            serialized_args[len(previous) :]
                            if serialized_args.startswith(previous)
                            else serialized_args
                        )
                        tool_buffer.buffer = serialized_args

                        if delta:
                            self._notify_tool_stream_listeners(
                                "delta",
                                {
                                    "tool_name": tool_buffer.name,
                                    "tool_use_id": tool_buffer.tool_use_id,
                                    "index": active_tool_index,
                                    "chunk": delta,
                                },
                            )

                finish_reason = getattr(candidate, "finish_reason", None)
                if finish_reason:
                    finish_value = str(finish_reason).split(".")[-1].upper()
                    if finish_value in {"FUNCTION_CALL", "STOP"} and active_tool_index is not None:
                        self._close_google_tool_stream(
                            tracker=tracker,
                            tool_index=active_tool_index,
                        )
                        active_tool_index = None
        finally:
            stream_close = getattr(response_stream, "aclose", None)
            if callable(stream_close):
                try:
                    await stream_close()
                except Exception:
                    pass

        if active_tool_index is not None:
            self._close_google_tool_stream(
                tracker=tracker,
                tool_index=active_tool_index,
            )

        incomplete_tools = tracker.incomplete()
        if incomplete_tools:
            raise RuntimeError(
                "Streaming completed but tool call(s) never finished: "
                + ", ".join(f"{tool.name}:{tool.tool_use_id}" for tool in incomplete_tools)
            )

        return self._build_google_final_response(
            last_chunk=last_chunk,
            usage_metadata=usage_metadata,
            timeline=timeline,
            tool_buffers=tool_buffers,
        )

    async def _google_completion(
        self,
        message: list[types.Content] | None,
        request_params: RequestParams | None = None,
        tools: list[McpTool] | None = None,
        *,
        response_mime_type: str | None = None,
        response_schema: object | None = None,
    ) -> PromptMessageExtended:
        """
        Process a query using Google's generate_content API and available tools.
        """
        request_params = self.get_request_params(request_params=request_params)
        responses: list[ContentBlock] = []

        # Caller supplies the full set of messages to send (history + turn)
        conversation_history: list[types.Content] = list(message or [])

        self.logger.debug(f"Google completion requested with messages: {conversation_history}")
        self._log_chat_progress(self.chat_turn(), model=request_params.model)

        available_tools: list[types.Tool] = (
            self._converter.convert_to_google_tools(tools or []) if tools else []
        )

        # 2. Prepare generate_content arguments
        thinking_budget, thinking_level = self._resolve_thinking_config()
        generate_content_config = self._converter.convert_request_params_to_google_config(
            request_params,
            thinking_budget=thinking_budget,
            thinking_level=thinking_level,
        )

        # Apply structured output config OR tool calling (mutually exclusive)
        if response_schema or response_mime_type:
            # Structured output mode: disable tool use
            if response_mime_type:
                generate_content_config.response_mime_type = response_mime_type
            if response_schema is not None:
                generate_content_config.response_schema = response_schema
        elif available_tools:
            # Tool calling enabled only when not doing structured output
            generate_content_config.tools = available_tools  # type: ignore[assignment]
            generate_content_config.tool_config = types.ToolConfig(
                function_calling_config=types.FunctionCallingConfig(
                    mode=types.FunctionCallingConfigMode.AUTO
                )
            )

        # 3. Call the google.genai API
        client = self._initialize_google_client()
        model_name = self._resolve_model_name(request_params.model or DEFAULT_GOOGLE_MODEL)
        try:
            async with client.aio:
                # Use the async client
                api_response = None
                streaming_supported = response_schema is None and response_mime_type is None
                if streaming_supported:
                    api_response = await self._stream_generate_content(
                        model=model_name,
                        contents=conversation_history,
                        config=generate_content_config,
                        client=client,
                    )
                if api_response is None:
                    api_response = await client.aio.models.generate_content(
                        model=model_name,
                        contents=cast("types.ContentListUnion", conversation_history),
                        config=generate_content_config,
                    )
                self.logger.debug("Google generate_content response:", data=api_response)

                # Track usage if response is valid and has usage data
                if (
                    hasattr(api_response, "usage_metadata")
                    and api_response.usage_metadata
                    and not isinstance(api_response, BaseException)
                ):
                    try:
                        turn_usage = TurnUsage.from_google(api_response.usage_metadata, model_name)
                        self._finalize_turn_usage(turn_usage)

                    except Exception as e:
                        self.logger.warning(f"Failed to track usage: {e}")

        except errors.APIError as e:
            # Handle specific Google API errors
            self.logger.error(f"Google API Error: {e.code} - {e.message}")
            raise ProviderKeyError(f"Google API Error: {e.code}", e.message or "") from e
        except Exception as e:
            self.logger.error(f"Error during Google generate_content call: {e}")
            raise e

        # 4. Process the API response
        if not api_response.candidates:
            # No response from the model, we're done
            self.logger.debug("No candidates returned.")
            return Prompt.assistant(stop_reason=LlmStopReason.END_TURN)

        candidate = api_response.candidates[0]  # Process the first candidate

        # Convert the model's response content to fast-agent types
        # Handle case where candidate.content might be None
        candidate_content = candidate.content
        if candidate_content is None:
            model_response_content_parts: list[ContentBlock | CallToolRequestParams] = []
        else:
            model_response_content_parts = self._converter.convert_from_google_content(
                candidate_content
            )
        stop_reason = LlmStopReason.END_TURN
        tool_calls: dict[str, CallToolRequest] | None = None
        # Add model's response to the working conversation history for this turn
        if candidate_content is not None:
            conversation_history.append(candidate_content)

        # Extract and process text content and tool calls
        assistant_message_parts = []
        tool_calls_to_execute = []

        for part in model_response_content_parts:
            if isinstance(part, TextContent):
                responses.append(part)  # Add text content to the final responses to be returned
                assistant_message_parts.append(
                    part
                )  # Collect text for potential assistant message display
            elif isinstance(part, CallToolRequestParams):
                # This is a function call requested by the model
                # If in structured mode, ignore tool calls per either-or rule
                if response_schema or response_mime_type:
                    continue
                tool_calls_to_execute.append(part)  # Collect tool calls to execute

        if tool_calls_to_execute:
            stop_reason = LlmStopReason.TOOL_USE
            tool_calls = {}
            for tool_call_params in tool_calls_to_execute:
                # Convert to CallToolRequest and execute
                tool_call_request = CallToolRequest(method="tools/call", params=tool_call_params)
                hex_string = secrets.token_hex(3)[:5]
                tool_calls[hex_string] = tool_call_request

            self.logger.debug("Tool call results processed.")
        else:
            stop_reason = self._map_finish_reason(getattr(candidate, "finish_reason", None))

        # Update diagnostic snapshot (never read again)
        # This provides a snapshot of what was sent to the provider for debugging
        self.history.set(conversation_history)

        self._log_chat_finished(model=model_name)  # Use resolved model name
        return Prompt.assistant(*responses, stop_reason=stop_reason, tool_calls=tool_calls)

    #        return responses  # Return the accumulated responses (fast-agent content types)

    async def _apply_prompt_provider_specific(
        self,
        multipart_messages: list[PromptMessageExtended],
        request_params: RequestParams | None = None,
        tools: list[McpTool] | None = None,
        is_template: bool = False,
    ) -> PromptMessageExtended:
        """
        Provider-specific prompt application.
        Templates are handled by the agent; messages already include them.
        """
        request_params = self.get_request_params(request_params=request_params)

        # Determine the last message
        last_message = multipart_messages[-1]

        if last_message.role == "assistant":
            # No generation required; the provided assistant message is the output
            return last_message

        # Build the provider-native message list for this turn from the last user message
        # This must handle tool results as function responses before any additional user content.
        turn_messages: list[types.Content] = []

        # 1) Convert tool results (if any) to google function responses
        if last_message.tool_results:
            # Map correlation IDs back to tool names using the last assistant tool_calls
            # found in our high-level message history
            id_to_name: dict[str, str] = {}
            for prev in reversed(multipart_messages):
                if prev.role == "assistant" and prev.tool_calls:
                    for call_id, call in prev.tool_calls.items():
                        try:
                            id_to_name[call_id] = call.params.name
                        except Exception:
                            pass
                    break

            tool_results_pairs = []
            for call_id, result in last_message.tool_results.items():
                tool_name = id_to_name.get(call_id, "tool")
                tool_results_pairs.append((tool_name, result))

            if tool_results_pairs:
                turn_messages.extend(
                    self._converter.convert_function_results_to_google(tool_results_pairs)
                )

        # 2) Convert any direct user content in the last message
        if last_message.content:
            user_contents = self._converter.convert_to_google_content([last_message])
            # convert_to_google_content returns a list; preserve order after tool responses
            turn_messages.extend(user_contents)

        # If we somehow have no provider-native parts, ensure we send an empty user content
        if not turn_messages:
            turn_messages.append(types.Content(role="user", parts=[types.Part.from_text(text="")]))

        conversation_history: list[types.Content] = []
        if request_params.use_history and len(multipart_messages) > 1:
            conversation_history.extend(self._convert_to_provider_format(multipart_messages[:-1]))
        conversation_history.extend(turn_messages)

        return await self._google_completion(
            conversation_history,
            request_params=request_params,
            tools=tools,
        )

    def _convert_extended_messages_to_provider(
        self, messages: list[PromptMessageExtended]
    ) -> list[types.Content]:
        """
        Convert PromptMessageExtended list to Google types.Content format.
        This is called fresh on every API call from _convert_to_provider_format().

        Args:
            messages: List of PromptMessageExtended objects

        Returns:
            List of Google types.Content objects
        """
        return self._converter.convert_to_google_content(messages)

    def _map_finish_reason(self, finish_reason: object) -> LlmStopReason:
        """Map Google finish reasons to LlmStopReason robustly."""
        # Normalize to string if it's an enum-like object
        reason = None
        try:
            reason = str(finish_reason) if finish_reason is not None else None
        except Exception:
            reason = None

        if not reason:
            return LlmStopReason.END_TURN

        # Extract last token after any dots or enum prefixes
        key = reason.split(".")[-1].upper()

        if key in {"STOP"}:
            return LlmStopReason.END_TURN
        if key in {"MAX_TOKENS", "LENGTH"}:
            return LlmStopReason.MAX_TOKENS
        if key in {
            "PROHIBITED_CONTENT",
            "SAFETY",
            "RECITATION",
            "BLOCKLIST",
            "SPII",
            "IMAGE_SAFETY",
            "IMAGE_PROHIBITED_CONTENT",
            "IMAGE_RECITATION",
        }:
            return LlmStopReason.SAFETY
        if key in {
            "MALFORMED_FUNCTION_CALL",
            "UNEXPECTED_TOOL_CALL",
            "TOO_MANY_TOOL_CALLS",
            "NO_IMAGE",
            "IMAGE_OTHER",
        }:
            return LlmStopReason.ERROR
        # Some SDKs include OTHER, LANGUAGE, GROUNDING, UNSPECIFIED, etc.
        return LlmStopReason.ERROR

    async def _apply_prompt_provider_specific_structured(
        self,
        multipart_messages,
        model,
        request_params=None,
    ):
        """
        Provider-specific structured output implementation.
        Note: Message history is managed by base class and converted via
        _convert_to_provider_format() on each call.
        """
        import json

        # Determine the last message
        last_message = multipart_messages[-1] if multipart_messages else None

        # If the last message is an assistant message, attempt to parse its JSON and return
        if last_message and last_message.role == "assistant":
            assistant_text = last_message.last_text()
            if assistant_text:
                try:
                    json_data = json.loads(assistant_text)
                    validated_model = model.model_validate(json_data)
                    return validated_model, last_message
                except (json.JSONDecodeError, Exception) as e:
                    self.logger.warning(
                        f"Failed to parse assistant message as structured response: {e}"
                    )
                    return None, last_message

        # Prepare request params
        request_params = self.get_request_params(request_params)

        # Build schema for structured output
        schema = None
        try:
            schema = model.model_json_schema()
        except Exception:
            pass
        response_schema = model if schema is None else schema

        # Convert the last user message to provider-native content for the current turn
        turn_messages: list[types.Content] = []
        if last_message:
            turn_messages = self._converter.convert_to_google_content([last_message])

        # Delegate to unified completion with structured options enabled (no tools)
        assistant_msg = await self._google_completion(
            turn_messages,
            request_params=request_params,
            tools=None,
            response_mime_type="application/json",
            response_schema=response_schema,
        )

        # Parse using shared helper for consistency
        parsed, _ = self._structured_from_multipart(assistant_msg, model)
        return parsed, assistant_msg
