from __future__ import annotations

import json
import os
from typing import Any

from nexus.protocols.providers.base import AbstractProvider, LLMRequest, LLMResponse
from nexus.resilience.errors import ProviderError

try:
    import google.generativeai as genai
except ImportError:  # optional dependency path
    genai = None


class GeminiProvider(AbstractProvider):
    """Google Gemini provider using google-generativeai SDK."""

    def __init__(self, model: str = "gemini-1.5-pro", api_key: str | None = None) -> None:
        if genai is None:
            raise ProviderError(
                "google-generativeai is not installed. Install nexus-ai[google].",
                provider_name="gemini",
            )
        self._model = model
        self._api_key = api_key or os.getenv("GOOGLE_API_KEY")
        genai.configure(api_key=self._api_key)
        self._client = genai.GenerativeModel(model_name=model)

    async def complete(self, request: LLMRequest) -> LLMResponse:
        """Execute Gemini request and normalize output."""
        prompt_parts: list[str] = []
        if request.system:
            prompt_parts.append(f"System: {request.system}")
        for m in request.messages:
            prompt_parts.append(f"{m.role}: {m.content}")
        prompt = "\n".join(prompt_parts)

        try:
            resp = await self._client.generate_content_async(prompt)
        except Exception as exc:
            raise ProviderError(f"Gemini call failed: {exc}", provider_name="gemini") from exc

        text = getattr(resp, "text", "") or ""
        in_tok = self.count_tokens(prompt)
        out_tok = self.count_tokens(text)
        return LLMResponse(
            content=text,
            input_tokens=in_tok,
            output_tokens=out_tok,
            model=self._model,
            finish_reason="stop",
            raw=json.loads(resp.to_json()),
        )

    def count_tokens(self, text: str) -> int:
        """Approximate Gemini token count."""
        return len(text.encode("utf-8")) // 4

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

    @property
    def provider_name(self) -> str:
        return "gemini"

    @property
    def context_window(self) -> int:
        return 1_000_000 if "1.5" in self._model else 128_000
