"""Google Generative AI (Gemini) provider implementation."""
from __future__ import annotations
import json
import logging
import time
from collections.abc import AsyncIterator, Awaitable, Callable
from dataclasses import replace
from typing import Any
logger = logging.getLogger(__name__)
from ..dtypes import (
    Actor,
    ChatCompletion,
    Choice,
    Message,
    StreamDone,
    StreamEvent,
    StreamStart,
    TextContent,
    TextDelta,
    TextEnd,
    TextStart,
    ThinkingContent,
    ThinkingDelta,
    ThinkingEnd,
    ThinkingStart,
    ToolCall,
    ToolCallContent,
    ToolCallDelta,
    ToolCallEnd,
    ToolCallStart,
    Usage,
)
from .base import calculate_cost_from_usage
async def aggregate_google_stream(
    stream: AsyncIterator,
    on_chunk: Callable[[StreamEvent], Awaitable[None]],
) -> tuple[Message, dict[str, Any]]:
    """Aggregate Google Generative AI streaming chunks into a complete Message.
    Handles Gemini models' streaming format with support for:
    - Text content streaming
    - Thinking/reasoning content (thought=True flag)
    - Tool calls (functionCall)
    Returns:
        tuple of (final_message, usage_dict)
    """
    assert stream is not None
    assert on_chunk is not None
    assert callable(on_chunk)
    # Emit start event
    await on_chunk(StreamStart())

    content_blocks: list[dict[str, Any]] = []
    current_block: dict[str, Any] | None = None
    finish_reason = "stop"
    # Usage tracking
    usage_data: dict[str, Any] = {}
    # Tool call counter for generating unique IDs
    tool_call_counter = 0
    async for chunk in stream:
        # Access the candidate from the chunk
        candidate = None
        if hasattr(chunk, "candidates") and chunk.candidates:
            candidate = chunk.candidates[0]
        # Process content parts
        if candidate and hasattr(candidate, "content") and candidate.content:
            parts = getattr(candidate.content, "parts", None)
            if parts:
                for part in parts:
                    if hasattr(part, "text") and part.text is not None:
                        is_thinking = getattr(part, "thought", False)
                        if (
                            not current_block
                            or (is_thinking and current_block.get("type") != "thinking")
                            or (not is_thinking and current_block.get("type") != "text")
                        ):
                            # Finalize previous block if exists
                            if current_block:
                                block_index = len(content_blocks) - 1
                                if current_block["type"] == "text":
                                    await on_chunk(
                                        TextEnd(
                                            content_index=block_index,
                                            content=current_block["text"],
                                        )
                                    )
                                else:
                                    await on_chunk(
                                        ThinkingEnd(
                                            content_index=block_index,
                                            content=current_block["thinking"],
                                        )
                                    )
                            # Start new block
                            if is_thinking:
                                current_block = {
                                    "type": "thinking",
                                    "thinking": "",
                                    "thinkingSignature": getattr(part, "thoughtSignature", None),
                                }
                                content_blocks.append(current_block)
                                await on_chunk(ThinkingStart(content_index=len(content_blocks) - 1))
                            else:
                                current_block = {"type": "text", "text": ""}
                                content_blocks.append(current_block)
                                await on_chunk(TextStart(content_index=len(content_blocks) - 1))
                        block_index = len(content_blocks) - 1
                        if current_block["type"] == "thinking":
                            current_block["thinking"] += part.text
                            if getattr(part, "thoughtSignature", None):
                                current_block["thinkingSignature"] = part.thoughtSignature
                            await on_chunk(
                                ThinkingDelta(
                                    content_index=block_index,
                                    delta=part.text,
                                )
                            )
                        else:
                            current_block["text"] += part.text
                            await on_chunk(
                                TextDelta(
                                    content_index=block_index,
                                    delta=part.text,
                                )
                            )
                    if hasattr(part, "functionCall") and part.functionCall:
                        # Finalize previous block if exists
                        if current_block:
                            block_index = len(content_blocks) - 1
                            if current_block["type"] == "text":
                                await on_chunk(
                                    TextEnd(
                                        content_index=block_index,
                                        content=current_block["text"],
                                    )
                                )
                            else:
                                await on_chunk(
                                    ThinkingEnd(
                                        content_index=block_index,
                                        content=current_block["thinking"],
                                    )
                                )
                            current_block = None
                        # Generate unique ID
                        fc = part.functionCall
                        provided_id: str | None = getattr(fc, "id", None)
                        needs_new_id = not provided_id or any(
                            b.get("type") == "toolCall" and b.get("id") == provided_id
                            for b in content_blocks
                        )
                        if needs_new_id:
                            tool_call_counter += 1
                            tool_call_id: str = f"{fc.name}_{int(time.time())}_{tool_call_counter}"
                        else:
                            assert (
                                provided_id is not None
                            )  # ensured by `not provided_id` check above
                            tool_call_id = provided_id
                        tool_call_block = {
                            "type": "toolCall",
                            "id": tool_call_id,
                            "name": getattr(fc, "name", ""),
                            "arguments": dict(getattr(fc, "args", {})),
                        }
                        if getattr(part, "thoughtSignature", None):
                            tool_call_block["thoughtSignature"] = part.thoughtSignature
                        content_blocks.append(tool_call_block)
                        block_index = len(content_blocks) - 1
                        # Emit tool call events
                        await on_chunk(
                            ToolCallStart(
                                content_index=block_index,
                                tool_call_id=tool_call_id,
                                tool_name=tool_call_block["name"],
                            )
                        )
                        await on_chunk(
                            ToolCallDelta(
                                content_index=block_index,
                                tool_call_id=tool_call_id,
                                delta=json.dumps(tool_call_block["arguments"]),
                                partial_args=tool_call_block["arguments"],
                            )
                        )
                        tool_call = ToolCall(
                            id=tool_call_id,
                            name=tool_call_block["name"],
                            args=tool_call_block["arguments"],
                        )
                        await on_chunk(
                            ToolCallEnd(
                                content_index=block_index,
                                tool_call=tool_call,
                            )
                        )
        if candidate and hasattr(candidate, "finishReason") and candidate.finishReason:

            finish_reason_value = candidate.finishReason
            if hasattr(finish_reason_value, "name"):
                finish_reason_name = finish_reason_value.name
            else:
                finish_reason_name = str(finish_reason_value)
            if finish_reason_name == "STOP":
                finish_reason = "stop"
            elif finish_reason_name == "MAX_TOKENS":
                finish_reason = "length"
            else:
                finish_reason = "stop"  # Default for other reasons
            # Override if we have tool calls
            has_tool_calls = any(b.get("type") == "toolCall" for b in content_blocks)
            if has_tool_calls:
                finish_reason = "tool_calls"
        if hasattr(chunk, "usageMetadata") and chunk.usageMetadata:
            metadata = chunk.usageMetadata
            cached_tokens = getattr(metadata, "cachedContentTokenCount", 0) or 0
            prompt_tokens = getattr(metadata, "promptTokenCount", 0) or 0
            usage_data = {
                "input_tokens": prompt_tokens - cached_tokens,  # Non-cached input
                "output_tokens": getattr(metadata, "candidatesTokenCount", 0) or 0,
                "reasoning_tokens": getattr(metadata, "thoughtsTokenCount", 0) or 0,
                "cache_read_tokens": cached_tokens,
                "cache_write_tokens": 0,
            }
    # Finalize current block if exists
    if current_block:
        block_index = len(content_blocks) - 1
        if current_block["type"] == "text":
            await on_chunk(
                TextEnd(
                    content_index=block_index,
                    content=current_block["text"],
                )
            )
        else:
            await on_chunk(
                ThinkingEnd(
                    content_index=block_index,
                    content=current_block["thinking"],
                )
            )
    # Emit done event
    await on_chunk(StreamDone(finish_reason=finish_reason))
    final_content_blocks: list = []
    for block in content_blocks:
        if block["type"] == "text":
            final_content_blocks.append(TextContent(text=block["text"]))
        elif block["type"] == "thinking":
            final_content_blocks.append(
                ThinkingContent(
                    thinking=block["thinking"],
                    thinking_signature=block.get("thinkingSignature"),
                )
            )
        elif block["type"] == "toolCall":
            final_content_blocks.append(
                ToolCallContent(
                    id=block["id"],
                    name=block["name"],
                    arguments=block["arguments"],
                    thought_signature=block.get("thoughtSignature"),
                )
            )
    final_message = Message(
        role="assistant",
        content=final_content_blocks,
    )
    return final_message, usage_data
