from __future__ import annotations

import logging
from typing import Any, Dict, List, Tuple

import httpx

from app.exceptions import ProviderAPICallError
from app.models import (
    LLMChoice,
    LLMMessage,
    LLMProvider,
    LLMRequest,
    LLMResponse,
    LLMRole,
)
from app.providers.base import LLMProviderClient


logger = logging.getLogger(__name__)


class GeminiProvider(LLMProviderClient):
    name = LLMProvider.GEMINI.value
    api_key_env = "GEMINI_API_KEY"
    base_url = "https://generativelanguage.googleapis.com/v1beta"

    async def chat(
        self, request: LLMRequest, client: httpx.AsyncClient
    ) -> LLMResponse:
        self.ensure_stream_supported(request.stream)

        system_instruction, contents = self._convert_messages(request.messages)
        config = {
            "temperature": request.temperature,
            "topP": request.top_p,
            "maxOutputTokens": request.max_tokens,
        }

        payload: Dict[str, Any] = self.merge_payload(
            {"contents": contents}, request.extra_params
        )

        generation_config = {k: v for k, v in config.items() if v is not None}
        if generation_config:
            payload["generationConfig"] = generation_config
        if system_instruction:
            payload["systemInstruction"] = {
                "role": "system",
                "parts": [{"text": system_instruction}],
            }

        endpoint = f"{self.base_url}/models/{request.model}:generateContent?key={self.api_key}"

        headers = {"Content-Type": "application/json"}

        try:
            response = await client.post(endpoint, json=payload, headers=headers)
            response.raise_for_status()
        except httpx.HTTPStatusError as exc:
            status_code = exc.response.status_code
            body = exc.response.text
            logger.error(
                "Gemini upstream returned %s: %s", status_code, body, exc_info=True
            )
            raise ProviderAPICallError(
                f"Gemini request failed with status {status_code}",
                status_code=status_code,
                response_text=body,
            ) from exc
        except httpx.HTTPError as exc:
            logger.error("Gemini transport error: %s", exc, exc_info=True)
            raise ProviderAPICallError(f"Gemini request failed: {exc}") from exc

        data: Dict[str, Any] = response.json()
        choices = self._build_choices(data.get("candidates", []))

        return LLMResponse(
            provider=LLMProvider.GEMINI,
            model=request.model,
            choices=choices,
            raw=data,
        )

    @staticmethod
    def _convert_messages(
        messages: List[LLMMessage],
    ) -> Tuple[str | None, List[dict[str, Any]]]:
        system_parts: List[str] = []
        contents: List[dict[str, Any]] = []

        for msg in messages:
            if msg.role == LLMRole.SYSTEM:
                system_parts.append(msg.content)
                continue

            role = "user" if msg.role == LLMRole.USER else "model"
            contents.append({"role": role, "parts": [{"text": msg.content}]})

        system_instruction = "\n\n".join(system_parts) if system_parts else None
        return system_instruction, contents

    @staticmethod
    def _build_choices(candidates: List[dict[str, Any]]) -> List[LLMChoice]:
        choices: List[LLMChoice] = []
        for idx, candidate in enumerate(candidates):
            content = candidate.get("content", {})
            parts = content.get("parts", [])
            text_parts = [
                part.get("text", "")
                for part in parts
                if isinstance(part, dict) and part.get("text")
            ]
            text = "\n\n".join(text_parts)
            choices.append(
                LLMChoice(
                    index=candidate.get("index", idx),
                    message=LLMMessage(role="assistant", content=text),
                )
            )
        if not choices:
            choices.append(
                LLMChoice(
                    index=0,
                    message=LLMMessage(role="assistant", content=""),
                )
            )
        return choices