from collections.abc import Sequence
from typing import Any

from sinapsis_chatbots_base.helpers.llm_conversation import build_chat_message
from sinapsis_chatbots_base.helpers.llm_keys import LLMChatKeys
from sinapsis_core.data_containers.data_packet import DataContainer, LLMConversationPacket
from sinapsis_core.data_containers.model_types import LLMContextItem, LLMMessage
from sinapsis_core.template_base import Template
from sinapsis_core.template_base.base_models import (
    OutputTypes,
    TemplateAttributes,
    TemplateAttributeType,
    UIPropertiesMetadata,
)
from transformers import AutoTokenizer

from sinapsis_chat_history.helpers.tags import Tags


class ChatHistoryTrimAttributes(TemplateAttributes):
    """Attributes that control token-aware history trimming.

    Attributes:
        tokenizer_model_name: Hugging Face tokenizer name or path used for prompt-size estimation.
        model_max_context: Maximum input-plus-output context length supported by the target model.
        reserved_output_tokens: Tokens reserved for model generation.
        safety_margin_tokens: Extra token buffer reserved to absorb formatting and tokenizer-estimation mismatches.
        preserve_system_messages: Whether history messages with `system` or `developer` roles should be retained
            whenever possible.

    """

    tokenizer_model_name: str
    model_max_context: int
    reserved_output_tokens: int = 1024
    safety_margin_tokens: int = 256
    preserve_system_messages: bool = True


