from limin import (
    ModelConfiguration,
    generate_text_completion_for_conversation,
    Conversation,
    Message,
)

from typing import Literal


def invert_role(
    role: Literal["system", "user", "assistant"],
) -> Literal["system", "user", "assistant"]:
    if role == "user":
        return "assistant"
    elif role == "assistant":
        return "user"
    return role


def get_inverted_conversation(conversation: Conversation) -> Conversation:
    inverted_conversation = Conversation()

    # Append an empty user message to the beginning of the conversation so that the first message is always a user message
    inverted_conversation.add_message(Message(role="user", content=""))

    for message in conversation.messages:
        inverted_conversation.add_message(
            Message(role=invert_role(message.role), content=message.content)
        )
    return inverted_conversation


class Character:
    def __init__(self, system_prompt: str, model_configuration: ModelConfiguration):
        self.system_prompt = system_prompt
        self.model_configuration = model_configuration

    async def generate_initial_message(self) -> str:
        conversation = Conversation(
            messages=[
                Message(role="system", content=self.system_prompt),
                Message(
                    role="user",
                    content="Generate an initial message for a conversation.",
                ),
            ]
        )

        completion = await generate_text_completion_for_conversation(
            conversation, model_configuration=self.model_configuration
        )
        return completion.content

    async def generate_next_message(self, conversation: Conversation) -> str:
        # The character system prompt is included.
        # For the user character, this will be the user character's system prompt.
        # For the assistant character, this will be the assistant character's system prompt.
        messages = Conversation(
            messages=[
                Message(role="system", content=self.system_prompt),
                *conversation.messages,
            ]
        )

        completion = await generate_text_completion_for_conversation(
            messages, model_configuration=self.model_configuration
        )
        return completion.content


async def talk(
    user_character: Character,
    assistant_character: Character,
    n_turns: int,
    initial_message: str | None = None,
) -> Conversation:
    conversation = Conversation()

    # The first message is generated by the user character
    initial_message = (
        await user_character.generate_initial_message()
        if initial_message is None
        else initial_message
    )
    conversation.add_message(Message(role="user", content=initial_message))

    turn_count = 0
    while turn_count < n_turns:
        # The assistant character responds
        assistant_character_message = await assistant_character.generate_next_message(
            conversation
        )
        conversation.add_message(
            Message(role="assistant", content=assistant_character_message)
        )

        turn_count += 1
        if turn_count == n_turns:
            break

        # The user character continues the conversation.
        # Note that from the perspective of the user character, the conversation is inverted.
        inverted_conversation = get_inverted_conversation(conversation)
        user_character_message = await user_character.generate_next_message(
            inverted_conversation
        )
        conversation.add_message(Message(role="user", content=user_character_message))

    return conversation
