# src/conversation.py
from collections import deque
from typing import List, Dict, Optional

class ConversationMemory:
    """
    Role-aware conversation memory stored as a deque of dicts:
      {"role": "user"|"assistant"|"system", "content": "..."}

    Default max_len keeps the most recent messages.
    """

    def __init__(self, max_len: int = 60):
        self.max_len = max_len
        self.history = deque(maxlen=max_len)

    def add(self, user: str, bot: str):
        """Add user message and assistant reply as two messages."""
        if user is not None:
            self.history.append({"role": "user", "content": user})
        if bot is not None:
            self.history.append({"role": "assistant", "content": bot})

    def add_message(self, role: str, content: str):
        if role not in ("user", "assistant", "system"):
            raise ValueError("role must be 'user', 'assistant' or 'system'")
        self.history.append({"role": role, "content": content})

    def get_formatted(self, separator: str = "\n") -> str:
        out_lines = []
        for msg in self.history:
            role_label = "System" if msg["role"] == "system" else ("User" if msg["role"] == "user" else "Assistant")
            out_lines.append(f"{role_label}: {msg['content']}")
        return separator.join(out_lines) + (separator if out_lines else "")

    def clear(self):
        self.history.clear()

    def as_list(self) -> List[Dict[str, str]]:
        return list(self.history)

    def trim_to_recent_pairs(self, max_pairs: int = 12):
        """
        Keep roughly the last `max_pairs` (user+assistant) pairs.
        """
        items = list(self.history)
        keep = items[-(2 * max_pairs):] if len(items) > 2 * max_pairs else items
        self.history = deque(keep, maxlen=self.max_len)