class ChatHistoryTrim(Template):
    """Template to trim fetched chat history to fit a target prompt budget.

    Usage example:
    agent:
      name: my_test_agent
    templates:
    - template_name: InputTemplate
      class_name: InputTemplate
      attributes: {}
    - template_name: ChatHistoryTrim
      class_name: ChatHistoryTrim
      template_input: InputTemplate
      attributes:
        tokenizer_model_name: 'Qwen/Qwen3.5-9B'
        model_max_context: 8192
        reserved_output_tokens: 2048

    """

    AttributesBaseModel = ChatHistoryTrimAttributes
    UIProperties = UIPropertiesMetadata(
        category="Chat History",
        output_type=OutputTypes.TEXT,
        tags=[Tags.CHATBOTS, Tags.CHAT_HISTORY, Tags.CONTEXT],
    )

    def __init__(self, attributes: TemplateAttributeType) -> None:
        super().__init__(attributes)
        self.tokenizer = AutoTokenizer.from_pretrained(self.attributes.tokenizer_model_name)

    @staticmethod
    def _format_context_item(context: LLMContextItem) -> str:
        """Format one packet context item like the downstream LLM templates do.

        Args:
            context (LLMContextItem): Context item to serialize.

        Returns:
            str: Formatted context string.
        """
        header = f"[{context.category}]"
        if context.source:
            header = f"{header} {context.source}"
        return f"{header}\n{context.content}"

    def _build_system_message(self, llm_conversation: LLMConversationPacket) -> dict[str, Any] | None:
        """Build the synthetic system message used for token estimation.

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

        Returns:
            dict[str, Any] | None: System message payload when available.
        """
        system_parts: list[str] = []

        if llm_conversation.system_prompt:
            system_parts.append(llm_conversation.system_prompt)

        if llm_conversation.contexts:
            context_lines = [self._format_context_item(context) for context in llm_conversation.contexts]
            context_text = "\n\n".join(context_lines)
            system_parts.append(f"Use the following context:\n{context_text}")

        if not system_parts:
            return None

        return build_chat_message(LLMChatKeys.system_value, "\n\n".join(system_parts))

    @staticmethod
    def _build_history_message(message: LLMMessage) -> dict[str, Any]:
        """Convert one stored message into the downstream chat payload format.

        Args:
            message (LLMMessage): History message to convert.

        Returns:
            dict[str, Any]: Provider-neutral chat message payload.
        """
        payload = build_chat_message(message.role, message.content)
        if message.name:
            payload[LLMChatKeys.name] = message.name
        return payload

    @staticmethod
    def _build_tool_messages(llm_conversation: LLMConversationPacket) -> list[dict[str, Any]]:
        """Build tool-result messages for token estimation.

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

        Returns:
            list[dict[str, Any]]: Tool result payloads.
        """
        tool_messages: list[dict[str, Any]] = []
        for tool_result in llm_conversation.tool_results:
            message = build_chat_message(LLMChatKeys.tool_value, tool_result.content)
            if tool_result.tool_call_id:
                message[LLMChatKeys.tool_call_id] = tool_result.tool_call_id
            elif tool_result.name:
                message[LLMChatKeys.name] = tool_result.name
            tool_messages.append(message)
        return tool_messages

    @staticmethod
    def _should_append_current_prompt(llm_conversation: LLMConversationPacket) -> bool:
        """Return whether the current packet content will be appended as a user turn.

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

        Returns:
            bool: `True` when the current prompt will be appended.
        """
        if llm_conversation.content is None:
            return False

        if not llm_conversation.messages:
            return True

        last_message = llm_conversation.messages[-1]
        return not (last_message.role == LLMChatKeys.user_value and last_message.content == llm_conversation.content)

    def _build_effective_messages(
        self,
        llm_conversation: LLMConversationPacket,
        history_messages: Sequence[LLMMessage],
    ) -> list[dict[str, Any]]:
        """Build the full message list that the downstream LLM call will see.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet to inspect.
            history_messages (Sequence[LLMMessage]): History messages to include after trimming.

        Returns:
            list[dict[str, Any]]: Provider-neutral backend messages.
        """
        messages: list[dict[str, Any]] = []

        system_message = self._build_system_message(llm_conversation)
        if system_message:
            messages.append(system_message)

        messages.extend(self._build_history_message(message) for message in history_messages)
        messages.extend(self._build_tool_messages(llm_conversation))

        if self._should_append_current_prompt(llm_conversation):
            messages.append(build_chat_message(LLMChatKeys.user_value, llm_conversation.content))

        return messages

    def estimate_prompt_tokens(
        self,
        llm_conversation: LLMConversationPacket,
        history_messages: Sequence[LLMMessage],
    ) -> int:
        """Estimate prompt tokens for one packet and one history window.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet to inspect.
            history_messages (Sequence[LLMMessage]): History messages to include in the estimate.

        Returns:
            int: Estimated prompt token count.
        """
        messages = self._build_effective_messages(llm_conversation, history_messages)
        apply_chat_template = getattr(self.tokenizer, "apply_chat_template", None)
        if callable(apply_chat_template):
            token_ids = apply_chat_template(
                list(messages),
                tokenize=True,
                add_generation_prompt=False,
            )
            return len(token_ids)

        raise RuntimeError(
            f"Tokenizer `{self.attributes.tokenizer_model_name}` does not implement `apply_chat_template`."
        )

    def _message_is_pinned(self, message: LLMMessage) -> bool:
        """Return whether one history message should be treated as pinned.

        Args:
            message (LLMMessage): History message to inspect.

        Returns:
            bool: `True` when the message should be retained whenever possible.
        """
        return self.attributes.preserve_system_messages and message.role in {
            LLMChatKeys.system_value,
            "developer",
        }

    @staticmethod
    def _starts_new_turn(message: LLMMessage) -> bool:
        """Return whether one message starts a new conversational turn.

        Args:
            message (LLMMessage): History message to inspect.

        Returns:
            bool: `True` when a new turn should start at this message.
        """
        return message.role in {LLMChatKeys.user_value, LLMChatKeys.system_value, "developer"}

    def _partition_history(self, history_messages: Sequence[LLMMessage]) -> tuple[set[int], list[tuple[int, ...]]]:
        """Partition history into pinned messages and flexible turn groups.

        Args:
            history_messages (Sequence[LLMMessage]): Stored conversation history.

        Returns:
            tuple[set[int], list[tuple[int, ...]]]: Pinned message indices and
                flexible turn groups expressed as original message indices.
        """
        pinned_indices: set[int] = set()
        flexible_turns: list[tuple[int, ...]] = []
        current_turn: list[int] = []

        for index, message in enumerate(history_messages):
            if self._message_is_pinned(message):
                if current_turn:
                    flexible_turns.append(tuple(current_turn))
                    current_turn = []
                pinned_indices.add(index)
                continue

            if current_turn and self._starts_new_turn(message):
                flexible_turns.append(tuple(current_turn))
                current_turn = [index]
                continue

            current_turn.append(index)

        if current_turn:
            flexible_turns.append(tuple(current_turn))

        return pinned_indices, flexible_turns

    def _select_history_indices(
        self,
        llm_conversation: LLMConversationPacket,
        history_messages: Sequence[LLMMessage],
        input_budget: int,
    ) -> set[int]:
        """Select the history message indices that best fit the prompt budget.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet whose history is being trimmed.
            history_messages (Sequence[LLMMessage]): Full current history for the packet.
            input_budget: Maximum prompt-token budget available for input.

        Returns:
            set[int]: Indices of the history messages that should be kept.
        """
        pinned_indices, flexible_turns = self._partition_history(history_messages)
        kept_indices = set(pinned_indices)

        for turn in reversed(flexible_turns):
            trial_indices = kept_indices.union(turn)
            trial_history = [message for index, message in enumerate(history_messages) if index in trial_indices]
            if self.estimate_prompt_tokens(llm_conversation, trial_history) <= input_budget:
                kept_indices = trial_indices

        return kept_indices

    def trim_packet_history(self, llm_conversation: LLMConversationPacket) -> None:
        """Trim one conversation packet in place.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet whose history should be
                trimmed.
        """
        history_messages = list(llm_conversation.messages)
        if not history_messages:
            return

        input_budget = max(
            self.attributes.model_max_context
            - self.attributes.reserved_output_tokens
            - self.attributes.safety_margin_tokens,
            0,
        )
        if input_budget <= 0:
            self.logger.warning("ChatHistoryTrim has no available input budget after reservations.")
            return

        original_token_estimate = self.estimate_prompt_tokens(llm_conversation, history_messages)
        if original_token_estimate <= input_budget:
            return

        kept_indices = self._select_history_indices(llm_conversation, history_messages, input_budget)
        llm_conversation.messages = [message for index, message in enumerate(history_messages) if index in kept_indices]

        final_token_estimate = self.estimate_prompt_tokens(llm_conversation, llm_conversation.messages)
        if final_token_estimate > input_budget:
            self.logger.warning(
                "ChatHistoryTrim could not fully fit the prompt budget. "
                f"Estimated prompt tokens={final_token_estimate}, input budget={input_budget}."
            )

    def execute(self, container: DataContainer) -> DataContainer:
        """Trim chat history for every conversation packet in the container.

        Args:
            container (DataContainer): Container holding conversation packets.

        Returns:
            DataContainer: Input container with trimmed packet histories.
        """
        if not container.llm_conversations:
            self.logger.debug("No llm conversation packet in the container.")
            return container

        for llm_conversation in container.llm_conversations:
            self.trim_packet_history(llm_conversation)

        return container
