"""
Agent Memory - Server-side message storage for inbox-based agents.

This module provides the AgentMemory abstraction that stores messages
and supports add_message, get_messages, and get_latest_message operations.
"""

from datetime import datetime
from typing import List, Optional
from pydantic import BaseModel


class InboxMessage(BaseModel):
    """Message structure for inbox."""
    from_agent: str
    content: str
    timestamp: Optional[str] = None


class AgentMemory:
    """
    Server-side memory storage for each agent.
    
    This class provides a simple in-memory storage for messages received
    by an agent. In production, this should be replaced with a proper
    database-backed implementation.
    
    Example:
        ```python
        memory = AgentMemory(agent_name="Writer")
        memory.add_message(from_agent="Editor", content="Great work!")
        messages = memory.get_messages()
        latest = memory.get_latest_message()
        ```
    """
    
    def __init__(self, agent_name: str):
        """
        Initialize memory for an agent.
        
        Args:
            agent_name: Name of the agent this memory belongs to
        """
        self.agent_name = agent_name
        self.messages: List[InboxMessage] = []
    
    def add_message(self, from_agent: str, content: str) -> None:
        """
        Add a message to the agent's memory.
        
        Args:
            from_agent: Name of the agent that sent the message
            content: Message content
        """
        message = InboxMessage(
            from_agent=from_agent,
            content=content,
            timestamp=datetime.now().isoformat()
        )
        self.messages.append(message)
    
    def get_messages(self) -> List[InboxMessage]:
        """
        Get all messages in the agent's memory.
        
        Returns:
            List of all messages (copy to prevent external modification)
        """
        return self.messages.copy()
    
    def get_latest_message(self) -> Optional[InboxMessage]:
        """
        Get the most recent message.
        
        Returns:
            Latest message, or None if no messages exist
        """
        return self.messages[-1] if self.messages else None
    
    def clear(self) -> None:
        """Clear all messages from memory."""
        self.messages.clear()
    
    def __repr__(self) -> str:
        """String representation."""
        return f"AgentMemory(agent_name='{self.agent_name}', messages={len(self.messages)})"

