"""Google Gemini provider implementation."""

from __future__ import annotations

import asyncio
import time
from typing import Any

from pi_agents.ai.env_api_keys import get_env_api_key
from pi_agents.ai.models import calculate_cost
from pi_agents.ai.types import (
    AssistantMessage,
    Context,
    Cost,
    Model,
    SimpleStreamOptions,
    TextContent,
    ToolCall,
    Usage,
)
from pi_agents.ai.utils.event_stream import AssistantMessageEventStream


def stream_google(
    model: Model,
    context: Context,
    options: dict[str, Any] | None = None,
) -> AssistantMessageEventStream:
    """Stream responses from the Google Generative AI API."""
    event_stream = AssistantMessageEventStream()
    opts = options or {}

    async def _run():
        try:
            from google import genai
            from google.genai import types as genai_types

            api_key = opts.get("api_key") or opts.get("apiKey") or get_env_api_key("google")
            if not api_key:
                raise ValueError("Google API key not found. Set GEMINI_API_KEY environment variable.")

            client = genai.Client(api_key=api_key)

            # Build contents
            contents: list[Any] = []
            for msg in context.messages:
                role = msg.role if hasattr(msg, "role") else msg.get("role", "user")
                if role == "user":
                    content = msg.content if hasattr(msg, "content") else msg.get("content", "")
                    if isinstance(content, str):
                        contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))
                    else:
                        parts = []
                        for block in content:
                            btype = block.type if hasattr(block, "type") else block.get("type", "text")
                            if btype == "text":
                                text = block.text if hasattr(block, "text") else block.get("text", "")
                                parts.append(genai_types.Part(text=text))
                        contents.append(genai_types.Content(role="user", parts=parts))
                elif role == "assistant":
                    content_list = msg.content if hasattr(msg, "content") else msg.get("content", [])
                    parts = []
                    for block in content_list:
                        btype = block.type if hasattr(block, "type") else block.get("type", "text")
                        if btype == "text":
                            text = block.text if hasattr(block, "text") else block.get("text", "")
                            parts.append(genai_types.Part(text=text))
                    contents.append(genai_types.Content(role="model", parts=parts))

            config: dict[str, Any] = {}
            if context.system_prompt:
                config["system_instruction"] = context.system_prompt

            max_tokens = opts.get("max_tokens") or opts.get("maxTokens") or model.max_tokens
            if max_tokens:
                config["max_output_tokens"] = max_tokens

            if opts.get("temperature") is not None:
                config["temperature"] = opts["temperature"]

            partial = AssistantMessage(
                role="assistant",
                content=[],
                api=model.api,
                provider=model.provider,
                model=model.id,
                usage=Usage(cost=Cost()),
                stop_reason="stop",
                timestamp=int(time.time() * 1000),
            )

            event_stream.push({"type": "start", "partial": partial})

            text_buffer = ""
            content_index = 0

            response = client.models.generate_content_stream(
                model=model.id,
                contents=contents,
                config=genai_types.GenerateContentConfig(**config) if config else None,
            )

            for chunk in response:
                if chunk.text:
                    if not text_buffer:
                        event_stream.push({"type": "text_start", "contentIndex": content_index, "partial": partial})
                    text_buffer += chunk.text
                    event_stream.push({
                        "type": "text_delta",
                        "contentIndex": content_index,
                        "delta": chunk.text,
                        "partial": partial,
                    })

                if chunk.usage_metadata:
                    partial.usage.input = chunk.usage_metadata.prompt_token_count or 0
                    partial.usage.output = chunk.usage_metadata.candidates_token_count or 0
                    partial.usage.total_tokens = chunk.usage_metadata.total_token_count or 0

            if text_buffer:
                event_stream.push({
                    "type": "text_end",
                    "contentIndex": content_index,
                    "content": text_buffer,
                    "partial": partial,
                })
                partial.content.append(TextContent(type="text", text=text_buffer))

            calculate_cost(model, partial.usage)
            event_stream.push({"type": "done", "reason": "stop", "message": partial})

        except Exception as e:
            error_msg = AssistantMessage(
                role="assistant",
                content=[],
                api=model.api,
                provider=model.provider,
                model=model.id,
                usage=Usage(cost=Cost()),
                stop_reason="error",
                error_message=str(e),
                timestamp=int(time.time() * 1000),
            )
            event_stream.push({"type": "error", "reason": "error", "error": error_msg})

        event_stream.end()

    asyncio.ensure_future(_run())
    return event_stream


def stream_simple_google(
    model: Model,
    context: Context,
    options: SimpleStreamOptions | None = None,
) -> AssistantMessageEventStream:
    """Stream with simplified options."""
    opts: dict[str, Any] = {}
    if options:
        if options.temperature is not None:
            opts["temperature"] = options.temperature
        if options.max_tokens is not None:
            opts["maxTokens"] = options.max_tokens
        if options.api_key:
            opts["apiKey"] = options.api_key
    return stream_google(model, context, opts)
