import json
from typing import Any

from anthropic.types import Message, ThinkingBlock, ToolUseBlock
from sinapsis_anthropic.helpers.anthropic_keys import AnthropicKeys
from sinapsis_anthropic.helpers.mcp_tool_helper import make_tools_anthropic_compatible
from sinapsis_anthropic.templates.anthropic_text_generation import AnthropicAttributes, AnthropicTextGeneration
from sinapsis_chatbots_base.helpers.llm_conversation import (
    build_chat_message,
    build_packet_message,
    combine_assistant_responses,
    is_partial_llm_turn,
    normalize_tool_definition,
)
from sinapsis_chatbots_base.helpers.llm_keys import LLMChatKeys, MCPKeys
from sinapsis_chatbots_base.helpers.tags import Tags
from sinapsis_core.data_containers.data_packet import DataContainer, LLMConversationPacket
from sinapsis_core.data_containers.model_types import (
    LLMMessage,
    LLMToolCall,
    LLMToolResult,
)
from sinapsis_core.template_base.base_models import OutputTypes, TemplateAttributeType, UIPropertiesMetadata


class AnthropicWithMCP(AnthropicTextGeneration):
    """Template for Anthropic chat processing with MCP tool support.

    Usage example:

    agent:
      name: my_test_agent
    templates:
    - template_name: InputTemplate
      class_name: InputTemplate
      attributes: {}
    - template_name: AnthropicWithMCP
      class_name: AnthropicWithMCP
      template_input: LLMConversationInput
      attributes:
        init_args:
          llm_model_name: claude-sonnet-4-6
        completion_args:
          max_tokens: 1024
          temperature: 0.2
          service_tier: standard_only
        extended_thinking:
          type: disabled
        web_search: false

    An upstream template or orchestrator must populate `LLMConversationPacket.tools`
    before this template runs.
    """

    UIProperties = UIPropertiesMetadata(
        category="Anthropic",
        output_type=OutputTypes.TEXT,
        tags=[Tags.ANTHROPIC, Tags.CHATBOTS, Tags.CONTEXT, Tags.LLM, Tags.TEXT_COMPLETION, Tags.TEXT, Tags.MCP],
    )

    class AttributesBaseModel(AnthropicAttributes):
        """Attributes for Anthropic MCP templates.

        Attributes:
            init_args (LLMInitArgs): Anthropic model configuration.
            completion_args (AnthropicCompletionArgs): Anthropic generation arguments.
            extended_thinking (AnthropicThinkingArgs): Extended-thinking configuration.
            web_search (bool): Whether to enable Anthropic's web search tool.
        """

    def __init__(self, attributes: TemplateAttributeType) -> None:
        super().__init__(attributes)
        self.has_pending_tool_calls = False

    @staticmethod
    def build_available_tools(llm_conversation: LLMConversationPacket) -> list[dict[str, Any]]:
        """Build Anthropic-compatible tools directly from the packet.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet being processed.

        Returns:
            list[dict[str, Any]]: Anthropic-compatible tools for the current packet.
        """
        raw_tools = [normalize_tool_definition(tool) for tool in llm_conversation.tools]
        return make_tools_anthropic_compatible(raw_tools)

    @classmethod
    def build_packet_message_from_chat_message(cls, message: dict[str, Any] | None) -> LLMMessage | None:
        """Convert a normalized chat dictionary into packet history.

        Args:
            message (dict[str, Any] | None): Chat message dictionary to convert.

        Returns:
            LLMMessage | None: Packet history message when content is present.
        """
        if message is None:
            return None

        return build_packet_message(
            role=message[LLMChatKeys.role],
            content=message[LLMChatKeys.content],
            name=message.get(LLMChatKeys.name),
        )

    @staticmethod
    def _serialize_tool_result_content(content: Any) -> str | list[dict[str, Any]]:
        """Serialize tool result content for Anthropic tool-result blocks.

        Args:
            content (Any): Tool result payload.

        Returns:
            str | list[dict[str, Any]]: Anthropic-compatible tool result content.
        """
        if isinstance(content, str):
            return content

        if isinstance(content, list) and all(isinstance(item, dict) and AnthropicKeys.type in item for item in content):
            return content

        return json.dumps(content, ensure_ascii=True, default=str)

    @classmethod
    def build_tool_result_message(cls, tool_results: list[LLMToolResult]) -> dict[str, Any] | None:
        """Build the user message that returns tool outputs to Anthropic.

        Args:
            tool_results (list[LLMToolResult]): Tool execution results.

        Returns:
            dict[str, Any] | None: Anthropic-compatible tool result message.
        """
        if not tool_results:
            return None

        content: list[dict[str, Any]] = []
        for tool_result in tool_results:
            tool_use_id = tool_result.tool_call_id or tool_result.name
            tool_content: dict[str, Any] = {
                AnthropicKeys.type: AnthropicKeys.tool_result,
                AnthropicKeys.tool_use_id: tool_use_id,
                AnthropicKeys.content: cls._serialize_tool_result_content(tool_result.content),
            }
            if tool_result.is_error:
                tool_content[MCPKeys.is_error] = True
            content.append(tool_content)

        return build_chat_message(LLMChatKeys.user_value, content)

    @staticmethod
    def build_tool_messages(llm_conversation: LLMConversationPacket) -> list[dict[str, Any]]:
        """Suppress base tool-result messages for MCP turns.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet to read.

        Returns:
            list[dict[str, Any]]: Always empty for Anthropic MCP requests.
        """
        _ = llm_conversation
        return []

    def build_current_turn_messages(
        self,
        llm_conversation: LLMConversationPacket,
        container: DataContainer | None = None,
    ) -> list[dict[str, Any]]:
        """Build the current MCP turn for Anthropic.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet to convert.
            container (DataContainer | None): Optional container holding shared context.

        Returns:
            list[dict[str, Any]]: Anthropic-compatible current-turn messages.
        """
        _ = container
        if is_partial_llm_turn(llm_conversation):
            tool_result_message = self.build_tool_result_message(llm_conversation.tool_results)
            if tool_result_message is None:
                return []
            return [tool_result_message]

        return super().build_current_turn_messages(llm_conversation, container)

    def build_request_tools(self, llm_conversation: LLMConversationPacket) -> list[dict[str, Any]]:
        """Build Anthropic MCP tools for the current packet.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet being processed.

        Returns:
            list[dict[str, Any]]: Anthropic-compatible tools for the current packet.
        """
        tools = self.build_available_tools(llm_conversation)
        tools.extend(super().build_request_tools(llm_conversation))
        return tools

    @staticmethod
    def _build_assistant_text_block(text: str) -> dict[str, str]:
        """Build an Anthropic assistant text block.

        Args:
            text (str): Text block content.

        Returns:
            dict[str, str]: Anthropic-compatible text block.
        """
        return {AnthropicKeys.type: AnthropicKeys.text, AnthropicKeys.text: text}

    @staticmethod
    def _build_assistant_tool_use_block(content_block: ToolUseBlock) -> dict[str, Any]:
        """Build an Anthropic assistant tool-use block.

        Args:
            content_block (ToolUseBlock): Anthropic tool-use block.

        Returns:
            dict[str, Any]: Anthropic-compatible tool-use block.
        """
        return {
            AnthropicKeys.type: MCPKeys.tool_use,
            AnthropicKeys.id: content_block.id,
            AnthropicKeys.name: content_block.name,
            AnthropicKeys.input: content_block.input,
        }

    @staticmethod
    def _handle_thinking_block(content_block: ThinkingBlock, reasoning_parts: list[str]) -> None:
        """Collect reasoning text from a thinking block.

        Args:
            content_block (ThinkingBlock): Anthropic thinking block.
            reasoning_parts (list[str]): Accumulated reasoning text.
        """
        reasoning_parts.append(content_block.thinking)

    @staticmethod
    def _build_assistant_message_content(assistant_content: list[dict[str, Any]]) -> str | list[dict[str, Any]] | None:
        """Normalize assistant content blocks for packet history.

        Args:
            assistant_content (list[dict[str, Any]]): Anthropic assistant content blocks.

        Returns:
            str | list[dict[str, Any]] | None: Persistable packet message content.
        """
        if not assistant_content:
            return None

        if any(content_block.get(AnthropicKeys.type) != AnthropicKeys.text for content_block in assistant_content):
            return assistant_content

        response_text = "".join(content_block[AnthropicKeys.text] for content_block in assistant_content).strip()
        return response_text or None

    def append_packet_history(
        self,
        llm_conversation: LLMConversationPacket,
        assistant_content: list[dict[str, Any]],
        was_partial: bool,
    ) -> None:
        """Persist the current MCP step directly on packet history.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet to update.
            assistant_content (list[dict[str, Any]]): Current assistant content blocks.
            was_partial (bool): Whether this response continues an in-progress tool turn.
        """
        if was_partial:
            tool_result_message = self.build_tool_result_message(llm_conversation.tool_results)
            packet_message = self.build_packet_message_from_chat_message(tool_result_message)
            if packet_message is not None:
                llm_conversation.messages.append(packet_message)
        else:
            user_message = self.build_user_message(llm_conversation)
            packet_message = self.build_packet_message_from_chat_message(user_message)
            if packet_message is not None:
                llm_conversation.messages.append(packet_message)

        assistant_message = build_packet_message(
            role=LLMChatKeys.assistant_value,
            content=self._build_assistant_message_content(assistant_content),
        )
        if assistant_message is not None:
            llm_conversation.messages.append(assistant_message)

    def update_packet_from_response(self, llm_conversation: LLMConversationPacket, response: Any) -> None:
        """Update a packet from an Anthropic MCP response.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet to update.
            response (Any): Raw Anthropic response object.
        """
        if not isinstance(response, Message):
            return

        was_partial = is_partial_llm_turn(llm_conversation)
        previous_response = llm_conversation.response
        reasoning_parts: list[str] = []
        response_parts: list[str] = []
        assistant_content: list[dict[str, Any]] = []
        tool_calls: list[LLMToolCall] = []

        for content_block in response.content:
            if content_block.type == AnthropicKeys.thinking:
                self._handle_thinking_block(content_block, reasoning_parts)
            elif content_block.type == AnthropicKeys.text:
                response_parts.append(content_block.text)
                assistant_content.append(self._build_assistant_text_block(content_block.text))
            elif content_block.type == MCPKeys.tool_use:
                assistant_content.append(self._build_assistant_tool_use_block(content_block))
                tool_calls.append(
                    LLMToolCall(name=content_block.name, arguments=content_block.input, id=content_block.id)
                )

        current_response = "".join(response_parts).strip() or None
        combined_response = combine_assistant_responses(previous_response, current_response)
        reasoning_text = "\n\n".join(reasoning_parts).strip() or None

        self.append_packet_history(llm_conversation, assistant_content, was_partial)

        llm_conversation.response = combined_response
        llm_conversation.reasoning = reasoning_text
        self.update_packet_metadata(llm_conversation, response, is_partial=bool(tool_calls))

        if tool_calls:
            llm_conversation.tool_calls = tool_calls
            llm_conversation.tool_results = []
            return

        self._cleanup_conversation_state(llm_conversation)

    def execute(self, container: DataContainer) -> DataContainer:
        """Process all conversation packets with MCP support.

        Args:
            container (DataContainer): Input container.

        Returns:
            DataContainer: Updated container with Anthropic MCP state.
        """
        self.has_pending_tool_calls = False
        if not container.llm_conversations:
            self.logger.debug("Container has no llm conversations to process. Returning.")
            return container

        for llm_conversation in container.llm_conversations:
            self.process_llm_packet(llm_conversation=llm_conversation, container=container)
            self.has_pending_tool_calls = self.has_pending_tool_calls or bool(llm_conversation.tool_calls)

        return container

    def num_elements(self) -> int:
        """Determine whether the MCP loop should continue.

        Returns:
            int: `1` when pending tool calls exist, `-1` otherwise.
        """
        return 1 if self.has_pending_tool_calls else -1

    @staticmethod
    def _cleanup_conversation_state(llm_conversation: LLMConversationPacket) -> None:
        """Clear MCP state from a completed conversation packet.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet to clean up.
        """
        llm_conversation.tool_calls = []
        llm_conversation.tool_results = []
