"""Google GenerateContent API client mixin."""

from collections.abc import AsyncIterator
from typing import Any, ClassVar

from celeste.client import APIMixin
from celeste.core import UsageField
from celeste.io import FinishReason

from ..auth import GoogleADC
from . import config


class GoogleGenerateContentClient(APIMixin):
    """Mixin for GenerateContent API capabilities.

    Provides shared implementation for all capabilities using the GenerateContent API:
    - _make_request() - HTTP POST to generateContent
    - _make_stream_request() - HTTP streaming to streamGenerateContent
    - _parse_usage() - Extract usage dict from usageMetadata
    - _parse_content() - Extract parts array from first candidate
    - _parse_finish_reason() - Extract finish reason string from candidates
    - _content_fields: ClassVar - Content field names to exclude from metadata

    Auth-based endpoint selection:
    - GoogleADC auth -> Vertex AI endpoints
    - API key auth -> Gemini API endpoints

    Capability clients extend parsing methods via super() to wrap/transform results.

    Usage:
        class GoogleTextGenerationClient(GoogleGenerateContentClient, TextGenerationClient):
            def _parse_content(self, response_data):
                parts = super()._parse_content(response_data)
                text = parts[0].get("text") or ""
                return text
    """

    _content_fields: ClassVar[set[str]] = {"candidates"}

    def _get_vertex_endpoint(self, gemini_endpoint: str) -> str:
        """Map Gemini endpoint to Vertex AI endpoint."""
        mapping: dict[str, str] = {
            config.GoogleGenerateContentEndpoint.GENERATE_CONTENT: config.VertexGenerateContentEndpoint.GENERATE_CONTENT,
            config.GoogleGenerateContentEndpoint.STREAM_GENERATE_CONTENT: config.VertexGenerateContentEndpoint.STREAM_GENERATE_CONTENT,
            config.GoogleGenerateContentEndpoint.COUNT_TOKENS: config.VertexGenerateContentEndpoint.COUNT_TOKENS,
        }
        vertex_endpoint = mapping.get(gemini_endpoint)
        if vertex_endpoint is None:
            raise ValueError(f"No Vertex AI endpoint mapping for: {gemini_endpoint}")
        return vertex_endpoint

    def _build_url(self, endpoint: str) -> str:
        """Build full URL based on auth type."""
        if isinstance(self.auth, GoogleADC):
            return self.auth.build_url(
                self._get_vertex_endpoint(endpoint), model_id=self.model.id
            )
        return f"{config.BASE_URL}{endpoint.format(model_id=self.model.id)}"

    async def _make_request(
        self,
        request_body: dict[str, Any],
        *,
        endpoint: str | None = None,
        extra_headers: dict[str, str] | None = None,
        **parameters: Any,
    ) -> dict[str, Any]:
        """Make HTTP request to generateContent endpoint."""
        if endpoint is None:
            endpoint = config.GoogleGenerateContentEndpoint.GENERATE_CONTENT

        headers = self._json_headers(extra_headers)
        response = await self.http_client.post(
            url=self._build_url(endpoint),
            headers=headers,
            json_body=request_body,
        )
        self._handle_error_response(response)
        data: dict[str, Any] = response.json()
        return data

    def _make_stream_request(
        self,
        request_body: dict[str, Any],
        *,
        endpoint: str | None = None,
        extra_headers: dict[str, str] | None = None,
        **parameters: Any,
    ) -> AsyncIterator[dict[str, Any]]:
        """Make streaming request to streamGenerateContent endpoint."""
        if endpoint is None:
            endpoint = config.GoogleGenerateContentEndpoint.STREAM_GENERATE_CONTENT

        headers = self._json_headers(extra_headers)
        return self.http_client.stream_post(
            url=self._build_url(endpoint),
            headers=headers,
            json_body=request_body,
        )

    @staticmethod
    def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | float | None]:
        """Map Google Gemini usage fields to unified names.

        Shared by client and streaming across all capabilities.
        """
        return {
            UsageField.INPUT_TOKENS: usage_data.get("promptTokenCount"),
            UsageField.OUTPUT_TOKENS: usage_data.get("candidatesTokenCount"),
            UsageField.TOTAL_TOKENS: usage_data.get("totalTokenCount"),
            UsageField.REASONING_TOKENS: usage_data.get("thoughtsTokenCount"),
        }

    def _parse_usage(
        self, response_data: dict[str, Any]
    ) -> dict[str, int | float | None]:
        """Extract usage data from Gemini usageMetadata."""
        usage_metadata = response_data.get("usageMetadata", {})
        return GoogleGenerateContentClient.map_usage_fields(usage_metadata)

    def _parse_content(self, response_data: dict[str, Any]) -> Any:
        """Return all candidates from response.

        Returns list of candidate objects that capability clients extract content from.
        """
        candidates = response_data.get("candidates", [])
        if not candidates:
            msg = "No candidates in response"
            raise ValueError(msg)
        return candidates

    def _parse_finish_reason(self, response_data: dict[str, Any]) -> FinishReason:
        """Extract finish reason from Gemini candidates.

        Returns FinishReason that capability clients wrap in their specific type.
        """
        candidates = response_data.get("candidates", [])
        if not candidates:
            reason = None
        else:
            reason = candidates[0].get("finishReason")
        return FinishReason(reason=reason)


__all__ = ["GoogleGenerateContentClient"]
