
# Copyright © 2023-2026 Cognizant Technology Solutions Corp, www.cognizant.com.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# END COPYRIGHT

from typing import Any
from typing import Dict
from typing import List

from langchain_core.messages.ai import AIMessage
from langchain_core.messages.base import BaseMessage
from langchain_core.messages.human import HumanMessage
from langchain_core.messages.system import SystemMessage
from langchain_core.messages.tool import ToolMessage

from leaf_common.serialization.interface.dictionary_converter import DictionaryConverter

from neuro_san.internals.messages.agent_message import AgentMessage
from neuro_san.internals.messages.agent_framework_message import AgentFrameworkMessage
from neuro_san.internals.messages.agent_progress_message import AgentProgressMessage
from neuro_san.internals.messages.agent_tool_result_message import AgentToolResultMessage
from neuro_san.internals.messages.chat_message_type import ChatMessageType


class BaseMessageDictionaryConverter(DictionaryConverter):
    """
    A DictionaryConverter implementation which can convert langchain BaseMessages
    back and forth to our own ChatMessage dictionary format (as defined in chat.proto).
    """

    def __init__(self, origin: List[Dict[str, Any]] = None,
                 langchain_only: bool = True):
        """
        Constructor

        :param origin: A List of origin dictionaries indicating the origin of the run.
            The origin can be considered a path to the original call to the front-man.
            Origin dictionaries themselves each have the following keys:
                "tool"                  The string name of the tool in the spec
                "instantiation_index"   An integer indicating which incarnation
                                        of the tool is being dealt with.
            This is only used for to_dict()
        :param langchain_only: When True (default) only convert langchain messages,
                        and not our own special variants.
            This is only used for from_dict()
        """
        self.origin: List[Dict[str, Any]] = origin
        self.langchain_only: bool = langchain_only

    def to_dict(self, obj: BaseMessage) -> Dict[str, Any]:
        """
        Convert the BaseMessage to a chat.ChatMessage dictionary

        :param obj: The BaseMessage to convert
        :return: The ChatMessage in dictionary form
        """

        message: BaseMessage = obj
        message_type: ChatMessageType = ChatMessageType.from_message(message)
        chat_message: Dict[str, Any] = {
            "type": message_type,
            # No mime_data for now
        }

        # Handle the origin information if we have it
        if self.origin is not None:
            chat_message["origin"] = self.origin

        # Dictionary of BaseMessage field sources to ChatMessage destinations
        # Anything in this dictionary is considered optional and we only populate
        # the field on ChatMessage if it has a value.
        optionals: Dict[str, str] = {
            "content": "text",
            "chat_context": "chat_context",
            "tool_result_origin": "tool_result_origin",
            "structure": "structure",
            "sly_data": "sly_data",
        }
        for src, dest in optionals.items():
            value: Any = None
            try:
                value = getattr(message, src)
            except AttributeError:
                # Not all BaseMessage subclasses have every field we are looking
                # for, and that is ok.
                value = None

            # For OpenAI and Ollama, content of AI message is a string but content from
            # Anthropic AI message can either be a single string or a list of content blocks.
            # If it is a list, "text" is a key of a dictionary which is the first element of
            # the list. For more details: https://python.langchain.com/docs/integrations/chat/anthropic/#content-blocks
            if src == "content" and isinstance(value, list):
                if len(value) > 0:
                    value = value[0].get("text", "")
                else:
                    value = None

            if value is not None:
                chat_message[dest] = value

        return chat_message

    def from_dict(self, obj_dict: Dict[str, Any]) -> BaseMessage:
        """
        The intended use for this method is for converting a neuro-san ChatMessage
        dictionary into a langchain BaseMessage such that it can be understood by
        a langchain agent within the context of a langchain-facing chat history.
        This means that some special messages we might have inserted into a chat
        history *might* need to be converted into something langchain agents will
        understand.

        :param obj_dict: A ChatMessage dictionary to convert into BaseMessage
        :return: A BaseMessage that was converted from the input.
                Can return None if conversion could not take place
        """
        chat_message: Dict[str, Any] = obj_dict
        base_message: BaseMessage = None
        if chat_message is None:
            return base_message

        content: str = chat_message.get("text")
        chat_message_type: ChatMessageType = ChatMessageType.from_response_type(chat_message.get("type"))

        if chat_message_type == ChatMessageType.SYSTEM:
            base_message = SystemMessage(content=content)
        elif chat_message_type == ChatMessageType.HUMAN:
            base_message = HumanMessage(content=content)
        elif chat_message_type == ChatMessageType.AI:
            base_message = AIMessage(content=content)
        elif chat_message_type == ChatMessageType.AGENT_TOOL_RESULT:
            base_message = AgentToolResultMessage(content=content,
                                                  tool_result_origin=chat_message.get("tool_result_origin"))

        if self.langchain_only:
            # Go no further
            return base_message

        if base_message is None:
            # convert_to_base_message() is primarily for populating langchain chat history.
            # so handle these other types differently.
            if chat_message_type == ChatMessageType.AGENT:
                base_message = AgentMessage(content=chat_message.get("text", ""),
                                            structure=chat_message.get("structure"))
            elif chat_message_type == ChatMessageType.AGENT_FRAMEWORK:
                # Don't bother passing on chat_context
                base_message = AgentFrameworkMessage(content=chat_message.get("text", ""))
            elif chat_message_type == ChatMessageType.AGENT_PROGRESS:
                base_message = AgentProgressMessage(content=chat_message.get("text", ""),
                                                    structure=chat_message.get("structure"))

        # Any other message type we do not want to send to any agent as chat history.

        return base_message

    @staticmethod
    def is_relevant_to_chat_history(base_message: BaseMessage) -> bool:
        """
        :param base_message: The BaseMessage to consider
        :return: True if the BaseMessage type is relevant to chat history (include).
                 False otherwise (do not include).
        """
        if isinstance(base_message, (AgentMessage, AgentFrameworkMessage, ToolMessage, AgentProgressMessage)):
            # These guys cannot be in chat history as langchain will not recognize them.
            return False
        return True
