from __future__ import annotations
from typing import Any, List

from .base import ProviderConfig, LLMMessage, LLMResponse

class GeminiProvider:
    """Google GenAI (Gemini) Provider (real implementation) using `google-genai` SDK."""
    def __init__(self, config: ProviderConfig):
        self.config = config
        try:
            from google import genai  # type: ignore
            from google.genai import types  # type: ignore
        except Exception as e:  # pragma: no cover
            raise RuntimeError("Install provider extra: pip install adi-agentic-agi[gemini]") from e

        self._genai = genai
        self._types = types

        # Gemini Developer API: provide api_key; Vertex can be configured via env vars or config.extra
        if config.api_key:
            self._client = genai.Client(api_key=config.api_key).aio
        else:
            # supports Vertex AI if env vars are set as per Google docs
            extra = config.extra or {}
            if extra.get("vertexai"):
                self._client = genai.Client(
                    vertexai=True,
                    project=extra.get("project"),
                    location=extra.get("location", "us-central1"),
                ).aio
            else:
                self._client = genai.Client().aio

    @staticmethod
    def _system_and_user(messages: List[LLMMessage]) -> tuple[str, str]:
        system_parts = [m.content for m in messages if m.role == "system" and m.content]
        user_parts = [m.content for m in messages if m.role == "user" and m.content]
        return "\n\n".join(system_parts).strip(), "\n\n".join(user_parts).strip()

    async def generate(self, messages: List[LLMMessage], **kwargs: Any) -> LLMResponse:
        system, user = self._system_and_user(messages)
        cfg = self._types.GenerateContentConfig(
            system_instruction=system or None,
            temperature=kwargs.pop("temperature", None),
            max_output_tokens=kwargs.pop("max_output_tokens", None),
        )
        resp = await self._client.models.generate_content(
            model=self.config.model,
            contents=user,
            config=cfg,
            **kwargs,
        )
        text = getattr(resp, "text", None) or ""
        return LLMResponse(text=text, raw={})