async def rollout_google(
    actor: Actor,
    on_chunk: Callable[[StreamEvent], Awaitable[None]],
    **kwargs: Any,
) -> Actor:
    """Make a Google Generative AI (Gemini) API call with streaming.
    Note: **kwargs accepts but ignores provider-specific params (e.g., anthropic thinking params)
    Note: Uses trio-asyncio to bridge trio event loop with asyncio-based google-genai SDK
    """
    assert actor is not None
    assert isinstance(actor, Actor)
    assert actor.endpoint is not None
    assert actor.trajectory is not None
    assert on_chunk is not None
    assert callable(on_chunk)
    try:
        from google import genai
        from google.genai import types
    except ImportError as e:
        raise ImportError(
            "Google Generative AI SDK not installed. Install with: pip install google-genai"
        ) from e
    try:
        import trio_asyncio
    except ImportError as e:
        raise ImportError(
            "trio-asyncio not installed. Install with: pip install trio-asyncio"
        ) from e
    api_key = actor.endpoint.api_key
    if not api_key:
        import os
        api_key = os.getenv("GEMINI_API_KEY")
        if not api_key:
            raise ValueError(
                "Gemini API key is required. Set GEMINI_API_KEY environment variable "
                "or provide it in endpoint.api_key"
            )
    # Transform messages for cross-provider compatibility (like pi-ai does)
    from ..transform_messages import transform_messages
    transformed_messages = transform_messages(
        actor.trajectory.messages,
        target_provider=actor.endpoint.provider,
        target_api="google-generative-ai",
    )
    # Prepare message conversion outside asyncio context
    contents = []
    for m in transformed_messages:
        if m.role == "user":
            # Tiger Style: Explicit control flow - handle both string and ContentBlock content
            if isinstance(m.content, str):
                user_text = m.content
            elif isinstance(m.content, list):
                text_blocks = [b for b in m.content if isinstance(b, TextContent)]
                user_text = "\n".join(b.text for b in text_blocks) if text_blocks else ""
            else:
                user_text = ""
            contents.append(types.Content(role="user", parts=[types.Part(text=user_text)]))
        elif m.role == "assistant":
            parts = []
            # Tiger Style: Explicit control flow - handle both string and ContentBlock content
            if isinstance(m.content, str):
                parts.append(types.Part(text=m.content))
            elif isinstance(m.content, list):
                # Process ContentBlocks
                for block in m.content:
                    if isinstance(block, TextContent):
                        parts.append(types.Part(text=block.text))
                    elif isinstance(block, ToolCallContent):
                        parts.append(
                            types.Part(
                                function_call=types.FunctionCall(
                                    name=block.name,
                                    args=block.arguments,
                                )
                            )
                        )
            if parts:
                contents.append(types.Content(role="model", parts=parts))
        elif m.role == "tool":
            # Tiger Style: Explicit control flow - handle both string and ContentBlock content
            if isinstance(m.content, str):
                tool_result_text = m.content
            elif isinstance(m.content, list):
                text_blocks = [b for b in m.content if isinstance(b, TextContent)]
                tool_result_text = "\n".join(b.text for b in text_blocks) if text_blocks else ""
            else:
                tool_result_text = ""
            contents.append(
                types.Content(
                    role="user",
                    parts=[
                        types.Part(
                            function_response=types.FunctionResponse(
                                name=m.tool_name if m.tool_name else "unknown",
                                response={"result": tool_result_text},
                            )
                        )
                    ],
                )
            )
    config = types.GenerateContentConfig(
        temperature=actor.endpoint.temperature if hasattr(actor.endpoint, "temperature") else None,
        max_output_tokens=actor.endpoint.max_tokens
        if hasattr(actor.endpoint, "max_tokens")
        else None,
    )
    if actor.tools:
        tools = []
        function_declarations = []
        for tool in actor.tools:
            function_declarations.append(
                types.FunctionDeclaration(
                    name=tool.function.name,
                    description=tool.function.description,
                    parameters={
                        "type": tool.function.parameters.type,
                        "properties": tool.function.parameters.properties,
                        "required": tool.function.required,
                    },
                )
            )
        tools.append(types.Tool(function_declarations=function_declarations))
        config.tools = tools
    # Wide event logging for API request
    from .base import log_api_request
    log_api_request(
        provider="google",
        model=actor.endpoint.model,
        api_base=None,  # Google doesn't use custom base URLs
        messages=[{"role": c.role, "parts": len(c.parts)} for c in contents],
        tools=[t.function.name for t in actor.tools] if actor.tools else [],
        temperature=actor.endpoint.temperature if hasattr(actor.endpoint, "temperature") else None,
        max_tokens=actor.endpoint.max_tokens if hasattr(actor.endpoint, "max_tokens") else None,
    )
    # The actual API call in asyncio context
    async def _call_google_in_asyncio() -> ChatCompletion:
        client = genai.Client(api_key=api_key)
        try:
            # Generate streaming response
            stream = await client.aio.models.generate_content_stream(
                model=actor.endpoint.model,
                contents=contents,
                config=config,
            )
            final_message, usage_data = await aggregate_google_stream(stream, on_chunk)
            return final_message, usage_data
        except Exception as e:
            from ..providers.base import ProviderError, log_api_response
            log_api_response(
                provider="google",
                model=actor.endpoint.model,
                attempt=1,
                success=False,
                error_type=type(e).__name__,
                error_message=str(e),
            )
            raise ProviderError(
                f"Google API error: {e}",
                original_error=e,
                attempts=1,  # Google SDK doesn't expose retry count
                provider="google",
            ) from e
    # Run the asyncio function from trio context using trio-asyncio loop
    async with trio_asyncio.open_loop() as _loop:
        final_message, usage_data = await trio_asyncio.aio_as_trio(_call_google_in_asyncio)()
    # Wide event for successful response
    from .base import log_api_response
    log_api_response(
        provider="google",
        model=actor.endpoint.model,
        attempt=1,
        success=True,
        input_tokens=usage_data.get("input_tokens", 0),
        output_tokens=usage_data.get("output_tokens", 0),
        reasoning_tokens=usage_data.get("reasoning_tokens", 0),
        cache_read_tokens=usage_data.get("cache_read_tokens", 0),
        has_tool_calls=any(
            isinstance(b, ToolCallContent)
            for b in (final_message.content if isinstance(final_message.content, list) else [])
        ),
    )
    usage = Usage(
        input_tokens=usage_data.get("input_tokens", 0),
        output_tokens=usage_data.get("output_tokens", 0),
        reasoning_tokens=usage_data.get("reasoning_tokens", 0),
        cache_read_tokens=usage_data.get("cache_read_tokens", 0),
        cache_write_tokens=usage_data.get("cache_write_tokens", 0),
    )
    from ..models import get_model
    model_meta = get_model(actor.endpoint.provider, actor.endpoint.model)
    if model_meta and model_meta.cost:
        cost = calculate_cost_from_usage(usage, model_meta.cost)
        usage = replace(usage, cost=cost)
    # Enrich message with provider/api/model metadata for cross-provider handoff
    final_message = replace(
        final_message,
        provider=actor.endpoint.provider,
        api="google-generative-ai",
        model=actor.endpoint.model,
    )
    completion = ChatCompletion(
        id="google-" + str(int(time.time())),
        object="chat.completion",
        created=int(time.time()),
        model=actor.endpoint.model,
        usage=usage,
        choices=[Choice(0, final_message, "stop")],
    )
    new_trajectory = replace(
        actor.trajectory,
        messages=actor.trajectory.messages + [final_message],
        completions=actor.trajectory.completions + [completion],
    )
    result_actor = replace(actor, trajectory=new_trajectory)
    assert result_actor is not None
    assert result_actor.trajectory is not None
    return result_actor
