from dataclasses import asdict, is_dataclass
from typing import Any

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

from sinapsis_chat_history.templates.chat_history_base import ChatHistoryBase, ChatHistoryBaseAttributes


class ChatHistorySaveArgs(BaseModel):
    """Arguments for saving chat history messages.

    Attributes:
        include_packet_metadata (bool): Whether packet-level metadata such as
            provider info, usage, performance, and termination should be merged
            into the stored message metadata.
    """

    include_packet_metadata: bool = False


class ChatHistorySave(ChatHistoryBase):
    """Template for saving chat messages into a storage backend.

    Usage example:
    agent:
      name: my_test_agent
    templates:
    - template_name: InputTemplate
      class_name: InputTemplate
      attributes: {}
    - template_name: ChatHistorySave
      class_name: ChatHistorySave
      template_input: InputTemplate
      attributes:
        provider: sqlite
        db_config:
          db_path: "chat_history.sqlite3"
          table: "chat"
    """

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

        Attributes:
            save_args (ChatHistorySaveArgs): Save-time options for message
                persistence.
        """

        save_args: ChatHistorySaveArgs = Field(default_factory=ChatHistorySaveArgs)

    def extract_packet_metadata(self, llm_conversation: LLMConversationPacket) -> dict[str, Any]:
        """Extract packet metadata to be stored with chat messages.

        Args:
            llm_conversation (LLMConversationPacket): Packet from which to
                extract metadata.

        Returns:
            dict[str, Any]: Extracted packet metadata or an empty dictionary
                when metadata persistence is disabled.
        """
        if not self.attributes.save_args.include_packet_metadata:
            return {}

        metadata: dict[str, Any] = {}
        for field_name in ("provider_info", "usage", "performance", "termination"):
            value = getattr(llm_conversation, field_name, None)
            if value is None:
                continue
            if is_dataclass(value) and not isinstance(value, type):
                metadata[field_name] = asdict(value)
            else:
                metadata[field_name] = value
        return metadata

    def save_messages(self, llm_conversation: LLMConversationPacket) -> None:
        """Save the current turn from a conversation packet.

        Args:
            llm_conversation (LLMConversationPacket): Conversation packet to
                persist.
        """
        identity = extract_packet_identity(llm_conversation, exclude={"run_id", "agent_id"})
        if not identity:
            self.logger.warning("Skipping chat history save because at least user_id needs to be set.")
            return

        messages = build_current_turn_messages(llm_conversation)
        if not messages:
            self.logger.debug("No chat history messages found in the packet to persist.")
            return
        metadata = self.extract_packet_metadata(llm_conversation)
        self.db.save_messages(messages=messages, metadata=metadata, **identity)

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

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

        Returns:
            DataContainer: Unmodified input container.
        """
        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.save_messages(llm_conversation)
        return container
