"""Short per-chat conversation history, kept in memory."""

from __future__ import annotations

import threading
from collections import deque
from typing import Literal, TypedDict


class ChatTurn(TypedDict):
    role: Literal["user", "assistant"]
    content: str


class ConversationMemory:
    """Keeps the last `max_messages` user/assistant messages of each chat.

    Tool calls and tool results are not stored: the next turn starts from plain text,
    which keeps prompts small and avoids replaying stale order or stock data.
    """

    def __init__(self, max_messages: int, max_chats: int = 10_000) -> None:
        if max_messages < 2:
            raise ValueError("max_messages must be at least 2")
        self._max_messages = max_messages
        self._max_chats = max_chats
        self._chats: dict[int, deque[ChatTurn]] = {}
        self._lock = threading.Lock()

    def history(self, chat_id: int) -> list[ChatTurn]:
        with self._lock:
            return list(self._chats.get(chat_id, ()))

    def add_exchange(self, chat_id: int, user_text: str, assistant_text: str) -> None:
        with self._lock:
            turns = self._chats.pop(chat_id, None) or deque(maxlen=self._max_messages)
            turns.append({"role": "user", "content": user_text})
            turns.append({"role": "assistant", "content": assistant_text})
            # Re-insert so the dict stays ordered from least to most recently used.
            self._chats[chat_id] = turns
            while len(self._chats) > self._max_chats:
                self._chats.pop(next(iter(self._chats)))

    def clear(self, chat_id: int) -> None:
        with self._lock:
            self._chats.pop(chat_id, None)