"""Google Gemini adapter with function calling."""

from __future__ import annotations

import httpx
import json
import logging
from typing import Any

from mcpbridge.adapters.base import BaseLLMAdapter, ToolCall, ToolResult
from mcpbridge.config import LLMConfig
from mcpbridge.utils.token_counter import estimate_messages_tokens

logger = logging.getLogger("mcpbridge")


class GeminiAdapter(BaseLLMAdapter):
    """Adapter for Gemini generateContent endpoint with function calling."""

    def __init__(self, llm_config: LLMConfig) -> None:
        try:
            import google.generativeai as _genai  # noqa: F401
        except ImportError as e:
            raise ImportError("Missing `google-generativeai` dependency. Install with: pip install mcpbridge[gemini]") from e

        self._provider = "gemini"
        self._model = llm_config.model
        self._api_key = llm_config.api_key
        self._timeout = int(llm_config.timeout)

        self._temperature = float(llm_config.temperature)
        self._max_tokens = int(llm_config.max_tokens)
        self._top_p = llm_config.top_p
        self._top_k = llm_config.top_k

    @property
    def provider_name(self) -> str:
        return self._provider

    @property
    def supports_parallel_tool_calls(self) -> bool:
        return True

    @property
    def supports_system_prompt(self) -> bool:
        return True

    def _normalize_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Convert our generic messages into Gemini `contents` format."""
        contents: list[dict[str, Any]] = []
        for m in messages:
            role = m.get("role")
            if role == "assistant":
                gem_role = "model"
            else:
                gem_role = "user"

            if "parts" in m and isinstance(m.get("parts"), list):
                parts = m["parts"]
            else:
                content = m.get("content", "")
                if isinstance(content, str):
                    parts = [{"text": content}]
                elif isinstance(content, list):
                    # Best-effort pass-through.
                    parts = content
                else:
                    parts = [{"text": str(content)}]

            contents.append({"role": gem_role, "parts": parts})
        return contents

    async def chat(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]], system: str | None = None, **kwargs: Any) -> Any:
        if not self._api_key:
            raise PermissionError("Gemini API key missing.")

        url = f"https://generativelanguage.googleapis.com/v1beta/models/{self._model}:generateContent"
        params = {"key": self._api_key}

        contents = self._normalize_messages(messages)
        payload: dict[str, Any] = {
            "contents": contents,
            "generationConfig": {
                "temperature": self._temperature,
                "maxOutputTokens": self._max_tokens,
            },
        }
        if system:
            payload["system_instruction"] = {"parts": [{"text": system}]}

        if tools:
            payload["tools"] = tools

        if self._top_p is not None:
            payload["generationConfig"]["topP"] = self._top_p
        if self._top_k is not None:
            payload["generationConfig"]["topK"] = self._top_k

        async with httpx.AsyncClient(timeout=self._timeout) as client:
            resp = await client.post(url, params=params, json=payload)
            resp.raise_for_status()
            return resp.json()

    def extract_tool_calls(self, response: Any) -> list[ToolCall]:
        candidates = response.get("candidates") if isinstance(response, dict) else None
        if not candidates:
            return []
        parts = (
            candidates[0]
            .get("content", {})
            .get("parts", [])
            if isinstance(candidates[0], dict)
            else []
        )
        tool_calls: list[ToolCall] = []
        for idx, part in enumerate(parts):
            if not isinstance(part, dict):
                continue
            fc = part.get("function_call") or part.get("functionCall")
            if not fc:
                continue
            name = fc.get("name") or ""
            # Gemini may omit `id`/`tool_call_id` for function calls (e.g. some
            # model variants). The agent loop requires a stable id, so we
            # generate a fallback using the part index.
            call_id = fc.get("id") or fc.get("tool_call_id") or str(idx)
            args = fc.get("args") or fc.get("arguments") or {}
            if isinstance(args, str):
                try:
                    args = json.loads(args)
                except Exception:
                    args = {}
            if not isinstance(args, dict):
                args = {}
            if name:
                tool_calls.append(ToolCall(id=str(call_id), name=str(name), arguments=args))
        return tool_calls

    def extract_text(self, response: Any) -> str:
        if not isinstance(response, dict):
            return ""
        candidates = response.get("candidates") or []
        if not candidates or not isinstance(candidates[0], dict):
            return ""
        parts = candidates[0].get("content", {}).get("parts") or []
        texts: list[str] = []
        for p in parts:
            if isinstance(p, dict) and isinstance(p.get("text"), str):
                texts.append(p["text"])
        return "".join(texts).strip()

    def is_done(self, response: Any) -> bool:
        return len(self.extract_tool_calls(response)) == 0

    def append_tool_results(self, messages: list[dict[str, Any]], response: Any, tool_results: list[ToolResult]) -> list[dict[str, Any]]:
        if not isinstance(response, dict):
            return messages
        # Preserve text preamble (if any).
        preamble = self.extract_text(response)
        if preamble:
            messages.append({"role": "assistant", "parts": [{"text": preamble}]})

        # Add function responses as a user message.
        parts: list[dict[str, Any]] = []
        for tr in tool_results:
            parts.append(
                {
                    "function_response": {
                        "name": tr.name,
                        "id": tr.tool_call_id,
                        "response": {"output": tr.content},
                    }
                }
            )
        if parts:
            messages.append({"role": "user", "parts": parts})
        return messages

    def count_tokens(self, messages: list[dict[str, Any]], system: str) -> int:
        return estimate_messages_tokens(messages, system, provider=self._provider)

