"""AWS Bedrock adapter (Converse API) with tool use."""

from __future__ import annotations

import asyncio
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 BedrockAdapter(BaseLLMAdapter):
    """Adapter for AWS Bedrock Converse API."""

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

        self._provider = "bedrock"
        self._model = llm_config.model
        self._timeout = int(llm_config.timeout)

        self._region = llm_config.aws_region
        self._profile = llm_config.aws_profile
        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
        self._extra_params = dict(llm_config.extra_params or {})

    @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 _convert_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
        out: list[dict[str, Any]] = []
        for m in messages:
            role = m.get("role")
            if role == "tool":
                # In our representation tool results become role=user blocks for bedrock.
                role = "user"
            content = m.get("content")
            if isinstance(content, list):
                # If already bedrock-shaped content blocks, keep.
                blocks = content
            else:
                blocks = [{"text": str(content or "")}]
            out.append({"role": role if role in {"user", "assistant"} else "user", "content": blocks})
        return out

    async def chat(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]], system: str | None = None, **kwargs: Any) -> Any:
        import boto3

        session = boto3.Session(profile_name=self._profile or None, region_name=self._region or None)
        client = session.client("bedrock-runtime")

        tool_config = {"tools": tools} if tools else None
        system_blocks = [{"text": system}] if system else None

        request: dict[str, Any] = {
            "modelId": self._model,
            "messages": self._convert_messages(messages),
            "inferenceConfig": {
                "maxTokens": self._max_tokens,
                "temperature": self._temperature,
            },
        }
        if self._top_p is not None:
            request["inferenceConfig"]["topP"] = self._top_p
        if tool_config:
            request["toolConfig"] = tool_config
        if system_blocks:
            request["system"] = system_blocks
        request.update(self._extra_params)

        # boto3 is blocking; run in a thread to avoid blocking the event loop.
        resp = await asyncio.to_thread(client.converse, **request)  # type: ignore[name-defined]
        return resp

    def extract_tool_calls(self, response: Any) -> list[ToolCall]:
        if not isinstance(response, dict):
            return []
        content = (
            response.get("output", {})
            .get("message", {})
            .get("content", [])
        )
        if not isinstance(content, list):
            return []
        out: list[ToolCall] = []
        for block in content:
            if not isinstance(block, dict):
                continue
            tu = block.get("toolUse") or block.get("tool_use")
            if not tu or not isinstance(tu, dict):
                continue
            tool_use_id = tu.get("toolUseId") or tu.get("tool_use_id") or ""
            name = tu.get("name") or ""
            args = tu.get("input") or tu.get("arguments") or {}
            if isinstance(args, str):
                try:
                    args = json.loads(args)
                except Exception:
                    args = {}
            if not isinstance(args, dict):
                args = {}
            if tool_use_id and name:
                out.append(ToolCall(id=str(tool_use_id), name=str(name), arguments=args))
        return out

    def extract_text(self, response: Any) -> str:
        if not isinstance(response, dict):
            return ""
        content = (
            response.get("output", {})
            .get("message", {})
            .get("content", [])
        )
        if not isinstance(content, list):
            return ""
        texts: list[str] = []
        for block in content:
            if not isinstance(block, dict):
                continue
            # Bedrock uses {'text': {'text': '...'}} sometimes.
            if "text" in block:
                txt_block = block.get("text")
                if isinstance(txt_block, str):
                    texts.append(txt_block)
                elif isinstance(txt_block, dict) and isinstance(txt_block.get("text"), str):
                    texts.append(txt_block["text"])
            elif "text" in block.get("content", {}):
                pass
        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
        out_msg = response.get("output", {}).get("message", {})
        content_blocks = out_msg.get("content") if isinstance(out_msg, dict) else None

        if content_blocks is not None:
            messages.append({"role": "assistant", "content": content_blocks})

        # Append tool results in a single user message.
        if tool_results:
            tool_result_blocks: list[dict[str, Any]] = []
            for tr in tool_results:
                tool_result_blocks.append(
                    {
                        "toolResult": {
                            "toolUseId": tr.tool_call_id,
                            "content": [{"text": tr.content}],
                        }
                    }
                )
            messages.append({"role": "user", "content": tool_result_blocks})

        return messages

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

