"""Google GenAI (Gemini) provider."""

from __future__ import annotations

import json
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import Any

from google import genai
from google.genai import types as genai_types

from ..types import ChatMessage, ChatResponse, StreamChunk, ToolCallRequest, Usage
from .errors import StreamFailedError


def _format_contents(
    messages: list[ChatMessage],
) -> tuple[str | None, list[genai_types.Content]]:
    """Convert ChatMessages to Google GenAI format.

    Returns (system_instruction, contents).
    """
    system_parts: list[str] = []
    contents: list[genai_types.Content] = []

    for msg in messages:
        if msg.role == "system":
            system_parts.append(msg.text)
            continue

        role = "model" if msg.role == "assistant" else "user"
        if msg.role == "tool":
            role = "user"

        parts: list[genai_types.Part] = []

        if isinstance(msg.content, str):
            if msg.content:
                parts.append(genai_types.Part(text=msg.content))
        else:
            for block in msg.content:
                if block.type == "text" and block.text:
                    parts.append(genai_types.Part(text=block.text))

        if msg.tool_calls:
            for tc in msg.tool_calls:
                parts.append(
                    genai_types.Part(
                        function_call=genai_types.FunctionCall(name=tc.name, args=tc.arguments)
                    )
                )

        if msg.role == "tool" and msg.tool_call_id:
            try:
                result = json.loads(msg.text)
            except (json.JSONDecodeError, TypeError):
                result = {"result": msg.text}
            parts = [
                genai_types.Part(
                    function_response=genai_types.FunctionResponse(
                        name=msg.name or "tool", response=result
                    )
                )
            ]

        if parts:
            contents.append(genai_types.Content(role=role, parts=parts))

    system_instruction = "\n\n".join(system_parts) if system_parts else None
    return system_instruction, contents


def _tools_to_google(tools: list[dict[str, Any]]) -> list[genai_types.Tool]:
    """Convert OpenAI-format tools to Google GenAI format."""
    declarations: list[genai_types.FunctionDeclaration] = []
    for tool in tools:
        if tool.get("type") == "function":
            func = tool["function"]
            declarations.append(
                genai_types.FunctionDeclaration(
                    name=func["name"],
                    description=func.get("description", ""),
                    parameters=func.get("parameters"),
                )
            )
    return [genai_types.Tool(function_declarations=declarations)] if declarations else []


def _strip_additional_properties(schema: Any) -> Any:
    """Recursively remove additionalProperties from a JSON schema.

    Gemini's response_schema uses a restricted OpenAPI 3.0 subset that does not
    support additionalProperties. Passing it causes a 400 INVALID_ARGUMENT error.
    """
    if isinstance(schema, Mapping):
        return {
            k: _strip_additional_properties(v)
            for k, v in schema.items()
            if k != "additionalProperties"
        }
    if isinstance(schema, Sequence) and not isinstance(schema, str):
        return [_strip_additional_properties(item) for item in schema]
    return schema


def _apply_response_format(
    config: genai_types.GenerateContentConfig,
    response_format: dict[str, Any],
) -> None:
    """Apply OpenAI-style response_format to Google GenAI config.

    Converts the OpenAI JSON Schema format to Google's response_mime_type
    and response_schema parameters.
    """
    fmt_type = response_format.get("type")
    if fmt_type == "json_schema":
        json_schema_spec = response_format.get("json_schema", {})
        schema = json_schema_spec.get("schema", {})
        if schema:
            config.response_mime_type = "application/json"
            config.response_schema = _strip_additional_properties(schema)
    elif fmt_type == "json_object":
        config.response_mime_type = "application/json"


