"""Google Gemini REST API client."""

from __future__ import annotations

import json
from collections.abc import AsyncIterator
from typing import Any

import httpx

from llm_rotator._types import LLMResponse, StreamChunk, ToolCall, Usage
from llm_rotator.clients import AbstractLLMClient
from llm_rotator.exceptions import KeyDeadError, ModelRateLimitError, ServerError

_BASE_URL = "https://generativelanguage.googleapis.com"


class GeminiClient(AbstractLLMClient):
    """Client for Google Gemini REST API."""

    def __init__(self, timeout: float = 60.0) -> None:
        self._timeout = timeout

    async def generate(
        self,
        messages: list[dict],
        model: str,
        api_key: str,
        base_url: str | None = None,
        **kwargs: Any,
    ) -> LLMResponse:
        api_base = base_url or _BASE_URL
        url = f"{api_base}/v1beta/models/{model}:generateContent"
        headers = {
            "x-goog-api-key": api_key,
            "content-type": "application/json",
        }
        body = _translate_messages(messages)
        _apply_tools_kwargs(body, kwargs)

        async with httpx.AsyncClient(timeout=self._timeout) as client:
            resp = await client.post(url, json=body, headers=headers)

        _check_status(resp, model=model, api_key=api_key)

        data = resp.json()
        return _parse_response(data, model=model, api_key=api_key)

    async def stream(
        self,
        messages: list[dict],
        model: str,
        api_key: str,
        base_url: str | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[StreamChunk]:
        api_base = base_url or _BASE_URL
        url = f"{api_base}/v1beta/models/{model}:streamGenerateContent"
        headers = {
            "x-goog-api-key": api_key,
            "content-type": "application/json",
        }
        body = _translate_messages(messages)
        _apply_tools_kwargs(body, kwargs)

        async with (
            httpx.AsyncClient(timeout=self._timeout) as client,
            client.stream(
                "POST", url, json=body, headers=headers, params={"alt": "sse"}
            ) as resp,
        ):
                if resp.status_code != 200:
                    await resp.aread()
                    _check_status(resp, model=model, api_key=api_key)

                async for line in resp.aiter_lines():
                    if not line.startswith("data: "):
                        continue
                    payload = line[6:].strip()
                    if not payload:
                        continue

                    data = json.loads(payload)
                    text = _extract_text(data)
                    usage = _extract_usage(data)

                    yield StreamChunk(
                        delta=text,
                        usage=usage,
                        done=False,
                    )


def _translate_messages(messages: list[dict]) -> dict:
    """Translate OpenAI message format to Gemini format."""
    body: dict = {}
    contents: list[dict] = []

    for msg in messages:
        role = msg["role"]
        content = msg["content"]

        if role == "system":
            body["systemInstruction"] = {"parts": [{"text": content}]}
        else:
            gemini_role = "model" if role == "assistant" else "user"
            contents.append({
                "role": gemini_role,
                "parts": [{"text": content}],
            })

    body["contents"] = contents
    return body


def _extract_text(data: dict) -> str:
    """Extract text from Gemini response."""
    candidates = data.get("candidates", [])
    if not candidates:
        return ""
    parts = candidates[0].get("content", {}).get("parts", [])
    texts = [p.get("text", "") for p in parts if "text" in p]
    return "".join(texts)


def _extract_usage(data: dict) -> Usage | None:
    """Extract usage from Gemini response."""
    meta = data.get("usageMetadata")
    if not meta:
        return None
    return Usage(
        prompt_tokens=meta.get("promptTokenCount", 0),
        completion_tokens=meta.get("candidatesTokenCount", 0),
        total_tokens=meta.get("totalTokenCount", 0),
    )


def _extract_tool_calls(data: dict) -> list[ToolCall] | None:
    """Extract function calls from Gemini response."""
    candidates = data.get("candidates", [])
    if not candidates:
        return None
    parts = candidates[0].get("content", {}).get("parts", [])
    calls = []
    for part in parts:
        fc = part.get("functionCall")
        if fc:
            calls.append(
                ToolCall(
                    id=fc.get("name", ""),  # Gemini has no call ID; use name
                    name=fc["name"],
                    arguments=json.dumps(fc.get("args", {})),
                )
            )
    return calls or None


def _parse_response(data: dict, *, model: str, api_key: str) -> LLMResponse:
    """Parse Gemini generateContent response."""
    text = _extract_text(data)
    usage = _extract_usage(data) or Usage(
        prompt_tokens=0, completion_tokens=0, total_tokens=0
    )
    tool_calls = _extract_tool_calls(data)
    return LLMResponse(
        content=text or None if tool_calls else text,
        usage=usage,
        model=model,
        provider="gemini",
        key_alias=api_key[:8] + "...",
        tool_calls=tool_calls,
        raw=data,
    )


def _apply_tools_kwargs(body: dict, kwargs: dict[str, Any]) -> None:
    """Apply tools, tool_choice, response_format from kwargs to Gemini body."""
    tools = kwargs.get("tools")
    if tools:
        # OpenAI format: [{"type": "function", "function": {...}}]
        # Gemini format: [{"functionDeclarations": [{name, description, parameters}]}]
        declarations = []
        for tool in tools:
            fn = tool.get("function", tool)
            declarations.append({
                "name": fn["name"],
                "description": fn.get("description", ""),
                "parameters": fn.get("parameters", {}),
            })
        body["tools"] = [{"functionDeclarations": declarations}]

    tool_choice = kwargs.get("tool_choice")
    if tool_choice is not None:
        config: dict[str, Any] = {}
        if tool_choice == "auto":
            config["mode"] = "AUTO"
        elif tool_choice == "none":
            config["mode"] = "NONE"
        elif tool_choice == "required":
            config["mode"] = "ANY"
        elif isinstance(tool_choice, dict) and tool_choice.get("function"):
            config["mode"] = "ANY"
            config["allowedFunctionNames"] = [tool_choice["function"]["name"]]
        body["toolConfig"] = {"functionCallingConfig": config}

    response_format = kwargs.get("response_format")
    if response_format is not None:
        gen_config = body.setdefault("generationConfig", {})
        if isinstance(response_format, dict):
            if response_format.get("type") == "json_object":
                gen_config["responseMimeType"] = "application/json"
        elif hasattr(response_format, "model_json_schema"):
            gen_config["responseMimeType"] = "application/json"
            gen_config["responseSchema"] = response_format.model_json_schema()


def _check_status(
    resp: httpx.Response, *, model: str, api_key: str
) -> None:
    """Raise typed exceptions based on HTTP status code."""
    if resp.status_code < 400:
        return

    alias = api_key[:8] + "..."
    status = resp.status_code

    if status in (401, 403, 402):
        raise KeyDeadError(key_alias=alias, status_code=status)
    if status == 429:
        retry_after_raw = resp.headers.get("retry-after")
        retry_after = int(retry_after_raw) if retry_after_raw else None
        raise ModelRateLimitError(
            key_alias=alias, model=model, retry_after=retry_after
        )
    if status >= 500:
        raise ServerError(provider="gemini", status_code=status)

    raise ServerError(
        provider="gemini",
        status_code=status,
        message=f"Unexpected HTTP {status}",
    )
