import time
from typing import Dict, List, Any, Optional, Union

class MemoryManager:
    """Memory management for the autonomous AI agent
    
    This module provides capabilities for:
    1. Storing and retrieving conversation history
    2. Managing context windows
    3. Implementing forgetting mechanisms
    4. Prioritizing important information
    """
    
    def __init__(self, max_history_length: int = 20):
        """Initialize the memory manager
        
        Args:
            max_history_length: Maximum number of conversation turns to store
        """
        self.conversation_history = []
        self.max_history_length = max_history_length
        self.important_facts = []
        self.max_facts = 50
        self.session_data = {}
    
    def add_message(self, role: str, content: str) -> None:
        """Add a message to the conversation history
        
        Args:
            role: The role of the message sender (user or assistant)
            content: The content of the message
        """
        self.conversation_history.append({
            "role": role,
            "content": content,
            "timestamp": time.time()
        })
        
        # Trim history if it gets too long
        if len(self.conversation_history) > self.max_history_length * 2:
            self.conversation_history = self.conversation_history[-self.max_history_length*2:]
    
    def get_conversation_history(self, max_turns: Optional[int] = None) -> List[Dict[str, Any]]:
        """Get the conversation history
        
        Args:
            max_turns: Maximum number of turns to retrieve (None for all)
            
        Returns:
            List of conversation messages
        """
        if max_turns is None:
            return self.conversation_history
        else:
            # Calculate the number of messages (2 messages per turn)
            max_messages = max_turns * 2
            return self.conversation_history[-max_messages:]
    
    def format_conversation_for_prompt(self, max_turns: Optional[int] = None) -> str:
        """Format the conversation history for inclusion in a prompt
        
        Args:
            max_turns: Maximum number of turns to include
            
        Returns:
            Formatted conversation string
        """
        history = self.get_conversation_history(max_turns)
        formatted = ""
        
        for msg in history:
            formatted += f"{msg['role']}: {msg['content']}\n"
        
        return formatted
    
    def add_important_fact(self, fact: str, source: str) -> None:
        """Add an important fact to memory
        
        Args:
            fact: The important fact to remember
            source: The source of the fact (e.g., user, inference)
        """
        self.important_facts.append({
            "fact": fact,
            "source": source,
            "timestamp": time.time()
        })
        
        # Trim facts if they get too numerous
        if len(self.important_facts) > self.max_facts:
            self.important_facts = self.important_facts[-self.max_facts:]
    
    def get_important_facts(self) -> List[Dict[str, Any]]:
        """Get the list of important facts
        
        Returns:
            List of important facts
        """
        return self.important_facts
    
    def format_facts_for_prompt(self) -> str:
        """Format important facts for inclusion in a prompt
        
        Returns:
            Formatted facts string
        """
        if not self.important_facts:
            return ""
        
        formatted = "Important information I know about the user and context:\n"
        
        # Sort facts by timestamp (newest first)
        sorted_facts = sorted(self.important_facts, key=lambda x: x.get('timestamp', 0), reverse=True)
        
        # Group facts by source
        user_facts = [fact for fact in sorted_facts if fact.get('source') == 'user']
        inference_facts = [fact for fact in sorted_facts if fact.get('source') == 'inference']
        
        # Add user facts first (they're more reliable)
        for i, fact in enumerate(user_facts):
            formatted += f"{i+1}. {fact['fact']} (from user)\n"
        
        # Then add inference facts
        start_idx = len(user_facts) + 1
        for i, fact in enumerate(inference_facts):
            formatted += f"{start_idx + i}. {fact['fact']} (inferred)\n"
        
        return formatted
    
    def store_session_data(self, key: str, value: Any) -> None:
        """Store data for the current session
        
        Args:
            key: The key to store the data under
            value: The data to store
        """
        self.session_data[key] = {
            "value": value,
            "timestamp": time.time()
        }
    
    def get_session_data(self, key: str) -> Optional[Any]:
        """Retrieve data from the current session
        
        Args:
            key: The key to retrieve data for
            
        Returns:
            The stored data, or None if not found
        """
        if key in self.session_data:
            return self.session_data[key]["value"]
        else:
            return None
    
    def clear_conversation_history(self) -> None:
        """Clear the conversation history"""
        self.conversation_history = []
    
    def clear_all_memory(self) -> None:
        """Clear all memory (conversation history, facts, and session data)"""
        self.conversation_history = []
        self.important_facts = []
        self.session_data = {}
    
    def get_memory_stats(self) -> Dict[str, Any]:
        """Get statistics about the agent's memory usage
        
        Returns:
            Dictionary containing memory statistics
        """
        return {
            "conversation_turns": len(self.conversation_history) // 2,
            "important_facts": len(self.important_facts),
            "session_data_keys": list(self.session_data.keys()),
            "memory_usage": {
                "conversation": len(str(self.conversation_history)),
                "facts": len(str(self.important_facts)),
                "session": len(str(self.session_data))
            }
        }