from typing import Any

from anthropic import Anthropic
from anthropic.types import Message
from pydantic import Field, model_validator
from sinapsis_anthropic.helpers.anthropic_keys import AnthropicKeys
from sinapsis_anthropic.helpers.env_var_keys import ANTHROPIC_API_KEY
from sinapsis_anthropic.helpers.schemas import AnthropicCompletionArgs, AnthropicThinkingArgs
from sinapsis_chatbots_base.helpers.llm_keys import LLMChatKeys
from sinapsis_chatbots_base.helpers.tags import Tags
from sinapsis_chatbots_base.templates.llm_template_base import LLMTemplateBase, LLMTemplateBaseAttributes
from sinapsis_core.data_containers.data_packet import DataContainer, LLMConversationPacket
from sinapsis_core.data_containers.model_types import LLMProviderInfo, LLMTermination, LLMUsage
from sinapsis_core.template_base.base_models import OutputTypes, UIPropertiesMetadata


class AnthropicAttributes(LLMTemplateBaseAttributes):
    """Base attributes for Anthropic 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.
    """

    completion_args: AnthropicCompletionArgs
    extended_thinking: AnthropicThinkingArgs = Field(default_factory=AnthropicThinkingArgs)
    web_search: bool = False

    @model_validator(mode="after")
    def validate_extended_thinking(self) -> "AnthropicAttributes":
        """Validate Anthropic extended-thinking constraints.

        Returns:
            AnthropicAttributes: Validated attributes.

        Raises:
            ValueError: If the extended-thinking configuration is incompatible
                with the selected completion arguments.
        """
        if self.extended_thinking.type != AnthropicKeys.enabled:
            return self

        completion_args = self.completion_args

        if completion_args.temperature is not None and completion_args.temperature != 1:
            raise ValueError("Anthropic extended thinking requires `temperature` to be omitted or set to `1`.")

        if completion_args.top_k is not None:
            raise ValueError("Anthropic extended thinking is not compatible with `top_k`.")

        if completion_args.top_p is not None and not (0.95 <= completion_args.top_p <= 1):
            raise ValueError("Anthropic extended thinking only supports `top_p` values between `0.95` and `1`.")

        if completion_args.max_tokens <= self.extended_thinking.budget_tokens:
            raise ValueError("Anthropic extended thinking requires `max_tokens` to be greater than `budget_tokens`.")

        return self


