from typing import List, Dict
import time

class MemoryManager:
    """
    A user-specific memory manager mimicking RAG state.
    Stores conversation history and context per user for the API.
    """
    def __init__(self, max_history: int = 20, session_timeout_seconds: int = 3600):
        self.max_history = max_history
        self.session_timeout_seconds = session_timeout_seconds
        # Shape: { user_id: { "last_active": timestamp, "history": [{role, content}] } }
        self._memory_store: Dict[str, Dict] = {}

    def _cleanup_stale_sessions(self):
        """Removes sessions that haven't been active within the timeout."""
        current_time = time.time()
        stale_users = [
            user_id for user_id, data in self._memory_store.items()
            if (current_time - data["last_active"]) > self.session_timeout_seconds
        ]
        for user_id in stale_users:
            del self._memory_store[user_id]

    def add_interaction(self, user_id: str, prompt: str, response: str, modality: str = "text"):
        """Adds a multi-turn interaction to a user's isolated memory."""
        self._cleanup_stale_sessions()
        
        if user_id not in self._memory_store:
            self._memory_store[user_id] = {
                "last_active": time.time(),
                "history": []
            }
            
        history = self._memory_store[user_id]["history"]
        
        # Append interaction
        history.append({
            "role": "user",
            "content": prompt,
            "modality_hint": modality
        })
        history.append({
            "role": "assistant",
            "content": response,
            "modality_hint": modality
        })
        
        # Prune if over max_history (keep last N turns)
        if len(history) > (self.max_history * 2):
            self._memory_store[user_id]["history"] = history[-(self.max_history * 2):]
            
        self._memory_store[user_id]["last_active"] = time.time()

    def get_context(self, user_id: str) -> List[Dict[str, str]]:
        """Retrieves the interaction history to be injected into the LLM context."""
        self._cleanup_stale_sessions()
        
        if user_id in self._memory_store:
            self._memory_store[user_id]["last_active"] = time.time()
            return self._memory_store[user_id]["history"]
        return []

    def clear_memory(self, user_id: str):
        """Manually wipe a user's memory."""
        if user_id in self._memory_store:
            del self._memory_store[user_id]
