import time
from typing import Any


class MemoryRecord:
    def __init__(self, role: str, content: str):
        self.role = role
        self.content = content
        self.timestamp = time.time()


class MemorySession:
    def __init__(self):
        self.records: list[MemoryRecord] = []

    def write(self, role: str, content: str) -> None:
        self.records.append(MemoryRecord(role, content))

    def read(self, limit: int = 10) -> list[dict]:
        return [
            {"role": r.role, "content": r.content}
            for r in self.records[-limit:]
        ]

    @property
    def size(self) -> int:
        return len(self.records)


class MemoryEngine:
    def __init__(self):
        self._users: dict[str, dict[str, MemorySession]] = {}

    def _ensure_session(self, user_id: str, session_id: str) -> MemorySession:
        if user_id not in self._users:
            self._users[user_id] = {}
        if session_id not in self._users[user_id]:
            self._users[user_id][session_id] = MemorySession()
        return self._users[user_id][session_id]

    def write(self, user_id: str, session_id: str, role: str, content: str) -> None:
        session = self._ensure_session(user_id, session_id)
        session.write(role, content)

    def read(self, user_id: str, session_id: str, limit: int = 10) -> list[dict]:
        session = self._ensure_session(user_id, session_id)
        return session.read(limit)

    def get_user_ids(self) -> list[str]:
        return list(self._users.keys())

    def get_session_ids(self, user_id: str) -> list[str]:
        return list(self._users.get(user_id, {}).keys())

    def get_stats(self) -> dict:
        total_users = len(self._users)
        total_sessions = sum(len(s) for s in self._users.values())
        total_records = sum(
            s.size for u in self._users.values() for s in u.values()
        )
        return {
            "users": total_users,
            "sessions": total_sessions,
            "records": total_records,
        }