from pydantic import BaseModel, Field
from sinapsis_chatbots_base.helpers.llm_conversation import extract_packet_identity
from sinapsis_core.data_containers.data_packet import DataContainer, LLMConversationPacket

from sinapsis_chat_history.helpers.schemas import ChatHistoryFilters, ChatHistoryOrder, ChatHistoryOrderBy
from sinapsis_chat_history.templates.chat_history_base import (
    ChatHistoryBase,
    ChatHistoryBaseAttributes,
)


class ChatHistoryFetchArgs(BaseModel):
    """Arguments for retrieving chat history messages.

    Attributes:
        limit (int): Maximum number of messages to fetch. Use `-1` to fetch
            all matching messages.
        offset (int): Number of matching messages to skip before returning
            results.
        order (ChatHistoryOrder): Sort direction for the returned messages.
        order_by (ChatHistoryOrderBy): Field used to order the returned
            messages.
        filters (ChatHistoryFilters): Optional message-level filters applied in
            addition to packet identity scope.
    """

    limit: int = 100
    offset: int = 0
    order: ChatHistoryOrder = "asc"
    order_by: ChatHistoryOrderBy = "created_at"
    filters: ChatHistoryFilters = Field(default_factory=ChatHistoryFilters)


class ChatHistoryFetch(ChatHistoryBase):
    """Template for loading chat history messages from a storage backend.

    Usage example:
    agent:
      name: my_test_agent
    templates:
    - template_name: InputTemplate
      class_name: InputTemplate
      attributes: {}
    - template_name: ChatHistoryFetch
      class_name: ChatHistoryFetch
      template_input: InputTemplate
      attributes:
        provider: sqlite
        db_config:
          db_path: "chat_history.sqlite3"
          table: "chat"
        fetch_args:
          limit: 10
          offset: 0
          order: asc
          order_by: created_at
    """

    class AttributesBaseModel(ChatHistoryBaseAttributes):
        """Attributes for configuring how chat history is fetched.

        Attributes:
            fetch_args (ChatHistoryFetchArgs): Query options used when loading
                chat history messages.
        """

        fetch_args: ChatHistoryFetchArgs = Field(default_factory=ChatHistoryFetchArgs)

    def fetch_histories(self, llm_conversation: LLMConversationPacket) -> None:
        """Fetch chat messages for a single conversation packet.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet whose
                scoped history should be loaded.
        """
        identity = extract_packet_identity(llm_conversation, exclude={"run_id", "agent_id"})
        if not identity:
            self.logger.warning("Skipping chat history fetch because packet identity has no user_id or session_id.")
            return

        llm_conversation.messages = self.db.fetch_messages(
            **identity, **self.attributes.fetch_args.model_dump(exclude_none=True)
        )

    def execute(self, container: DataContainer) -> DataContainer:
        """Fetch chat history for all conversation packets in the container.

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

        Returns:
            DataContainer: Input container with fetched histories attached to
                each packet's `messages`.
        """
        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.fetch_histories(llm_conversation)
        return container
