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

class Message:
    """Represents a single message in a conversation"""
    
    def __init__(self, content: str, role: str, timestamp: Optional[float] = None):
        """
        Initialize a message
        
        Args:
            content: The message content
            role: The role of the message sender (user or assistant)
            timestamp: Optional timestamp (defaults to current time)
        """
        self.content = content
        self.role = role
        self.timestamp = timestamp or time.time()
        
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            "content": self.content,
            "role": self.role,
            "timestamp": self.timestamp
        }
        
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Message':
        """Create a message from dictionary"""
        return cls(
            content=data["content"],
            role=data["role"],
            timestamp=data.get("timestamp", time.time())
        )

class Conversation:
    """Represents a conversation with message history"""
    
    def __init__(self, conversation_id: Optional[str] = None, max_history: int = 50):
        """
        Initialize a conversation
        
        Args:
            conversation_id: Optional conversation ID (defaults to a new UUID)
            max_history: Maximum number of messages to keep in history
        """
        self.conversation_id = conversation_id or str(uuid.uuid4())
        self.messages: List[Message] = []
        self.max_history = max_history
        self.created_at = time.time()
        self.last_updated = time.time()
        
    def add_message(self, content: str, role: str) -> Message:
        """
        Add a message to the conversation
        
        Args:
            content: The message content
            role: The role of the message sender (user or assistant)
            
        Returns:
            Message: The created message
        """
        message = Message(content, role)
        self.messages.append(message)
        
        # Trim history if needed
        if len(self.messages) > self.max_history:
            self.messages = self.messages[-self.max_history:]
            
        self.last_updated = time.time()
        return message
        
    def get_history(self) -> List[Dict[str, Any]]:
        """Get the conversation history as a list of dictionaries"""
        return [msg.to_dict() for msg in self.messages]
        
    def get_context_string(self) -> str:
        """Get the conversation history as a formatted string for context"""
        context = []
        for msg in self.messages:
            prefix = "User: " if msg.role == "user" else "Sam: "
            context.append(f"{prefix}{msg.content}")
        return "\n".join(context)
        
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            "conversation_id": self.conversation_id,
            "messages": [msg.to_dict() for msg in self.messages],
            "created_at": self.created_at,
            "last_updated": self.last_updated
        }
        
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Conversation':
        """Create a conversation from dictionary"""
        conversation = cls(
            conversation_id=data["conversation_id"],
            max_history=len(data.get("messages", []))
        )
        conversation.created_at = data.get("created_at", time.time())
        conversation.last_updated = data.get("last_updated", time.time())
        conversation.messages = [Message.from_dict(msg) for msg in data.get("messages", [])]
        return conversation

class ConversationMemoryManager:
    """Manages conversations and their history"""
    
    def __init__(self, max_conversations: int = 1000, max_history: int = 50):
        """
        Initialize the conversation memory manager
        
        Args:
            max_conversations: Maximum number of conversations to store
            max_history: Maximum number of messages per conversation
        """
        self.conversations: Dict[str, Conversation] = {}
        self.max_conversations = max_conversations
        self.max_history = max_history
        
    def get_conversation(self, conversation_id: str) -> Optional[Conversation]:
        """
        Get a conversation by ID
        
        Args:
            conversation_id: The conversation ID
            
        Returns:
            Optional[Conversation]: The conversation if found, None otherwise
        """
        return self.conversations.get(conversation_id)
        
    def create_conversation(self) -> Conversation:
        """
        Create a new conversation
        
        Returns:
            Conversation: The created conversation
        """
        conversation = Conversation(max_history=self.max_history)
        self.conversations[conversation.conversation_id] = conversation
        
        # Clean up old conversations if needed
        if len(self.conversations) > self.max_conversations:
            # Sort by last updated time and remove oldest
            sorted_convos = sorted(
                self.conversations.items(), 
                key=lambda x: x[1].last_updated
            )
            for conv_id, _ in sorted_convos[:len(self.conversations) - self.max_conversations]:
                del self.conversations[conv_id]
                
        return conversation
        
    def add_message(self, conversation_id: str, content: str, role: str) -> Optional[Message]:
        """
        Add a message to a conversation
        
        Args:
            conversation_id: The conversation ID
            content: The message content
            role: The role of the message sender
            
        Returns:
            Optional[Message]: The created message if successful, None otherwise
        """
        conversation = self.get_conversation(conversation_id)
        if not conversation:
            return None
            
        return conversation.add_message(content, role)
        
    def get_history(self, conversation_id: str) -> List[Dict[str, Any]]:
        """
        Get the history of a conversation
        
        Args:
            conversation_id: The conversation ID
            
        Returns:
            List[Dict[str, Any]]: The conversation history
        """
        conversation = self.get_conversation(conversation_id)
        if not conversation:
            return []
            
        return conversation.get_history()
        
    def get_context_string(self, conversation_id: str) -> str:
        """
        Get the conversation history as a formatted string for context
        
        Args:
            conversation_id: The conversation ID
            
        Returns:
            str: The formatted conversation history
        """
        conversation = self.get_conversation(conversation_id)
        if not conversation:
            return ""
            
        return conversation.get_context_string()

# Create a singleton instance
conversation_memory = ConversationMemoryManager()