from typing import Annotated, Any, Literal
from uuid import UUID

from beanie import Document, Indexed
from pydantic import BaseModel

type Role = Literal["user", "assistant", "system", "tool"]


class Function(BaseModel):
    name: str
    arguments: str


class Tool(BaseModel):
    id: str
    type: Literal["function"]
    function: Function


class Message(BaseModel):
    role: Role
    content: str | None = None
    tool_call_id: str | None = None
    tool_calls: list[Tool] | None = None


class Conversation(Document):
    conversation_id: Annotated[UUID, Indexed(unique=True)]
    messages: list[Message]

    class Settings:
        name = "conversations"

    def to_openai_format(self, limit: int, langchain_compat: bool = False):
        history = [m.model_dump(exclude_none=True) for m in self.messages]

        if langchain_compat:
            for msg in history:
                if "tool_calls" in msg:
                    msg["additional_kwargs"] = {"tool_calls": msg.pop("tool_calls")}
                if "content" not in msg:
                    msg["content"] = ""

        if len(history) > limit:
            # Truncate the history to the last `limit` messages
            # Always keep the first message (system prompt)
            history = history[:1] + history[-limit:]

        return history

    def add(
        self,
        role: Role,
        *,
        content: str | None = None,
        tool_call_id: str | None = None,
        tool_calls: list[Any] | None = None,
    ):
        self.messages.append(
            Message(
                role=role,
                content=content,
                tool_call_id=tool_call_id,
                tool_calls=tool_calls,
            )
        )
