"""Google (Gemini) provider client."""

import json
import re
import time
import uuid
from collections.abc import AsyncGenerator
from typing import Any

import httpx
import structlog

from robot_resources.proxy.models import (
    ChatCompletionChunk,
    ChatCompletionRequest,
    ChatCompletionResponse,
    Choice,
    Message,
    StreamChoice,
    StreamDelta,
    Usage,
)
from robot_resources.proxy.providers.base import BaseProvider

logger = structlog.get_logger()


def _redact_key(url: str) -> str:
    """Redact API key from Google API URLs for safe logging."""
    return re.sub(r"key=[^&]+", "key=REDACTED", url)


# Model name mappings
MODEL_MAPPING = {
    "gemini-2.5-pro": "gemini-2.5-pro",
    "gemini-2.5-flash": "gemini-2.5-flash",
    "gemini-2.5-flash-lite": "gemini-2.5-flash-lite",
    # Legacy aliases — route to current models
    "gemini-pro": "gemini-2.5-pro",
    "gemini-flash": "gemini-2.5-flash-lite",
    "gemini-1.5-pro": "gemini-2.5-pro",
    "gemini-1.5-flash": "gemini-2.5-flash-lite",
    "gemini-1.5-flash-8b": "gemini-2.5-flash-lite",
    "gemini-2.0-flash": "gemini-2.5-flash",
    "gemini-2.0-flash-exp": "gemini-2.5-flash",
    "gemini-2.0-flash-lite": "gemini-2.5-flash-lite",
}


