"""
Memory Manager
Provides utilities for managing agent memory and state.
"""

from typing import Any, Dict, List, Optional
import logging
import json

from agentscope.message import Msg

logger = logging.getLogger(__name__)


class MemoryManager:
    """
    Manages agent memory and conversation history.
    Provides utilities for loading, saving, and manipulating memory.
    """
    
    def __init__(self, max_history: int = 50):
        """ 初始化 the memory manager.
        
        Args:
            max_history: Maximum number of messages to keep in history
        """
        self.max_history = max_history
        self._history: List[Msg] = []
        self._metadata: Dict[str, Any] = {}
    
    def add_message(self, message: Msg):
        """
        Add a message to the history.
        
        Args:
            message: The message to add
        """
        self._history.append(message)
        
        # Trim history if exceeds max
        if len(self._history) > self.max_history:
            self._history = self._history[-self.max_history:]
    
    def add_messages(self, messages: List[Msg]):
        """
        Add multiple messages to the history.
        
        Args:
            messages: List of messages to add
        """
        for msg in messages:
            self.add_message(msg)
    
    def get_history(self, limit: Optional[int] = None) -> List[Msg]:
        """ 获取 conversation history.
        
        Args:
            limit: Optional limit on number of messages to return
        
        Returns:
            List of messages
        """
        if limit is not None:
            return self._history[-limit:]
        return self._history.copy()
    
    def clear(self):
        """ 清除 all history."""
        self._history.clear()
        logger.info("Memory cleared")
    
    def set_metadata(self, key: str, value: Any):
        """ 集合 metadata value.
        
        Args:
            key: Metadata key
            value: Metadata value
        """
        self._metadata[key] = value
    
    def get_metadata(self, key: str, default: Any = None) -> Any:
        """ 获取 metadata value.
        
        Args:
            key: Metadata key
            default: Default value if key not found
        
        Returns:
            Metadata value or default
        """ 
        return self._metadata.get(key, default)
    
    def to_dict(self) -> Dict[str, Any]:
        """ 序列化 memory state to dictionary.
        
        Returns:
            Dictionary representation of memory state
        """ 
        return {
            "history": [
                {
                    "name": msg.name,
                    "content": msg.content,
                    "role": msg.role,
                }
                for msg in self._history
            ],
            "metadata": self._metadata,
        }
    
    def from_dict(self, data: Dict[str, Any]):
        """
        Load memory state from dictionary.
        
        Args:
            data: Dictionary representation of memory state
        """
        self._history.clear()
        for msg_data in data.get("history", []):
            self._history.append(Msg(
                name=msg_data.get("name", "unknown"),
                content=msg_data.get("content", ""),
                role=msg_data.get("role", "user"),
            ))
        self._metadata = data.get("metadata", {})
    
    @staticmethod
    def convert_messages_to_agentscope(
        messages: List[Dict[str, Any]]
    ) -> List[Msg]:
        """ 转换 standard message format to AgentScope Msg format.
        
        Args:
            messages: List of message dictionaries with 'role' and 'content'
        
        Returns:
            List of Msg objects
        """
        result = []
        for m in messages:
            role = m.get("role", "user")
            # 映射 roles
            if role in ["assistant", "model"]:
                name = "assistant"
                ag_role = "assistant"
            elif role == "system":
                name = "system"
                ag_role = "system"
            else:
                name = "user"
                ag_role = "user"
            
            content = m.get("content", "")
            result.append(Msg(name=name, content=content, role=ag_role))
        
        return result


class SharedMemory:
    """
    Shared memory for agent collaboration (DEPRECATED).
    
    This class is deprecated in favor of single-agent architecture.
    Kept for reference only.
    
    Original purpose: Allows agents to share information during complex tasks.
    """
    
    def __init__(self):
        self._storage: Dict[str, Any] = {}
        self._agent_outputs: Dict[str, List[Any]] = {}
    
    def store(self, key: str, value: Any):
        """Store a value with the given key."""
        self._storage[key] = value
    
    def retrieve(self, key: str, default: Any = None) -> Any:
        """Retrieve a value by key.""" 
        return self._storage.get(key, default)
    
    def add_agent_output(self, agent_name: str, output: Any):
        """Add output from a specific agent."""
        if agent_name not in self._agent_outputs:
            self._agent_outputs[agent_name] = []
        self._agent_outputs[agent_name].append(output)
    
    def get_agent_outputs(self, agent_name: str) -> List[Any]:
        """ 获取 all outputs from a specific agent.""" 
        return self._agent_outputs.get(agent_name, [])
    
    def get_latest_agent_output(self, agent_name: str) -> Optional[Any]:
        """ 获取 the latest output from a specific agent."""
        outputs = self._agent_outputs.get(agent_name, [])
        return outputs[-1] if outputs else None
    
    def clear(self):
        """ 清除 all shared memory."""
        self._storage.clear()
        self._agent_outputs.clear()
