# Define memory function tools
def search_memory(query: str, user_id: str) -> dict:
    """Search through past conversations and memories"""
    filters = {"user_id": user_id}
    memories = mem0.search(query, filters=filters)
    # ... return formatted results

def save_memory(content: str, user_id: str) -> dict:
    """Save important information to memory"""
    result = mem0.add([{"role": "user", "content": content}], user_id=user_id)
    # ... return status

# Add as tools to agent
agent = Agent(
    name="personal_assistant",
    tools=[search_memory, save_memory]  # Manual tool management
)

from google.adk.memory.base_memory_service import BaseMemoryService, SearchMemoryResponse
from google.adk.memory.memory_entry import MemoryEntry
from google.adk.tools.preload_memory_tool import PreloadMemoryTool
from google.adk.agents.callback_context import CallbackContext
from mem0 import MemoryClient

class Mem0MemoryService(BaseMemoryService):
    """MemoryService implementation using Mem0 Platform."""
    
    def __init__(self, api_key: Optional[str] = None):
        super().__init__()
        api_key = api_key or os.environ.get("MEM0_API_KEY")
        self._client = MemoryClient(api_key=api_key) if api_key else None
    
    @override
    async def search_memory(
        self, *, app_name: str, user_id: str, query: str
    ) -> SearchMemoryResponse:
        """Search for relevant memories."""
        if not self.is_enabled():
            return SearchMemoryResponse(memories=[])
        
        filters = {"user_id": user_id}
        memories = self._client.search(query, filters=filters, limit=5)
        
        # Convert Mem0 memories to ADK MemoryEntry format
        memory_entries = []
        for mem in memories.get("results", []):
            memory_text = mem.get("memory", "")
            if memory_text:
                content = Content(parts=[Part(text=memory_text)])
                memory_entry = MemoryEntry(
                    content=content,
                    author=mem.get("metadata", {}).get("author", "user"),
                    timestamp=mem.get("created_at"),
                )
                memory_entries.append(memory_entry)
        
        return SearchMemoryResponse(memories=memory_entries)
    
    @override
    async def add_session_to_memory(self, session: Session) -> None:
        """Add session content to memory."""
        if not self.is_enabled():
            return
        
        user_id = getattr(session, "user_id", None)
        if not user_id:
            return
        
        # Extract messages from session events
        messages = []
        for event in session.events:
            if event.content and event.content.parts:
                role = getattr(event.content, "role", "user")
                # Map ADK's "model" role to Mem0's "assistant"
                if role == "model":
                    role = "assistant"
                
                text_parts = [
                    part.text for part in event.content.parts 
                    if hasattr(part, "text") and part.text
                ]
                if text_parts:
                    messages.append({
                        "role": role,
                        "content": " ".join(text_parts)
                    })
        
        if messages:
            self._client.add(messages, user_id=user_id)

# Setup with automatic memory management
async def auto_save_session_to_memory_callback(callback_context: CallbackContext) -> None:
    """Automatically save session to memory after agent completes."""
    invocation_context = getattr(callback_context, "_invocation_context", None)
    if not invocation_context:
        return
    
    session = getattr(invocation_context, "session", None)
    memory_service = getattr(invocation_context, "memory_service", None)
    
    if session and memory_service:
        await memory_service.add_session_to_memory(session)

# Create memory service
memory_service = Mem0MemoryService()

# Create agent with PreloadMemoryTool (automatically injects memories)
agent = Agent(
    name="personal_assistant",
    tools=[PreloadMemoryTool()],  # Automatically searches and injects memories
    after_agent_callback=auto_save_session_to_memory_callback,  # Auto-saves sessions
)

# Pass memory_service to Runner
runner = Runner(
    agent=agent,
    session_service=session_service,
    memory_service=memory_service,  # Native ADK integration
    app_name="memory_assistant",
)