Session Buddy

"""Google Gemini API provider implementation.

This module provides the Gemini provider implementation using the Google
Generative AI SDK for chat completions and streaming.
"""

from __future__ import annotations

from datetime import datetime
from typing import TYPE_CHECKING, Any

from session_buddy.llm.base import LLMProvider
from session_buddy.llm.models import LLMMessage, LLMResponse

if TYPE_CHECKING:
    from collections.abc import AsyncGenerator


class GeminiProvider(LLMProvider):
    """Google Gemini API provider."""

    def __init__(self, config: dict[str, Any]) -> None:
        super().__init__(config)
        self.api_key = config.get("api_key")
        self.default_model = config.get("default_model", "gemini-pro")
        self._client = None

    async def _get_client(self) -> Any:
        """Get or create Gemini client."""
        if self._client is None:
            try:
                import google.generativeai as genai

                genai.configure(api_key=self.api_key)
                self._client = genai
            except ImportError:
                msg = "Google Generative AI package not installed. Install with: pip install google-generativeai"
                raise ImportError(
                    msg,
                )
        return self._client

    def _convert_messages(self, messages: list[LLMMessage]) -> list[dict[str, Any]]:
        """Convert LLMMessage objects to Gemini format using modern pattern matching."""
        converted: list[dict[str, Any]] = []

        for msg in messages:
            match msg.role:
                case "system":
                    # Gemini doesn't have system role, prepend to first user message
                    if converted and converted[-1]["role"] == "user":
                        converted[-1]["parts"] = [
                            f"System: {msg.content}\n\nUser: {converted[-1]['parts'][0]}",
                        ]
                    else:
                        converted.append(
                            {"role": "user", "parts": [f"System: {msg.content}"]},
                        )
                case "user":
                    converted.append({"role": "user", "parts": [msg.content]})
                case "assistant":
                    converted.append({"role": "model", "parts": [msg.content]})
                case _:
                    # Unknown role - default to user for safety
                    converted.append({"role": "user", "parts": [msg.content]})

        return converted

    async def generate(
        self,
        messages: list[LLMMessage],
        model: str | None = None,
        temperature: float = 0.7,
        max_tokens: int | None = None,
        **kwargs: Any,
    ) -> LLMResponse:
        """Generate response using Gemini API."""
        if not await self.is_available():
            msg = "Gemini provider not available"
            raise RuntimeError(msg)

        genai = await self._get_client()
        model_name = model or self.default_model

        try:
            model_instance = genai.GenerativeModel(model_name)

            # Convert messages to Gemini chat format
            chat_messages = self._convert_messages(messages)

            # Create chat or generate single response
            if len(chat_messages) > 1:
                chat = model_instance.start_chat(history=chat_messages[:-1])
                response = await chat.send_message_async(
                    chat_messages[-1]["parts"][0],
                    generation_config={
                        "temperature": temperature,
                        "max_output_tokens": max_tokens,
                    },
                )
            else:
                response = await model_instance.generate_content_async(
                    chat_messages[0]["parts"][0],
                    generation_config={
                        "temperature": temperature,
                        "max_output_tokens": max_tokens,
                    },
                )

            return LLMResponse(
                content=response.text,
                model=model_name,
                provider="gemini",
                usage={
                    "prompt_tokens": response.usage_metadata.prompt_token_count
                    if hasattr(response, "usage_metadata")
                    else 0,
                    "completion_tokens": response.usage_metadata.candidates_token_count
                    if hasattr(response, "usage_metadata")
                    else 0,
                    "total_tokens": response.usage_metadata.total_token_count
                    if hasattr(response, "usage_metadata")
                    else 0,
                },
                finish_reason="stop",  # Gemini doesn't provide detailed finish reasons
                timestamp=datetime.now().isoformat(),
            )

        except Exception as e:
            self.logger.exception(f"Gemini generation failed: {e}")
            raise

    async def stream_generate(  # type: ignore[override]
        self,
        messages: list[LLMMessage],
        model: str | None = None,
        temperature: float = 0.7,
        max_tokens: int | None = None,
        **kwargs: Any,
    ) -> AsyncGenerator[str]:
        """Stream response using Gemini API."""
        if not await self.is_available():
            msg = "Gemini provider not available"
            raise RuntimeError(msg)

        genai = await self._get_client()
        model_name = model or self.default_model

        try:
            model_instance = genai.GenerativeModel(model_name)
            chat_messages = self._convert_messages(messages)

            if len(chat_messages) > 1:
                chat = model_instance.start_chat(history=chat_messages[:-1])
                response = chat.send_message(
                    chat_messages[-1]["parts"][0],
                    generation_config={
                        "temperature": temperature,
                        "max_output_tokens": max_tokens,
                    },
                    stream=True,
                )
            else:
                response = model_instance.generate_content(
                    chat_messages[0]["parts"][0],
                    generation_config={
                        "temperature": temperature,
                        "max_output_tokens": max_tokens,
                    },
                    stream=True,
                )

            for chunk in response:
                if chunk.text:
                    yield chunk.text

        except Exception as e:
            self.logger.exception(f"Gemini streaming failed: {e}")
            raise

    async def is_available(self) -> bool:
        """Check if Gemini API is available."""
        if not self.api_key:
            return False

        try:
            genai = await self._get_client()
            # Test with a simple model list request
            list(genai.list_models())
            return True
        except Exception:
            return False

    def get_models(self) -> list[str]:
        """Get available Gemini models."""
        return [
            "gemini-pro",
            "gemini-pro-vision",
            "gemini-1.5-pro",
            "gemini-1.5-flash",
            "gemini-1.0-pro",
        ]