import json

import typer

from termcolor import colored

app = typer.Typer()


def pretty_print_conversation(messages):
    role_to_color = {
        "system": "red",
        "user": "green",
        "assistant": "blue",
        "function": "magenta",
    }
    formatted_messages = []
    for message in messages:
        if message["role"] == "function":
            formatted_messages.append(
                f"function ({message['name']}): {message['content']}\n"
            )
        else:
            assistant_content = (
                message["function_call"]
                if message.get("function_call")
                else message["content"]
            )
            role_to_message = {
                "system": f"system: {message['content']}\n",
                "user": f"user: {message['content']}\n",
                "assistant": f"assistant: {assistant_content}\n",
            }
            formatted_messages.append(role_to_message[message["role"]])

    for formatted_message in formatted_messages:
        role = messages[formatted_messages.index(formatted_message)]["role"]
        color = role_to_color[role]
        print(colored(formatted_message, color))


@app.command()
def main(
    messages_path: str,
):
    with open(messages_path) as f:
        messages = json.load(f)

    pretty_print_conversation(messages)


if __name__ == "__main__":
    app()