from uuid import UUID

from .models import Conversation, Message


async def create_conversation(
    conversation_id: UUID, system_prompt: str
) -> Conversation:
    conversation = Conversation(
        conversation_id=conversation_id,
        messages=[Message(role="system", content=system_prompt)],
    )
    await conversation.create()

    return conversation


async def get_conversation(conversation_id: UUID) -> Conversation | None:
    conversation = await Conversation.find_one(
        Conversation.conversation_id == conversation_id
    )
    return conversation


async def get_or_create_conversation(
    conversation_id: UUID, system_prompt: str
) -> Conversation:
    conversation = await get_conversation(conversation_id)

    if not conversation:
        conversation = await create_conversation(conversation_id, system_prompt)
    else:
        conversation.messages[0].content = system_prompt

    return conversation
