"""Google Gemini provider with native multimodal inputs."""

from __future__ import annotations

import os
from typing import Any, AsyncIterator

import httpx

from logos_adk.errors import AuthenticationError, ProviderError
from logos_adk.providers import Provider, provider_registry
from logos_adk.providers.multimodal import gemini_message, parse_gemini_parts
from logos_adk.types import MediaChunk, Message, Response, TextChunk, Usage


class GeminiProvider(Provider):
    """Provider for Gemini generateContent API."""

    _http_client_pool: dict[str | None, httpx.AsyncClient] = {}

    @classmethod
    def _shared_http_client(cls, proxy: str | None) -> httpx.AsyncClient:
        if proxy not in cls._http_client_pool:
            kwargs: dict[str, Any] = {"timeout": 120.0}
            if proxy:
                kwargs["proxy"] = proxy
            cls._http_client_pool[proxy] = httpx.AsyncClient(**kwargs)
        return cls._http_client_pool[proxy]

    def __init__(
        self,
        model: str = "gemini-2.0-flash",
        *,
        api_key: str | None = None,
        base_url: str | None = None,
        proxy: str | None = None,
        client: httpx.AsyncClient | None = None,
    ) -> None:
        resolved_key = api_key or os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
        if not resolved_key:
            raise AuthenticationError(
                "Gemini API key not found. Set GEMINI_API_KEY or GOOGLE_API_KEY, or pass api_key=."
            )
        self._api_key = resolved_key
        self._model = model
        self._base_url = (base_url or "https://generativelanguage.googleapis.com/v1beta").rstrip("/")
        self.proxy = proxy
        self._client = client

    def _payload(
        self,
        messages: list[Message],
        *,
        tools: list[dict] | None,
        output_schema: Any,
        **kwargs: Any,
    ) -> dict[str, Any]:
        system_instruction = None
        contents: list[dict[str, Any]] = []
        for message in messages:
            if message.role == "system":
                system_instruction = {"parts": [{"text": message.content}]}
                continue
            contents.append(gemini_message(message))

        payload: dict[str, Any] = {"contents": contents}
        if system_instruction is not None:
            payload["systemInstruction"] = system_instruction
        if tools:
            payload["tools"] = [
                {
                    "functionDeclarations": [
                        {
                            "name": tool["name"],
                            "description": tool.get("description", ""),
                            "parameters": tool.get("parameters", {"type": "object", "properties": {}}),
                        }
                    ]
                }
                for tool in tools
            ]
        generation_config = dict(kwargs.pop("generationConfig", {}) or {})
        if output_schema is not None:
            generation_config.setdefault("responseMimeType", "application/json")
        if generation_config:
            payload["generationConfig"] = generation_config
        payload.update(kwargs)
        return payload

    async def _post(self, payload: dict[str, Any]) -> dict[str, Any]:
        owns_client = self._client is None
        client = self._client or self._shared_http_client(self.proxy)
        try:
            response = await client.post(
                f"{self._base_url}/models/{self._model}:generateContent",
                params={"key": self._api_key},
                json=payload,
            )
            response.raise_for_status()
            return response.json()
        except httpx.HTTPError as exc:
            raise ProviderError(f"Gemini API error: {exc}") from exc
        finally:
            if owns_client and self._client is not None:
                await client.aclose()

    def _parse(self, data: dict[str, Any]) -> Response:
        candidates = data.get("candidates") or []
        first = candidates[0] if candidates else {}
        content = first.get("content") or {}
        text, parts, tool_calls = parse_gemini_parts(content.get("parts") or [])
        usage = data.get("usageMetadata") or {}
        return Response(
            content=text,
            tool_calls=tool_calls,
            usage=Usage(
                input_tokens=int(usage.get("promptTokenCount") or 0),
                output_tokens=int(usage.get("candidatesTokenCount") or 0),
            ),
            raw=data,
            content_parts=parts,
        )

    async def generate(
        self,
        messages: list[Message],
        tools: list[dict] | None = None,
        output_schema: Any = None,
        **kwargs: Any,
    ) -> Response:
        payload = self._payload(messages, tools=tools, output_schema=output_schema, **kwargs)
        data = await self._post(payload)
        return self._parse(data)

    async def stream(
        self,
        messages: list[Message],
        tools: list[dict] | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[Any]:
        response = await self.generate(messages, tools=tools, **kwargs)
        if response.content:
            yield TextChunk(text=response.content)
        for part in response.content_parts:
            if part.type != "text":
                yield MediaChunk(part=part)
        yield response


def _gemini_factory(model_name: str) -> GeminiProvider:
    return GeminiProvider(model=model_name)


provider_registry.register("gemini-*", _gemini_factory)

