from langchain.schema import (
    AgentAction,
    AIMessage,
    BaseMessage,
    HumanMessage,
    SystemMessage,
)

from .colorful import *
from .my_typing import *


def convert_to_message(content: str, type: str) -> BaseMessage:
    if type == "human":
        return HumanMessage(content=content)
    elif type == "ai":
        return AIMessage(content=content)
    elif type == "system":
        return SystemMessage(content=content)
    else:
        raise ValueError(f"Unknown message type: {type}")


def print_prompt(prompt: Union[List[BaseMessage], str], raw_text: bool = False):
    if isinstance(prompt, str):
        print(prompt)
    else:
        for msg in prompt:
            if raw_text:
                print(f"# {msg.type} message".upper() + "\n" + msg.content)
            else:
                print(make_colorful(msg.type, msg.content))


# AgentAction is not json serializable due to it does not inherit from NamedTuple anymore...
def replace_agent_action_with_list(x: Any) -> Any:
    if isinstance(x, AgentAction):
        return [x.tool, x.tool_input, x.log]
    elif isinstance(x, dict):
        return {k: replace_agent_action_with_list(v) for k, v in x.items()}
    elif isinstance(x, list):
        return [replace_agent_action_with_list(v) for v in x]
    elif isinstance(x, tuple):
        return tuple(replace_agent_action_with_list(v) for v in x)
    elif isinstance(x, set):
        return {replace_agent_action_with_list(v) for v in x}
    else:
        return x