class AnthropicTextGeneration(LLMTemplateBase):
    """Template for Anthropic text generation using `LLMConversationPacket`.

    Usage example:

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

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

    @staticmethod
    def init_llm_model() -> Anthropic:
        """Initialize the Anthropic client.

        Returns:
            Anthropic: Configured Anthropic client.

        Raises:
            ValueError: If `ANTHROPIC_API_KEY` is not configured.
        """
        if not ANTHROPIC_API_KEY:
            raise ValueError("ANTHROPIC_API_KEY is not configured.")
        return Anthropic(api_key=ANTHROPIC_API_KEY)

    def build_system_prompt(self, llm_conversation: LLMConversationPacket) -> str | None:
        """Build the Anthropic system prompt from packet system data.

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

        Returns:
            str | None: System prompt text, if any.
        """
        system_message = self.build_system_message(llm_conversation)
        if system_message is None:
            return None

        return system_message[LLMChatKeys.content]

    @staticmethod
    def convert_message_to_anthropic(message: dict[str, Any]) -> dict[str, Any]:
        """Convert a normalized chat message into Anthropic's message format.

        Args:
            message (dict[str, Any]): Normalized chat message.

        Returns:
            dict[str, Any]: Anthropic-compatible message payload.
        """
        role = message[LLMChatKeys.role]
        content = message[LLMChatKeys.content]

        if role == LLMChatKeys.tool_value:
            tool_call_id = message.get(LLMChatKeys.tool_call_id)
            if tool_call_id:
                return {
                    LLMChatKeys.role: LLMChatKeys.user_value,
                    LLMChatKeys.content: [
                        {
                            AnthropicKeys.type: "tool_result",
                            "tool_use_id": tool_call_id,
                            AnthropicKeys.content: content,
                        }
                    ],
                }
            return {
                LLMChatKeys.role: LLMChatKeys.user_value,
                LLMChatKeys.content: content,
            }

        return {
            LLMChatKeys.role: role,
            LLMChatKeys.content: content,
        }

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

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

        Returns:
            list[dict[str, Any]]: Anthropic tool definitions.
        """
        _ = llm_conversation
        if not self.attributes.web_search:
            return []

        return [
            {
                AnthropicKeys.name: AnthropicKeys.web_search,
                AnthropicKeys.type: AnthropicKeys.web_search_20250305,
            }
        ]

    def build_output_config(self) -> dict[str, Any] | None:
        """Build Anthropic structured-output configuration when requested.

        Returns:
            dict[str, Any] | None: Anthropic `output_config` payload, if structured output is enabled.
        """
        response_format = self.attributes.completion_args.response_format
        if not response_format.is_structured():
            return None

        schema = response_format.to_json_schema()
        if schema is None:
            return None

        schema = self._close_object_schema(schema)

        return {
            AnthropicKeys.format: {
                AnthropicKeys.type: AnthropicKeys.json_schema,
                "schema": schema,
            }
        }

    @classmethod
    def _close_object_schema(cls, schema: Any) -> Any:
        """Recursively set `additionalProperties: false` on object schemas.

        Anthropic requires every object node in a structured-output schema to
        declare `additionalProperties: false`.

        Args:
            schema (Any): JSON schema fragment to normalize.

        Returns:
            Any: Normalized schema fragment.
        """
        if isinstance(schema, list):
            return [cls._close_object_schema(item) for item in schema]

        if not isinstance(schema, dict):
            return schema

        normalized = dict(schema)
        if normalized.get("type") == "object":
            normalized.setdefault("additionalProperties", False)

        properties = normalized.get("properties")
        if isinstance(properties, dict):
            normalized["properties"] = {key: cls._close_object_schema(value) for key, value in properties.items()}

        if "items" in normalized:
            normalized["items"] = cls._close_object_schema(normalized["items"])

        for key in ("allOf", "anyOf", "oneOf", "prefixItems"):
            if key in normalized and isinstance(normalized[key], list):
                normalized[key] = [cls._close_object_schema(item) for item in normalized[key]]

        if "$defs" in normalized and isinstance(normalized["$defs"], dict):
            normalized["$defs"] = {key: cls._close_object_schema(value) for key, value in normalized["$defs"].items()}

        return normalized

    def build_backend_payload(
        self,
        llm_conversation: LLMConversationPacket,
        container: DataContainer | None = None,
    ) -> dict[str, Any]:
        """Build the Anthropic request payload.

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

        Returns:
            dict[str, Any]: Anthropic `messages.create()` payload.
        """
        messages = self.build_non_system_messages(llm_conversation, container)
        system_prompt = self.build_system_prompt(llm_conversation)

        payload: dict[str, Any] = {
            AnthropicKeys.model: self.attributes.init_args.llm_model_name,
            AnthropicKeys.messages: [self.convert_message_to_anthropic(message) for message in messages],
        }

        if system_prompt:
            payload[AnthropicKeys.system] = system_prompt

        if self.attributes.extended_thinking.type == AnthropicKeys.enabled:
            payload[AnthropicKeys.thinking] = self.attributes.extended_thinking.model_dump()

        tools = self.build_request_tools(llm_conversation)
        if tools:
            payload[AnthropicKeys.tools] = tools

        output_config = self.build_output_config()
        if output_config is not None:
            payload[AnthropicKeys.output_config] = output_config

        payload.update(self.attributes.completion_args.model_dump(exclude_none=True, exclude={"response_format"}))
        return payload

    def invoke_backend(self, payload: Any, llm_conversation: LLMConversationPacket) -> Message:
        """Invoke the Anthropic API.

        Args:
            payload (Any): Anthropic request payload.
            llm_conversation (LLMConversationPacket): Conversation packet being processed.

        Returns:
            Message: Raw Anthropic message response.
        """
        _ = llm_conversation
        self.logger.debug(f"Anthropic request payload: {payload}")
        return self.llm.messages.create(**payload)

    @staticmethod
    def extract_response_text(message_response: Message) -> str | None:
        """Extract the final answer text from an Anthropic response.

        Args:
            message_response (Message): Anthropic response object.

        Returns:
            str | None: Final response text, if any.
        """
        response_parts = [
            content_block.text for content_block in message_response.content if content_block.type == AnthropicKeys.text
        ]

        response_text = "".join(response_parts).strip()
        return response_text or None

    @staticmethod
    def extract_reasoning_text(message_response: Message) -> str | None:
        """Extract reasoning text from Anthropic thinking blocks.

        Args:
            message_response (Message): Anthropic response object.

        Returns:
            str | None: Extracted reasoning text, if any.
        """
        reasoning_parts = [
            content_block.thinking
            for content_block in message_response.content
            if content_block.type == AnthropicKeys.thinking
        ]

        reasoning_text = "\n\n".join(reasoning_parts).strip()
        return reasoning_text or None

    @staticmethod
    def update_packet_metadata(
        llm_conversation: LLMConversationPacket,
        response: Message,
        *,
        is_partial: bool,
    ) -> None:
        """Update shared Anthropic packet metadata.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet to update.
            response (Message): Raw Anthropic response object.
            is_partial (bool): Whether the response represents a partial turn.
        """
        llm_conversation.provider_info = LLMProviderInfo(
            provider="anthropic",
            model_name=str(response.model),
        )
        llm_conversation.termination = LLMTermination(
            finish_reason=response.stop_reason,
            stop_reason=response.stop_sequence,
            is_partial=is_partial,
        )

        if response.usage is not None:
            input_tokens = response.usage.input_tokens
            output_tokens = response.usage.output_tokens
            llm_conversation.usage = LLMUsage(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                reasoning_tokens=None,
                total_tokens=input_tokens + output_tokens,
            )

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

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

        llm_conversation.response = self.extract_response_text(response)
        llm_conversation.reasoning = self.extract_reasoning_text(response)
        self.update_packet_metadata(
            llm_conversation,
            response,
            is_partial=response.stop_reason in {"max_tokens", "pause_turn"},
        )
