from typing import List, Optional
from .message import Message, Role


class Memory:
    def __init__(self, max_tokens: int = 128000):
        self.messages: List[Message] = []
        self.max_tokens = max_tokens

    def add_message(self, message: Message):
        self.messages.append(message)

    def get_messages(self) -> List[Message]:
        return self.messages

    def clear(self):
        self.messages = []

    def get_openai_messages(self) -> List[dict]:
        return [m.to_openai_format() for m in self.messages]

    def add_user_message(self, content: str):
        self.add_message(Message(role=Role.USER, content=content))

    def add_assistant_message(
        self, content: Optional[str] = None, tool_calls: Optional[List[dict]] = None
    ):
        self.add_message(
            Message(role=Role.ASSISTANT, content=content, tool_calls=tool_calls)
        )

    def add_tool_message(self, tool_call_id: str, content: str, name: str):
        self.add_message(
            Message(
                role=Role.TOOL, content=content, name=name, tool_call_id=tool_call_id
            )
        )