class GoogleProvider(BaseProvider):
    """Google Gemini API provider."""

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

    @property
    def base_url(self) -> str:
        return "https://generativelanguage.googleapis.com/v1beta"

    def _get_headers(self, *, api_key: str | None = None) -> dict:
        return {
            "content-type": "application/json",
        }

    def _get_completion_endpoint(self) -> str:
        # Note: Model is in URL for Google
        return ""

    def map_model(self, model: str) -> str:
        """Map model name to Google-specific name."""
        return MODEL_MAPPING.get(model, model)

    @staticmethod
    def _safe_raise_for_status(response: httpx.Response) -> None:
        """Check response status, redacting API keys from any error.

        Google API URLs contain ``?key=...`` query parameters. If the
        request fails, ``raise_for_status()`` embeds the full URL —
        including the key — in the exception message.  This method
        catches that and re-raises with the key redacted.
        """
        try:
            response.raise_for_status()
        except httpx.HTTPStatusError as exc:
            redacted_msg = _redact_key(str(exc))
            logger.warning(
                "google_api_error",
                status=exc.response.status_code,
                url=_redact_key(str(exc.request.url)),
            )
            raise httpx.HTTPStatusError(
                message=redacted_msg,
                request=exc.request,
                response=exc.response,
            ) from None

    def format_request(self, request: ChatCompletionRequest) -> dict[str, Any]:
        """Format request for Google Gemini API."""
        # Convert messages to Gemini format
        contents: list[dict[str, Any]] = []
        system_instruction = None

        for msg in request.messages:
            if msg.role == "system":
                system_instruction = msg.content
            else:
                # Map role: user -> user, assistant -> model
                role = "model" if msg.role == "assistant" else "user"
                contents.append(
                    {
                        "role": role,
                        "parts": [{"text": msg.content}],
                    }
                )

        generation_config: dict[str, Any] = {}
        if request.max_tokens:
            generation_config["maxOutputTokens"] = request.max_tokens
        if request.temperature != 1.0:
            generation_config["temperature"] = request.temperature

        formatted: dict[str, Any] = {
            "contents": contents,
            "generationConfig": generation_config,
        }

        if system_instruction:
            formatted["systemInstruction"] = {"parts": [{"text": system_instruction}]}

        return formatted

    async def complete(
        self, request: ChatCompletionRequest, *, api_key: str | None = None,
    ) -> ChatCompletionResponse:
        """Execute a non-streaming chat completion.

        Overrides base ``complete()`` to build the Google API URL with the
        correct mapped model name.  Follows the same pattern as ``stream()``
        which already constructs the URL correctly.
        """
        formatted = self.format_request(request)
        mapped_model = self.map_model(request.model)
        key = api_key or self.config.api_key
        url = (
            f"{self.config.base_url or self.base_url}"
            f"/models/{mapped_model}:generateContent"
            f"?key={key}"
        )

        async with self._create_client() as client:
            response = await client.post(url, json=formatted)
            self._safe_raise_for_status(response)
            result = response.json()

        return self.parse_response(result)

    def parse_response(self, response: dict) -> ChatCompletionResponse:
        """Parse Google response to standard format."""
        candidates = response.get("candidates", [])
        if not candidates:
            return ChatCompletionResponse(
                model="gemini",
                choices=[
                    Choice(
                        message=Message(role="assistant", content=""),
                        finish_reason="stop",
                    )
                ],
                usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
            )

        candidate = candidates[0]
        content_parts = candidate.get("content", {}).get("parts", [])
        content = "".join(part.get("text", "") for part in content_parts)

        usage_meta = response.get("usageMetadata", {})

        return ChatCompletionResponse(
            model="gemini",
            choices=[
                Choice(
                    index=0,
                    message=Message(role="assistant", content=content),
                    finish_reason=self._map_finish_reason(candidate.get("finishReason")),
                )
            ],
            usage=Usage(
                prompt_tokens=usage_meta.get("promptTokenCount", 0),
                completion_tokens=usage_meta.get("candidatesTokenCount", 0),
                total_tokens=usage_meta.get("totalTokenCount", 0),
            ),
        )

    def _map_finish_reason(self, reason: str | None) -> str:
        """Map Google finish reason to OpenAI format."""
        mapping = {
            "STOP": "stop",
            "MAX_TOKENS": "length",
            "SAFETY": "content_filter",
            "RECITATION": "content_filter",
        }
        return mapping.get(reason or "", "stop")

    async def stream(
        self, request: ChatCompletionRequest, *, api_key: str | None = None,
    ) -> AsyncGenerator[str, None]:
        """Stream a chat completion from Google Gemini.

        Google uses ``streamGenerateContent?alt=sse`` which returns SSE
        with Gemini-format JSON.  We convert each chunk to an
        OpenAI-compatible ``ChatCompletionChunk``.
        """
        formatted = self.format_request(request)
        mapped_model = self.map_model(request.model)
        key = api_key or self.config.api_key
        url = (
            f"{self.config.base_url or self.base_url}"
            f"/models/{mapped_model}:streamGenerateContent"
            f"?alt=sse&key={key}"
        )
        self._last_stream_usage = None

        chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
        created = int(time.time())
        first_chunk = True
        prompt_tokens = 0
        completion_tokens = 0

        async with self._create_client() as client:
            async with client.stream("POST", url, json=formatted) as response:
                self._safe_raise_for_status(response)

                async for line in response.aiter_lines():
                    line = line.strip()
                    if not line or not line.startswith("data: "):
                        continue

                    raw_data = line[len("data: "):]
                    try:
                        data = json.loads(raw_data)
                    except json.JSONDecodeError:
                        logger.warning(
                            "stream_json_decode_error",
                            provider="google",
                            raw_line=raw_data[:200],
                        )
                        continue

                    # Emit initial role chunk on first data
                    if first_chunk:
                        role_chunk = ChatCompletionChunk(
                            id=chunk_id,
                            created=created,
                            model=mapped_model,
                            choices=[
                                StreamChoice(
                                    delta=StreamDelta(role="assistant", content=""),
                                    finish_reason=None,
                                )
                            ],
                        )
                        yield f"data: {role_chunk.model_dump_json()}\n\n"
                        first_chunk = False

                    # Extract text from candidates
                    candidates = data.get("candidates", [])
                    if candidates:
                        parts = candidates[0].get("content", {}).get("parts", [])
                        text = "".join(p.get("text", "") for p in parts)

                        finish_reason_raw = candidates[0].get("finishReason")
                        finish_reason = None
                        if finish_reason_raw and finish_reason_raw != "STOP":
                            finish_reason = self._map_finish_reason(finish_reason_raw)

                        if text:
                            chunk = ChatCompletionChunk(
                                id=chunk_id,
                                created=created,
                                model=mapped_model,
                                choices=[
                                    StreamChoice(
                                        delta=StreamDelta(content=text),
                                        finish_reason=finish_reason,
                                    )
                                ],
                            )
                            yield f"data: {chunk.model_dump_json()}\n\n"

                    # Extract usage metadata (typically in last chunk)
                    usage_meta = data.get("usageMetadata")
                    if usage_meta:
                        prompt_tokens = usage_meta.get("promptTokenCount", 0)
                        completion_tokens = usage_meta.get("candidatesTokenCount", 0)

        # Emit final stop chunk and [DONE]
        final_chunk = ChatCompletionChunk(
            id=chunk_id,
            created=created,
            model=mapped_model,
            choices=[
                StreamChoice(
                    delta=StreamDelta(),
                    finish_reason="stop",
                )
            ],
        )
        yield f"data: {final_chunk.model_dump_json()}\n\n"

        self._last_stream_usage = Usage(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=prompt_tokens + completion_tokens,
        )
        yield "data: [DONE]\n\n"
