"""Async Mem0 storage implementation for Upsonic agent framework.

This module provides an asynchronous Mem0 storage backend that supports both:
- Self-hosted Mem0 (via AsyncMemory class from mem0 library)
- Mem0 Platform (via AsyncMemoryClient class from mem0 library)

Mem0 stores data as memories with metadata, which we use to simulate
table-like behavior for sessions and user memories.
"""
from __future__ import annotations

import json
import os
import time
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union

if TYPE_CHECKING:
    from mem0 import AsyncMemory, AsyncMemoryClient
    from mem0.configs.base import MemoryConfig
    from upsonic.session.base import Session, SessionType

from upsonic.storage.base import AsyncStorage
from upsonic.storage.mem0.utils import (
    apply_pagination,
    build_session_filters,
    build_user_memory_filters,
    deserialize_session_from_mem0,
    deserialize_session_to_object,
    deserialize_user_memory_from_mem0,
    generate_session_memory_id,
    generate_user_memory_id,
    serialize_session_to_mem0,
    serialize_user_memory_to_mem0,
    sort_records_by_field,
)
from upsonic.storage.schemas import UserMemory
from upsonic.utils.logging_config import get_logger

_logger = get_logger("upsonic.storage.mem0.async")


class AsyncMem0Storage(AsyncStorage):
    """Asynchronous Mem0 storage implementation for Upsonic agent framework.
    
    This storage backend uses the Mem0 library for memory persistence.
    It supports both self-hosted Mem0 (AsyncMemory class) and Mem0 Platform
    (AsyncMemoryClient class).
    
    Mem0 stores data as memories with content and metadata. This implementation
    uses metadata filtering to simulate table-like behavior for:
    - Sessions (AgentSession, TeamSession, WorkflowSession)
    - User memories
    
    Client Creation Modes:
        1. Pass an existing memory_client (AsyncMemory or AsyncMemoryClient instance)
        2. Pass api_key to create an AsyncMemoryClient (Platform)
        3. Pass config to create an AsyncMemory instance (Self-hosted)
    
    Example usage with internal AsyncMemoryClient creation (Platform):
        ```python
        storage = AsyncMem0Storage(api_key="your-api-key")
        
        # With additional platform options
        storage = AsyncMem0Storage(
            api_key="your-api-key",
            org_id="your-org-id",
            project_id="your-project-id",
        )
        ```
    
    Example usage with internal AsyncMemory creation (Self-hosted):
        ```python
        from mem0.configs.base import MemoryConfig
        
        config = MemoryConfig(...)
        storage = AsyncMem0Storage(config=config)
        ```
    
    Example usage with external client:
        ```python
        from mem0 import AsyncMemory, AsyncMemoryClient
        
        # With self-hosted
        memory = AsyncMemory()
        storage = AsyncMem0Storage(memory_client=memory)
        
        # With platform
        client = AsyncMemoryClient(api_key="your-api-key")
        storage = AsyncMem0Storage(memory_client=client)
        ```
    """

    def __init__(
        self,
        memory_client: Optional[Union["AsyncMemory", "AsyncMemoryClient"]] = None,
        # Platform client (AsyncMemoryClient) parameters
        api_key: Optional[str] = None,
        host: Optional[str] = None,
        org_id: Optional[str] = None,
        project_id: Optional[str] = None,
        # Self-hosted (AsyncMemory) parameters
        config: Optional["MemoryConfig"] = None,
        # Common parameters
        session_table: Optional[str] = None,
        user_memory_table: Optional[str] = None,
        default_user_id: str = "upsonic_default",
        id: Optional[str] = None,
    ) -> None:
        """
        Initialize the async Mem0 storage backend.
        
        The client can be created in three ways (in priority order):
        1. Pass an existing memory_client (AsyncMemory or AsyncMemoryClient instance)
        2. Pass api_key to create an AsyncMemoryClient (Platform)
        3. Pass config to create an AsyncMemory instance (Self-hosted)
        
        Args:
            memory_client: Pre-configured AsyncMemory or AsyncMemoryClient instance.
            api_key: Mem0 Platform API key. Creates AsyncMemoryClient if provided.
                    Can also use MEM0_API_KEY environment variable.
            host: Mem0 Platform host URL (default: "https://api.mem0.ai").
            org_id: Mem0 Platform organization ID.
            project_id: Mem0 Platform project ID.
            config: MemoryConfig for self-hosted AsyncMemory instance.
            session_table: Name of the session table (used in metadata for filtering).
            user_memory_table: Name of the user memory table (used in metadata).
            default_user_id: Default user_id for Mem0 operations (required by Mem0).
            id: Unique identifier for this storage instance.
        
        Raises:
            ValueError: If no valid client configuration is provided.
            ImportError: If mem0 library is not installed.
        """
        super().__init__(
            session_table=session_table,
            user_memory_table=user_memory_table,
            id=id,
        )
        
        # Create or use provided memory client
        self.memory_client: Union["AsyncMemory", "AsyncMemoryClient"] = self._create_or_use_client(
            memory_client=memory_client,
            api_key=api_key,
            host=host,
            org_id=org_id,
            project_id=project_id,
            config=config,
        )
        self.default_user_id: str = default_user_id
        
        # Determine client type for API compatibility
        self._is_platform_client: bool = self._check_is_platform_client()
        
        _logger.debug(
            f"Initialized AsyncMem0Storage with "
            f"{'platform' if self._is_platform_client else 'self-hosted'} client"
        )

    def _create_or_use_client(
        self,
        memory_client: Optional[Union["AsyncMemory", "AsyncMemoryClient"]],
        api_key: Optional[str],
        host: Optional[str],
        org_id: Optional[str],
        project_id: Optional[str],
        config: Optional["MemoryConfig"],
    ) -> Union["AsyncMemory", "AsyncMemoryClient"]:
        """
        Create a new async memory client or use the provided one.
        
        Priority order:
        1. Use memory_client if provided
        2. Create AsyncMemoryClient if api_key is provided
        3. Create AsyncMemory if config is provided
        4. Create AsyncMemory with default config
        
        Args:
            memory_client: Existing client instance
            api_key: Platform API key
            host: Platform host URL
            org_id: Platform org ID
            project_id: Platform project ID
            config: Self-hosted Memory config
        
        Returns:
            AsyncMemory or AsyncMemoryClient instance
        
        Raises:
            ValueError: If configuration is invalid
            ImportError: If mem0 is not installed
        """
        # 1. Use existing client if provided
        if memory_client is not None:
            _logger.debug("Using provided async memory client")
            return memory_client
        
        # Try to import mem0
        try:
            from mem0 import AsyncMemory, AsyncMemoryClient
        except ImportError:
            raise ImportError(
                "mem0 library is not installed. Please install it using "
                "`pip install mem0ai`"
            )
        
        # 2. Create AsyncMemoryClient if api_key provided (or from env)
        effective_api_key = api_key or os.getenv("MEM0_API_KEY")
        if effective_api_key:
            _logger.debug("Creating AsyncMemoryClient (platform)")
            client_kwargs: Dict[str, Any] = {"api_key": effective_api_key}
            if host:
                client_kwargs["host"] = host
            if org_id:
                client_kwargs["org_id"] = org_id
            if project_id:
                client_kwargs["project_id"] = project_id
            return AsyncMemoryClient(**client_kwargs)
        
        # 3. Create AsyncMemory with config if provided
        if config is not None:
            _logger.debug("Creating AsyncMemory (self-hosted) with custom config")
            return AsyncMemory(config=config)
        
        # 4. Create AsyncMemory with default config
        _logger.debug("Creating AsyncMemory (self-hosted) with default config")
        return AsyncMemory()

    def _check_is_platform_client(self) -> bool:
        """Check if the memory client is a platform client or self-hosted."""
        client_class_name = self.memory_client.__class__.__name__
        return client_class_name == "AsyncMemoryClient"

    # ======================== Table Management ========================

    async def table_exists(self, table_name: str) -> bool:
        """
        Check if a table exists (always returns True for Mem0).
        
        Mem0 doesn't have explicit tables - data is organized via metadata.
        This method returns True as tables are "virtual" and always available.
        
        Args:
            table_name: Name of the table to check.
        
        Returns:
            Always True for Mem0 storage.
        """
        return True

    async def _create_all_tables(self) -> None:
        """
        Create all required tables (no-op for Mem0).
        
        Mem0 doesn't require explicit table creation. Tables are virtual
        and defined through metadata filtering.
        """
        _logger.debug("Mem0 does not require explicit table creation")

    async def close(self) -> None:
        """
        Close the Mem0 client connection.
        
        For AsyncMemoryClient (platform), closes the underlying httpx client.
        For AsyncMemory (self-hosted), performs any necessary cleanup.
        """
        if self._is_platform_client:
            # AsyncMemoryClient has an async_client that may need closing
            if hasattr(self.memory_client, 'async_client'):
                await self.memory_client.async_client.aclose()
        _logger.debug("AsyncMem0 storage closed")

    # ======================== Session Methods ========================

    async def aupsert_session(
        self,
        session: "Session",
        deserialize: bool = True,
    ) -> Optional[Union["Session", Dict[str, Any]]]:
        """
        Insert or update a session in Mem0 (async).
        
        Args:
            session: The session to upsert. Must have session_id set.
            deserialize: If True, return deserialized Session object.
                        If False, return raw dictionary.
        
        Returns:
            The upserted session (deserialized or dict based on flag).
        
        Raises:
            ValueError: If session_id is not provided.
        """
        if not hasattr(session, "session_id") or session.session_id is None:
            raise ValueError("Session must have session_id set for upsert")
        
        try:
            # Serialize session for Mem0
            serialized = serialize_session_to_mem0(session, self.session_table_name)
            memory_id = serialized["memory_id"]
            
            # Check if memory already exists to determine if this is update or insert
            existing = await self._aget_memory_by_id(memory_id)
            
            if existing:
                # Update existing memory
                await self._aupdate_memory(
                    memory_id=memory_id,
                    content=serialized["content"],
                    metadata=serialized["metadata"],
                )
                _logger.debug(f"Updated session: {session.session_id}")
            else:
                # Add new memory
                await self._aadd_memory(
                    memory_id=memory_id,
                    content=serialized["content"],
                    metadata=serialized["metadata"],
                )
                _logger.debug(f"Added session: {session.session_id}")
            
            # Retrieve and return the stored session
            stored = await self._aget_memory_by_id(memory_id)
            if stored is None:
                return None
            
            session_dict = deserialize_session_from_mem0(stored)
            
            if not deserialize:
                return session_dict
            
            return deserialize_session_to_object(session_dict)
        
        except Exception as e:
            _logger.error(f"Error upserting session: {e}")
            raise e

    async def aupsert_sessions(
        self,
        sessions: List["Session"],
        deserialize: bool = True,
    ) -> List[Union["Session", Dict[str, Any]]]:
        """
        Bulk insert or update multiple sessions (async).
        
        Args:
            sessions: List of sessions to upsert. Each must have session_id set.
            deserialize: If True, return deserialized Session objects.
                        If False, return raw dictionaries.
        
        Returns:
            List of upserted sessions.
        
        Raises:
            ValueError: If any session is missing session_id.
        """
        if not sessions:
            return []
        
        # Validate all sessions have session_id
        for session in sessions:
            if not hasattr(session, "session_id") or session.session_id is None:
                raise ValueError("All sessions must have session_id set for upsert")
        
        results: List[Union["Session", Dict[str, Any]]] = []
        
        for session in sessions:
            try:
                result = await self.aupsert_session(session, deserialize=deserialize)
                if result is not None:
                    results.append(result)
            except Exception as e:
                _logger.error(f"Error upserting session {session.session_id}: {e}")
                raise e
        
        _logger.debug(f"Upserted {len(results)} sessions")
        return results

    async def aget_session(
        self,
        session_id: Optional[str] = None,
        session_type: Optional["SessionType"] = None,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        deserialize: bool = True,
    ) -> Optional[Union["Session", Dict[str, Any]]]:
        """
        Get a session from Mem0 (async).
        
        Args:
            session_id: ID of the session to retrieve. If None, returns latest session.
            session_type: Type of session (AGENT, TEAM, WORKFLOW).
            user_id: Filter by user ID.
            agent_id: Filter by agent ID.
            deserialize: If True, return deserialized Session object.
                        If False, return raw dictionary.
        
        Returns:
            The session if found, None otherwise.
        """
        try:
            if session_id is not None:
                # Direct lookup by memory_id
                memory_id = generate_session_memory_id(session_id, self.session_table_name)
                stored = await self._aget_memory_by_id(memory_id)
                
                if stored is None:
                    return None
                
                session_dict = deserialize_session_from_mem0(stored)
                
                if not deserialize:
                    return session_dict
                
                return deserialize_session_to_object(session_dict, session_type)
            
            # Search with filters and return latest
            filters = build_session_filters(
                table_name=self.session_table_name,
                session_type=session_type,
                user_id=user_id,
                agent_id=agent_id,
            )
            
            records = await self._asearch_memories(filters=filters)
            
            if not records:
                return None
            
            # Sort by created_at desc and get first
            sorted_records = sort_records_by_field(records, "created_at", "desc")
            
            if not sorted_records:
                return None
            
            session_dict = deserialize_session_from_mem0(sorted_records[0])
            
            if not deserialize:
                return session_dict
            
            return deserialize_session_to_object(session_dict, session_type)
        
        except Exception as e:
            _logger.error(f"Error getting session: {e}")
            raise e

    async def aget_sessions(
        self,
        session_ids: Optional[List[str]] = None,
        session_type: Optional["SessionType"] = None,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        sort_by: Optional[str] = None,
        sort_order: Optional[str] = None,
        deserialize: bool = True,
    ) -> Union[List["Session"], Tuple[List[Dict[str, Any]], int]]:
        """
        Get multiple sessions from Mem0 (async).
        
        Args:
            session_ids: List of session IDs to retrieve. If None, returns all sessions.
            session_type: Filter by session type.
            user_id: Filter by user ID.
            agent_id: Filter by agent ID.
            limit: Maximum number of sessions to return.
            offset: Number of sessions to skip.
            sort_by: Field to sort by.
            sort_order: Sort order ('asc' or 'desc').
            deserialize: If True, return deserialized Session objects.
                        If False, return tuple of (raw dictionaries, total count).
        
        Returns:
            List of sessions or tuple of (list of dicts, total count).
        """
        try:
            records: List[Dict[str, Any]] = []
            
            if session_ids is not None and len(session_ids) > 0:
                # Fetch specific sessions by ID
                for session_id in session_ids:
                    memory_id = generate_session_memory_id(
                        session_id, self.session_table_name
                    )
                    stored = await self._aget_memory_by_id(memory_id)
                    if stored:
                        records.append(stored)
            else:
                # Search with filters
                filters = build_session_filters(
                    table_name=self.session_table_name,
                    session_type=session_type,
                    user_id=user_id,
                    agent_id=agent_id,
                )
                records = await self._asearch_memories(filters=filters)
            
            if not records:
                return [] if deserialize else ([], 0)
            
            # Get total count before pagination
            total_count = len(records)
            
            # Sort records
            sorted_records = sort_records_by_field(
                records,
                sort_by=sort_by or "created_at",
                sort_order=sort_order or "desc",
            )
            
            # Apply pagination
            paginated_records = apply_pagination(sorted_records, limit, offset)
            
            # Deserialize all records
            session_dicts = [
                deserialize_session_from_mem0(r) for r in paginated_records
            ]
            
            if not deserialize:
                return session_dicts, total_count
            
            return [
                deserialize_session_to_object(s, session_type)
                for s in session_dicts
            ]
        
        except Exception as e:
            _logger.error(f"Error getting sessions: {e}")
            raise e

    async def adelete_session(self, session_id: str) -> bool:
        """
        Delete a session from Mem0 (async).
        
        Args:
            session_id: ID of the session to delete.
        
        Returns:
            True if deleted successfully, False otherwise.
        
        Raises:
            ValueError: If session_id is not provided.
        """
        if not session_id:
            raise ValueError("session_id is required for delete")
        
        try:
            memory_id = generate_session_memory_id(session_id, self.session_table_name)
            return await self._adelete_memory(memory_id)
        
        except Exception as e:
            _logger.error(f"Error deleting session: {e}")
            raise e

    async def adelete_sessions(self, session_ids: List[str]) -> int:
        """
        Delete multiple sessions from Mem0 (async).
        
        Args:
            session_ids: List of session IDs to delete.
        
        Returns:
            Number of sessions deleted.
        
        Raises:
            ValueError: If session_ids is empty.
        """
        if not session_ids:
            raise ValueError("session_ids is required and cannot be empty")
        
        deleted_count = 0
        
        for session_id in session_ids:
            try:
                if await self.adelete_session(session_id):
                    deleted_count += 1
            except Exception as e:
                _logger.warning(f"Error deleting session {session_id}: {e}")
        
        _logger.debug(f"Deleted {deleted_count} sessions")
        return deleted_count

    # ======================== User Memory Methods ========================

    async def aupsert_user_memory(
        self,
        user_memory: UserMemory,
        deserialize: bool = True,
    ) -> Optional[Union[UserMemory, Dict[str, Any]]]:
        """
        Insert or update user memory in Mem0 (async).
        
        Args:
            user_memory: The memory data to store. Must have 'user_id' key.
            deserialize: If True, return UserMemory object. If False, return raw dict.
        
        Returns:
            The upserted user memory record.
        
        Raises:
            ValueError: If user_id is not provided in user_memory.
        """
        if not user_memory.user_id:
            raise ValueError("user_memory must have user_id set for upsert")
        
        try:
            # Serialize user memory for Mem0
            serialized = serialize_user_memory_to_mem0(
                user_id=user_memory.user_id,
                user_memory=user_memory.user_memory,
                table_name=self.user_memory_table_name,
                agent_id=user_memory.agent_id,
                team_id=user_memory.team_id,
            )
            memory_id = serialized["memory_id"]
            
            # Check if memory already exists
            existing = await self._aget_memory_by_id(memory_id)
            
            if existing:
                # Update existing - preserve created_at
                existing_data = deserialize_user_memory_from_mem0(existing)
                created_at = existing_data.get("created_at", user_memory.created_at or int(time.time()))
                
                # Update metadata with preserved created_at
                serialized["metadata"]["created_at"] = created_at
                
                # Update the _data in metadata with preserved created_at
                data_str = serialized["metadata"].get("_data")
                if data_str:
                    data_dict = json.loads(data_str)
                    data_dict["created_at"] = created_at
                    serialized["metadata"]["_data"] = json.dumps(data_dict)
                
                await self._aupdate_memory(
                    memory_id=memory_id,
                    content=serialized["content"],
                    metadata=serialized["metadata"],
                )
                _logger.debug(f"Updated user memory: {user_memory.user_id}")
            else:
                # Add new memory
                await self._aadd_memory(
                    memory_id=memory_id,
                    content=serialized["content"],
                    metadata=serialized["metadata"],
                )
                _logger.debug(f"Added user memory: {user_memory.user_id}")
            
            # Retrieve and return the stored memory
            stored = await self._aget_memory_by_id(memory_id)
            if stored is None:
                return None
            
            memory_dict = deserialize_user_memory_from_mem0(stored)
            if deserialize:
                return UserMemory.from_dict(memory_dict)
            return memory_dict
        
        except Exception as e:
            _logger.error(f"Error upserting user memory: {e}")
            raise e

    async def aupsert_user_memories(
        self,
        user_memories: List[UserMemory],
        deserialize: bool = True,
    ) -> List[Union[UserMemory, Dict[str, Any]]]:
        """
        Bulk insert or update multiple user memories (async).
        
        Args:
            user_memories: List of user memory dicts. Each must have 'user_id'.
            deserialize: If True, return UserMemory objects. If False, return raw dicts.
        
        Returns:
            List of upserted user memory records.
        
        Raises:
            ValueError: If any record is missing user_id.
        """
        if not user_memories:
            return []
        
        # Validate all have user_id
        for memory in user_memories:
            if not memory.user_id:
                raise ValueError("All user memories must have user_id set")
        
        results: List[Union[UserMemory, Dict[str, Any]]] = []
        
        for memory in user_memories:
            try:
                result = await self.aupsert_user_memory(
                    user_memory=memory,
                    deserialize=deserialize,
                )
                if result is not None:
                    results.append(result)
            except Exception as e:
                _logger.error(
                    f"Error upserting user memory for {memory.user_id}: {e}"
                )
                raise e
        
        _logger.debug(f"Upserted {len(results)} user memories")
        return results

    async def aget_user_memory(
        self,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        team_id: Optional[str] = None,
        deserialize: bool = True,
    ) -> Optional[Union[UserMemory, Dict[str, Any]]]:
        """
        Get user memory from Mem0 (async).
        
        Args:
            user_id: User ID to retrieve. If None, returns latest memory.
            agent_id: Filter by agent ID.
            team_id: Filter by team ID.
            deserialize: If True, return UserMemory object. If False, return raw dict.
        
        Returns:
            The user memory record if found, None otherwise.
        """
        try:
            if user_id is not None:
                # Direct lookup by memory_id
                memory_id = generate_user_memory_id(
                    user_id, self.user_memory_table_name
                )
                stored = await self._aget_memory_by_id(memory_id)
                
                if stored is None:
                    return None
                
                memory_dict = deserialize_user_memory_from_mem0(stored)
                if deserialize:
                    return UserMemory.from_dict(memory_dict)
                return memory_dict
            
            # Search with filters and return latest
            filters = build_user_memory_filters(
                table_name=self.user_memory_table_name,
                agent_id=agent_id,
                team_id=team_id,
            )
            
            records = await self._asearch_memories(filters=filters)
            
            if not records:
                return None
            
            # Sort by updated_at desc and get first
            sorted_records = sort_records_by_field(records, "updated_at", "desc")
            
            if not sorted_records:
                return None
            
            memory_dict = deserialize_user_memory_from_mem0(sorted_records[0])
            if deserialize:
                return UserMemory.from_dict(memory_dict)
            return memory_dict
        
        except Exception as e:
            _logger.error(f"Error getting user memory: {e}")
            raise e

    async def aget_user_memories(
        self,
        user_ids: Optional[List[str]] = None,
        agent_id: Optional[str] = None,
        team_id: Optional[str] = None,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        deserialize: bool = True,
    ) -> Union[List[UserMemory], Tuple[List[Dict[str, Any]], int]]:
        """
        Get multiple user memories from Mem0 (async).
        
        Args:
            user_ids: List of user IDs to retrieve. If None, returns all.
            agent_id: Filter by agent ID.
            team_id: Filter by team ID.
            limit: Maximum number of records to return.
            offset: Number of records to skip.
            deserialize: If True, return list of UserMemory objects. 
                         If False, return tuple of (list of dicts, total count).
        
        Returns:
            List of UserMemory objects if deserialize=True, otherwise tuple of (dicts, count).
        """
        try:
            records: List[Dict[str, Any]] = []
            
            if user_ids is not None and len(user_ids) > 0:
                # Fetch specific user memories by ID
                for user_id in user_ids:
                    memory_id = generate_user_memory_id(
                        user_id, self.user_memory_table_name
                    )
                    stored = await self._aget_memory_by_id(memory_id)
                    if stored:
                        records.append(stored)
            else:
                # Search with filters
                filters = build_user_memory_filters(
                    table_name=self.user_memory_table_name,
                    agent_id=agent_id,
                    team_id=team_id,
                )
                records = await self._asearch_memories(filters=filters)
            
            if not records:
                if deserialize:
                    return []
                return [], 0
            
            # Get total count before pagination
            total_count = len(records)
            
            # Sort by updated_at desc
            sorted_records = sort_records_by_field(records, "updated_at", "desc")
            
            # Apply pagination
            paginated_records = apply_pagination(sorted_records, limit, offset)
            
            # Deserialize all records
            memory_dicts = [
                deserialize_user_memory_from_mem0(r) for r in paginated_records
            ]
            
            if deserialize:
                return [UserMemory.from_dict(m) for m in memory_dicts]
            return memory_dicts, total_count
        
        except Exception as e:
            _logger.error(f"Error getting user memories: {e}")
            raise e

    async def adelete_user_memory(self, user_id: str) -> bool:
        """
        Delete user memory from Mem0 (async).
        
        Args:
            user_id: User ID of the memory to delete.
        
        Returns:
            True if deleted successfully, False otherwise.
        
        Raises:
            ValueError: If user_id is not provided.
        """
        if not user_id:
            raise ValueError("user_id is required for delete")
        
        try:
            memory_id = generate_user_memory_id(user_id, self.user_memory_table_name)
            return await self._adelete_memory(memory_id)
        
        except Exception as e:
            _logger.error(f"Error deleting user memory: {e}")
            raise e

    async def adelete_user_memories(self, user_ids: List[str]) -> int:
        """
        Delete multiple user memories from Mem0 (async).
        
        Args:
            user_ids: List of user IDs to delete.
        
        Returns:
            Number of records deleted.
        
        Raises:
            ValueError: If user_ids is empty.
        """
        if not user_ids:
            raise ValueError("user_ids is required and cannot be empty")
        
        deleted_count = 0
        
        for user_id in user_ids:
            try:
                if await self.adelete_user_memory(user_id):
                    deleted_count += 1
            except Exception as e:
                _logger.warning(f"Error deleting user memory {user_id}: {e}")
        
        _logger.debug(f"Deleted {deleted_count} user memories")
        return deleted_count

    # ======================== Utility Methods ========================

    async def aclear_all(self) -> None:
        """
        Clear all data from all tables (async).
        
        This removes all sessions and user memories from the storage.
        Use with caution.
        """
        try:
            # Get all sessions and delete them
            session_filters = build_session_filters(table_name=self.session_table_name)
            session_records = await self._asearch_memories(filters=session_filters)
            
            for record in session_records:
                memory_id = record.get("id") or record.get("memory_id")
                if memory_id:
                    await self._adelete_memory(memory_id)
            
            _logger.debug(f"Cleared {len(session_records)} sessions")
            
            # Get all user memories and delete them
            user_memory_filters = build_user_memory_filters(
                table_name=self.user_memory_table_name
            )
            user_memory_records = await self._asearch_memories(
                filters=user_memory_filters
            )
            
            for record in user_memory_records:
                memory_id = record.get("id") or record.get("memory_id")
                if memory_id:
                    await self._adelete_memory(memory_id)
            
            _logger.debug(f"Cleared {len(user_memory_records)} user memories")
            _logger.info("Cleared all data from AsyncMem0 storage")
        
        except Exception as e:
            _logger.error(f"Error clearing all data: {e}")
            raise e

    # ======================== Private Async Mem0 API Methods ========================
    # 
    # These methods handle the API differences between:
    # - AsyncMemoryClient (Platform): Supports memory_id on add(), uses text param
    # - AsyncMemory (Self-hosted): Does NOT support memory_id on add(), uses data param
    #
    # For self-hosted AsyncMemory, we use metadata filtering to find records by our 
    # session_id since we cannot specify our own memory_id during creation.

    async def _aadd_memory(
        self,
        memory_id: str,
        content: str,
        metadata: Dict[str, Any],
    ) -> Optional[Dict[str, Any]]:
        """
        Add a new memory to Mem0 (async).
        
        Args:
            memory_id: Unique identifier for the memory (stored in metadata for lookup)
            content: The content to store
            metadata: Metadata for filtering (must include _upsonic_memory_id)
        
        Returns:
            The created memory record or None
        
        Note:
            Both Platform and Self-hosted use metadata-based ID tracking because:
            - Platform API's memory_id handling is unreliable for custom IDs
            - Self-hosted doesn't support memory_id parameter at all
            We store our custom ID in metadata._upsonic_memory_id for lookup.
        """
        try:
            if self._is_platform_client:
                # Platform API: Store in metadata, use async_mode=False for immediate availability
                # Note: async_mode controls Mem0's background processing, not Python async/await.
                # With async_mode=False, data is immediately available after the await completes.
                result = await self.memory_client.add(
                    messages=content,
                    user_id=self.default_user_id,
                    metadata=metadata,
                    async_mode=False,  # Sync processing for immediate data availability
                )
            else:
                # Self-hosted API: Store in metadata, use infer=False
                result = await self.memory_client.add(
                    messages=content,
                    user_id=self.default_user_id,
                    metadata=metadata,
                    infer=False,
                )
            
            _logger.debug(f"Added memory: {memory_id}")
            return result
        
        except Exception as e:
            _logger.error(f"Error adding memory {memory_id}: {e}")
            raise e

    async def _aupdate_memory(
        self,
        memory_id: str,
        content: str,
        metadata: Dict[str, Any],
    ) -> Optional[Dict[str, Any]]:
        """
        Update an existing memory in Mem0 (async).
        
        Args:
            memory_id: Our custom memory ID (stored in metadata as _upsonic_memory_id)
            content: The new content
            metadata: The new metadata
        
        Returns:
            The updated memory record or None
        
        Note:
            Both Platform and Self-hosted use metadata-based ID resolution.
            Platform uses 'text' param, Self-hosted uses 'data' param.
        """
        try:
            # Find the actual Mem0 memory_id from our custom ID
            actual_memory_id = await self._aresolve_memory_id(memory_id)
            
            if actual_memory_id is None:
                _logger.warning(f"Memory not found for update: {memory_id}")
                return None
            
            if self._is_platform_client:
                # Platform API: Uses 'text' parameter for content
                result = await self.memory_client.update(
                    memory_id=actual_memory_id,
                    text=content,
                    metadata=metadata,
                )
            else:
                # Self-hosted API: Uses 'data' parameter for content
                result = await self.memory_client.update(
                    memory_id=actual_memory_id,
                    data=content,
                )
            
            _logger.debug(f"Updated memory: {memory_id}")
            return result
        
        except Exception as e:
            _logger.error(f"Error updating memory {memory_id}: {e}")
            raise e

    async def _aget_memory_by_id(self, memory_id: str) -> Optional[Dict[str, Any]]:
        """
        Get a specific memory by its ID (async).
        
        Args:
            memory_id: Our custom memory ID (stored in metadata as _upsonic_memory_id)
        
        Returns:
            The memory record or None if not found
        
        Note:
            Both Platform and Self-hosted search by _upsonic_memory_id in metadata.
        """
        try:
            return await self._afind_memory_by_upsonic_id(memory_id)
        except Exception as e:
            _logger.debug(f"Error getting memory {memory_id}: {e}")
            return None

    async def _adelete_memory(self, memory_id: str) -> bool:
        """
        Delete a memory by its ID (async).
        
        Args:
            memory_id: Our custom memory ID (stored in metadata as _upsonic_memory_id)
        
        Returns:
            True if deleted, False otherwise
        """
        try:
            # Find actual Mem0 ID first
            actual_memory_id = await self._aresolve_memory_id(memory_id)
            
            if actual_memory_id is None:
                _logger.debug(f"Memory not found for delete: {memory_id}")
                return False
            
            await self.memory_client.delete(memory_id=actual_memory_id)
            
            _logger.debug(f"Deleted memory: {memory_id}")
            return True
        
        except Exception as e:
            if "not found" in str(e).lower():
                return False
            _logger.warning(f"Error deleting memory {memory_id}: {e}")
            return False

    async def _aresolve_memory_id(self, upsonic_memory_id: str) -> Optional[str]:
        """
        Resolve our custom memory_id to the actual Mem0 memory_id (async).
        
        We store our ID in metadata and need to find the actual
        Mem0-generated ID for operations.
        
        Args:
            upsonic_memory_id: Our custom memory ID (stored in metadata)
        
        Returns:
            The actual Mem0 memory_id or None if not found
        """
        memory = await self._afind_memory_by_upsonic_id(upsonic_memory_id)
        if memory:
            return memory.get("id")
        return None

    async def _afind_memory_by_upsonic_id(
        self, upsonic_memory_id: str
    ) -> Optional[Dict[str, Any]]:
        """
        Find a memory record by our custom upsonic_memory_id in metadata (async).
        
        Args:
            upsonic_memory_id: Our custom memory ID to search for
        
        Returns:
            The memory record or None if not found
        """
        try:
            if self._is_platform_client:
                # Platform API: Use proper filter format with AND wrapper
                result = await self.memory_client.get_all(
                    filters={
                        "AND": [
                            {"user_id": self.default_user_id},
                            {"metadata": {"_upsonic_memory_id": upsonic_memory_id}},
                        ]
                    }
                )
            else:
                # Self-hosted API: Get all for user and filter manually
                result = await self.memory_client.get_all(
                    user_id=self.default_user_id,
                    limit=10000,
                )
            
            # Handle different return formats
            if isinstance(result, dict):
                memories = result.get("results", result.get("memories", []))
            elif isinstance(result, list):
                memories = result
            else:
                memories = []
            
            # For Platform API with proper filters, we should get exact match
            if self._is_platform_client and memories:
                return memories[0] if memories else None
            
            # For Self-hosted, find the memory with matching _upsonic_memory_id in metadata
            for memory in memories:
                metadata = memory.get("metadata", {})
                if metadata.get("_upsonic_memory_id") == upsonic_memory_id:
                    return memory
            
            return None
        
        except Exception as e:
            _logger.debug(f"Error finding memory by upsonic ID: {e}")
            return None

    async def _asearch_memories(
        self,
        filters: Optional[Dict[str, Any]] = None,
        query: Optional[str] = None,
        limit: int = 1000,
    ) -> List[Dict[str, Any]]:
        """
        Search memories with optional filters (async).
        
        Args:
            filters: Metadata filters (for client-side filtering)
            query: Optional search query (for semantic search)
            limit: Maximum number of results
        
        Returns:
            List of matching memory records
        """
        try:
            if self._is_platform_client:
                # Platform API: Use proper filter format with AND wrapper
                api_filters: Dict[str, Any] = {
                    "AND": [
                        {"user_id": self.default_user_id},
                    ]
                }
                
                # Add metadata filters if provided
                if filters:
                    for key, value in filters.items():
                        api_filters["AND"].append({"metadata": {key: value}})
                
                result = await self.memory_client.get_all(filters=api_filters)
            else:
                # Self-hosted API: Get all for user
                result = await self.memory_client.get_all(
                    user_id=self.default_user_id,
                    limit=limit,
                )
            
            # Handle different return formats
            if isinstance(result, dict):
                memories = result.get("results", result.get("memories", []))
            elif isinstance(result, list):
                memories = result
            else:
                memories = []
            
            # For Self-hosted, apply metadata filters manually
            if not self._is_platform_client and filters and memories:
                filtered = []
                for memory in memories:
                    metadata = memory.get("metadata", {})
                    
                    # Check if all filter conditions match
                    matches = True
                    for key, value in filters.items():
                        if metadata.get(key) != value:
                            matches = False
                            break
                    
                    if matches:
                        filtered.append(memory)
                
                return filtered
            
            return memories
        
        except Exception as e:
            _logger.error(f"Error searching memories: {e}")
            return []

