"""
Google Gemini provider adapter.

Translates between the unified OpenAI-compatible format and the
Google Generative Language (Gemini) API format.

- Auth: API key as query param (?key=...) — Google AI uses key, not Bearer.
- Request: messages -> contents[] with role/parts[].text; system -> systemInstruction.
- Response: candidates[].content.parts[].text, usageMetadata.
- Streaming: streamGenerateContent with alt=sse; chunks are full candidate objects.
"""

from __future__ import annotations

import json
import urllib.parse
from typing import Any, Dict, List, Tuple

from proxy.adapters.base import BaseProviderAdapter
from proxy.pricing import GEMINI_PRICING, ZERO_PRICING
from proxy.schemas import (
    ChatCompletionChunk,
    ChatCompletionRequest,
    ChatCompletionResponse,
    Choice,
    ChoiceMessage,
    StreamChoice,
    StreamDelta,
    TokenPricing,
    Usage,
)


def _role_to_gemini(role: str) -> str:
    """Map OpenAI role to Gemini role (user -> user, assistant -> model)."""
    if role == "assistant":
        return "model"
    return role


class GeminiAdapter(BaseProviderAdapter):
    """Adapter for the Google Gemini (Generative Language) API.

    Uses the REST API at generativelanguage.googleapis.com/v1beta.
    API key is passed as query parameter (required by Google AI).
    """

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

    # ----- request translation -----

    def translate_request(
        self,
        request: ChatCompletionRequest,
        api_key: str,
    ) -> Tuple[str, Dict[str, str], Dict[str, Any]]:
        # Gemini uses :generateContent or :streamGenerateContent with ?key=API_KEY
        # Streaming requires alt=sse for Server-Sent Events format
        base = self._base_url.rstrip("/")
        model = request.model
        action = "streamGenerateContent" if request.stream else "generateContent"
        path = f"/models/{model}:{action}"
        params = [("key", api_key)]
        if request.stream:
            params.append(("alt", "sse"))
        query = urllib.parse.urlencode(params)
        url = f"{base}{path}?{query}"

        headers: Dict[str, str] = {
            "Content-Type": "application/json",
        }

        # Extract system prompt and build contents
        system_instruction: str | None = None
        contents: List[Dict[str, Any]] = []
        for msg in request.messages:
            if msg.role == "system":
                system_instruction = msg.content or ""
            else:
                role = _role_to_gemini(msg.role)
                text = msg.content or ""
                contents.append(
                    {
                        "role": role,
                        "parts": [{"text": text}],
                    }
                )

        body: Dict[str, Any] = {
            "contents": contents,
            "generationConfig": {
                "maxOutputTokens": request.max_tokens
                or request.max_completion_tokens
                or 8192,
            },
        }
        if system_instruction is not None:
            body["systemInstruction"] = {
                "parts": [{"text": system_instruction}],
            }
        if request.temperature is not None:
            body["generationConfig"]["temperature"] = request.temperature
        if request.top_p is not None:
            body["generationConfig"]["topP"] = request.top_p
        if request.stop is not None:
            stop_list = (
                request.stop
                if isinstance(request.stop, list)
                else [request.stop]
            )
            body["generationConfig"]["stopSequences"] = stop_list

        return url, headers, body

    # ----- response translation -----

    def translate_response(
        self,
        status_code: int,
        headers: Dict[str, str],
        body: Dict[str, Any],
    ) -> ChatCompletionResponse:
        # Gemini: candidates[].content.parts[].text, usageMetadata
        candidates = body.get("candidates", [])
        text_parts: List[str] = []
        finish_reason = "stop"
        if candidates:
            c0 = candidates[0]
            content = c0.get("content", {})
            for part in content.get("parts", []):
                if "text" in part:
                    text_parts.append(part["text"])
            finish_reason = c0.get("finishReason", "STOP").lower()
            if finish_reason == "stop":
                finish_reason = "stop"
            elif finish_reason == "max_tokens":
                finish_reason = "length"

        usage_meta = body.get("usageMetadata", {})
        prompt_tokens = usage_meta.get("promptTokenCount", 0)
        completion_tokens = usage_meta.get("candidatesTokenCount", 0)

        return ChatCompletionResponse(
            id=body.get("id", ""),
            object="chat.completion",
            created=0,
            model=body.get("model", ""),
            choices=[
                Choice(
                    index=0,
                    message=ChoiceMessage(
                        role="assistant",
                        content="\n".join(text_parts) if text_parts else None,
                    ),
                    finish_reason=finish_reason,
                )
            ],
            usage=Usage(
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                total_tokens=prompt_tokens + completion_tokens,
            ),
        )

    # ----- streaming -----

    def parse_stream_chunk(self, raw_line: str) -> ChatCompletionChunk | None:
        line = raw_line.strip()
        if not line:
            return None

        try:
            data = json.loads(line)
        except json.JSONDecodeError:
            return None

        candidates = data.get("candidates", [])
        if not candidates:
            return None

        c0 = candidates[0]
        content = c0.get("content", {})
        parts = content.get("parts", [])
        text = ""
        for part in parts:
            if "text" in part:
                text += part.get("text", "")

        finish = c0.get("finishReason")
        finish_reason = None
        if finish:
            finish_reason = "stop" if finish == "STOP" else "length"

        if not text and not finish_reason:
            return None

        return ChatCompletionChunk(
            id="",
            object="chat.completion.chunk",
            created=0,
            model=data.get("model", ""),
            choices=[
                StreamChoice(
                    index=0,
                    delta=StreamDelta(content=text if text else None),
                    finish_reason=finish_reason,
                )
            ],
        )

    # ----- pricing -----

    def get_pricing(self, model: str) -> TokenPricing:
        return GEMINI_PRICING.get(model, ZERO_PRICING)
