"""Cohere adapter for chat + tool 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 CohereAdapter(BaseLLMAdapter):
    """Adapter for Cohere chat with tool calling.

    Note: Cohere's API shape differs from OpenAI; this adapter provides
    best-effort request/response handling.
    """

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

        self._provider = "cohere"
        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._stop_sequences = llm_config.stop_sequences

    @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:
        # Cohere uses system prompt as part of message content (best-effort).
        return False

    def _last_user_text(self, messages: list[dict[str, Any]]) -> str:
        for m in reversed(messages):
            if m.get("role") == "user":
                content = m.get("content", "")
                return content if isinstance(content, str) else str(content)
        return ""

    def _extract_tool_results_from_messages(self, messages: list[dict[str, Any]]) -> list[ToolResult]:
        tool_results: list[ToolResult] = []
        for m in messages:
            if m.get("role") == "tool":
                tid = str(m.get("tool_call_id") or m.get("id") or "")
                name = str(m.get("name") or m.get("tool_name") or "")
                content = m.get("content") or ""
                if not isinstance(content, str):
                    content = str(content)
                tool_results.append(ToolResult(tool_call_id=tid, name=name, content=content, is_error=False))
        return tool_results

    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("Cohere API key missing.")

        headers = {"Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json"}
        last_user = self._last_user_text(messages)

        # Best-effort: treat `system` as prefix to the last user message.
        if system:
            last_user = f"{system}\n\n{last_user}" if last_user else system

        tool_results = self._extract_tool_results_from_messages(messages)
        tool_results_payload = []
        if tool_results:
            for tr in tool_results:
                tool_results_payload.append({"call": {"id": tr.tool_call_id, "name": tr.name}, "outputs": [tr.content]})

        payload: dict[str, Any] = {
            "model": self._model,
            "message": last_user,
            "max_tokens": self._max_tokens,
            "temperature": self._temperature,
            "stream": False,
        }
        if self._top_p is not None:
            payload["p"] = self._top_p
        if self._stop_sequences:
            payload["stop_sequences"] = self._stop_sequences
        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = "auto"
        if tool_results_payload:
            payload["tool_results"] = tool_results_payload

        url = "https://api.cohere.ai/v1/chat"
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            resp = await client.post(url, headers=headers, json=payload)
            resp.raise_for_status()
            return resp.json()

    def extract_tool_calls(self, response: Any) -> list[ToolCall]:
        if not isinstance(response, dict):
            return []
        raw = response.get("tool_calls") or response.get("toolCalls") or []
        if not isinstance(raw, list):
            return []
        out: list[ToolCall] = []
        for call in raw:
            if not isinstance(call, dict):
                continue
            call_id = str(call.get("id") or call.get("call_id") or "")
            name = str(call.get("name") or call.get("tool_name") or "")
            args = call.get("arguments") or call.get("parameters") or {}
            if isinstance(args, str):
                try:
                    args = json.loads(args)
                except Exception:
                    args = {}
            if not isinstance(args, dict):
                args = {}
            if call_id and name:
                out.append(ToolCall(id=call_id, name=name, arguments=args))
        return out

    def extract_text(self, response: Any) -> str:
        if not isinstance(response, dict):
            return ""
        text = response.get("text") or response.get("message") or response.get("output") or ""
        if isinstance(text, str):
            return text.strip()
        if isinstance(text, dict) and isinstance(text.get("content"), str):
            return str(text["content"]).strip()
        return str(text).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]]:
        # Append assistant message with any content we have, then tool outputs.
        if isinstance(response, dict) and (response.get("text") or response.get("output")):
            messages.append({"role": "assistant", "content": self.extract_text(response)})
        else:
            messages.append({"role": "assistant", "content": ""})

        for tr in tool_results:
            messages.append({"role": "tool", "tool_call_id": tr.tool_call_id, "name": tr.name, "content": tr.content})
        return messages

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

