from __future__ import annotations

"""Gemini provider implementation."""

import os
from typing import Any, AsyncIterator

import httpx

from agentrspy.core import (
    CompletionRequest,
    CompletionResponse,
    LlmProvider,
    Message,
    ProviderError,
    Role,
    StopReason,
    StreamChunk,
    Usage,
)


class GeminiProvider(LlmProvider):
    """Gemini provider."""

    def __init__(
        self, *, client: httpx.AsyncClient, api_key: str, base_url: str, model: str
    ) -> None:
        self.client = client
        self.api_key = api_key
        self.base_url = base_url
        self.model = model

    @classmethod
    def from_env(cls) -> "GeminiBuilder":
        return GeminiBuilder(api_key=os.getenv("GEMINI_API_KEY"))

    @classmethod
    def builder(cls) -> "GeminiBuilder":
        return GeminiBuilder()

    async def complete(self, request: CompletionRequest) -> CompletionResponse:
        contents: list[dict[str, Any]] = []
        if request.system is not None:
            contents.append(
                {
                    "role": "user",
                    "parts": [{"text": f"System instruction: {request.system}"}],
                }
            )
        for message in request.messages:
            contents.append(
                {
                    "role": "model" if message.role is Role.ASSISTANT else "user",
                    "parts": [{"text": message.text_content()}],
                }
            )
        try:
            response = await self.client.post(
                f"{self.base_url.rstrip('/')}/{request.model or self.model}:generateContent?key={self.api_key}",
                json={
                    "contents": contents,
                    "generationConfig": {
                        "temperature": request.temperature,
                        "maxOutputTokens": request.max_tokens,
                    },
                },
            )
        except httpx.HTTPError as error:
            raise ProviderError.http(str(error)) from error
        if response.status_code >= 400:
            raise ProviderError.api(response.status_code, response.text)
        try:
            payload = response.json()
        except ValueError as error:
            raise ProviderError.parse(str(error)) from error
        text = ""
        candidates = payload.get("candidates")
        if isinstance(candidates, list) and candidates:
            parts = (((candidates[0] or {}).get("content") or {}).get("parts")) or []
            for part in parts:
                if isinstance(part.get("text"), str):
                    text += part["text"]
        return CompletionResponse(
            message=Message.assistant(text),
            stop_reason=StopReason.STOP,
            usage=Usage(),
            model=request.model,
            raw=payload,
        )

    async def stream(self, request: CompletionRequest) -> AsyncIterator[StreamChunk]:
        del request
        raise ProviderError.unsupported(
            "Gemini streaming is not yet implemented in this SDK scaffold"
        )

    def name(self) -> str:
        return "gemini"


class GeminiBuilder:
    """Builder for `GeminiProvider`."""

    def __init__(
        self,
        *,
        api_key: str | None = None,
        base_url: str = "https://generativelanguage.googleapis.com/v1beta/models",
        model: str = "gemini-2.0-flash",
        client: httpx.AsyncClient | None = None,
    ) -> None:
        self._api_key = api_key
        self._base_url = base_url
        self._model = model
        self._client = client or httpx.AsyncClient()

    def api_key(self, api_key: str) -> "GeminiBuilder":
        self._api_key = api_key
        return self

    def model(self, model: str) -> "GeminiBuilder":
        self._model = model
        return self

    def base_url(self, base_url: str) -> "GeminiBuilder":
        self._base_url = base_url
        return self

    def build(self) -> GeminiProvider:
        if not self._api_key:
            raise ProviderError.missing_api_key("GEMINI_API_KEY")
        return GeminiProvider(
            client=self._client,
            api_key=self._api_key,
            base_url=self._base_url,
            model=self._model,
        )