class GoogleClient:
    """Google GenAI (Gemini) client."""

    def __init__(
        self,
        api_key: str,
        model: str,
        temperature: float = 0.7,
        max_tokens: int = 4096,
    ) -> None:
        self._client = genai.Client(api_key=api_key)
        self._model = model
        self._temperature = temperature
        self._max_tokens = max_tokens

    @property
    def model_name(self) -> str:
        return self._model

    async def chat(
        self,
        messages: list[ChatMessage],
        *,
        tools: list[dict[str, Any]] | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        response_format: dict[str, Any] | None = None,
    ) -> ChatResponse:
        system_instruction, contents = _format_contents(messages)

        config = genai_types.GenerateContentConfig(
            temperature=temperature if temperature is not None else self._temperature,
            max_output_tokens=max_tokens if max_tokens is not None else self._max_tokens,
        )

        if system_instruction:
            config.system_instruction = system_instruction

        if tools:
            config.tools = _tools_to_google(tools)

        if response_format:
            _apply_response_format(config, response_format)

        response = await self._client.aio.models.generate_content(
            model=self._model,
            contents=contents,
            config=config,
        )

        # Extract text and tool calls
        text_parts: list[str] = []
        tool_calls: list[ToolCallRequest] = []

        if response.candidates and response.candidates[0].content:
            for part in response.candidates[0].content.parts or []:
                if part.text:
                    text_parts.append(part.text)
                if part.function_call:
                    tool_calls.append(
                        ToolCallRequest(
                            id=f"call_{part.function_call.name}",
                            name=part.function_call.name,
                            arguments=(
                                dict(part.function_call.args) if part.function_call.args else {}
                            ),
                        )
                    )

        # Extract usage
        usage = Usage()
        if response.usage_metadata:
            usage = Usage(
                prompt_tokens=response.usage_metadata.prompt_token_count or 0,
                completion_tokens=response.usage_metadata.candidates_token_count or 0,
                total_tokens=response.usage_metadata.total_token_count or 0,
            )

        return ChatResponse(
            message=ChatMessage(
                role="assistant",
                content="\n".join(text_parts),
                tool_calls=tool_calls if tool_calls else None,
            ),
            usage=usage,
            model=self._model,
        )

    async def chat_stream(
        self,
        messages: list[ChatMessage],
        *,
        tools: list[dict[str, Any]] | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
    ) -> AsyncIterator[StreamChunk]:
        system_instruction, contents = _format_contents(messages)

        config = genai_types.GenerateContentConfig(
            temperature=temperature if temperature is not None else self._temperature,
            max_output_tokens=max_tokens if max_tokens is not None else self._max_tokens,
        )

        if system_instruction:
            config.system_instruction = system_instruction

        if tools:
            config.tools = _tools_to_google(tools)

        # A legitimate Gemini stream emits at least one of: a text part,
        # a function-call part, a candidate-level finish_reason, or a
        # usage_metadata frame. If none of these arrive, the upstream
        # silently rejected the request (revoked API key, quota miss,
        # certain safety blocks that bypass the normal block_reason path)
        # and we raise so callers don't render a blank reply.
        produced_anything = False

        async for chunk in await self._client.aio.models.generate_content_stream(
            model=self._model,
            contents=contents,
            config=config,
        ):
            text = ""
            tool_call_list: list[ToolCallRequest] = []

            if chunk.candidates and chunk.candidates[0].content:
                for part in chunk.candidates[0].content.parts or []:
                    if part.text:
                        text += part.text
                    if part.function_call:
                        tool_call_list.append(
                            ToolCallRequest(
                                id=f"call_{part.function_call.name}",
                                name=part.function_call.name,
                                arguments=(
                                    dict(part.function_call.args) if part.function_call.args else {}
                                ),
                            )
                        )

            if text or tool_call_list:
                produced_anything = True
            if chunk.candidates and chunk.candidates[0].finish_reason:
                produced_anything = True

            usage_chunk: Usage | None = None
            if chunk.usage_metadata:
                usage_chunk = Usage(
                    prompt_tokens=chunk.usage_metadata.prompt_token_count or 0,
                    completion_tokens=chunk.usage_metadata.candidates_token_count or 0,
                    total_tokens=chunk.usage_metadata.total_token_count or 0,
                )
                # usage_metadata is Gemini's most reliable terminal
                # signal — present on every legitimate completion,
                # absent on silent rejections.
                produced_anything = True

            yield StreamChunk(
                delta_text=text,
                delta_tool_calls=tool_call_list if tool_call_list else None,
                usage=usage_chunk,
            )

        if not produced_anything:
            raise StreamFailedError(
                f"Google stream closed with no content for model {self._model!r}",
                provider="google",
            )


__all__ = ["GoogleClient"]
