"""
Asynchronous memory management implementation

This module provides the asynchronous memory management interface.
"""

import asyncio
import logging
import warnings
import hashlib
import json
from typing import Any, Dict, List, Optional, Union
from datetime import datetime
from powermem.utils.utils import get_current_datetime
from copy import deepcopy

from .base import MemoryBase
from ..configs import MemoryConfig
from ..storage.factory import VectorStoreFactory, GraphStoreFactory
from ..storage.adapter import StorageAdapter, SubStorageAdapter
from ..intelligence.manager import IntelligenceManager
from ..integrations.llm.factory import LLMFactory
from ..integrations.embeddings.factory import EmbedderFactory
from .telemetry import TelemetryManager
from .audit import AuditLogger
from ..intelligence.plugin import IntelligentMemoryPlugin, EbbinghausIntelligencePlugin
from ..utils.utils import (
    convert_config_object_to_dict,
    parse_vision_messages,
    llm_json_text_with_fallback,
    parse_fact_extraction_json,
    parse_memory_actions_json,
)
from ..prompts.intelligent_memory_prompts import (
    FACT_RETRIEVAL_PROMPT,
    FACT_EXTRACTION_PROMPT,
    get_memory_update_prompt,
    parse_messages_for_facts
)

logger = logging.getLogger(__name__)


def _forget_marker_updates() -> Dict[str, Any]:
    return {
        "should_forget": True,
        "marked_for_forgetting_at": get_current_datetime().isoformat(),
    }


def _auto_convert_config(config: Dict[str, Any]) -> Dict[str, Any]:
    """
    Convert legacy powermem config to format for compatibility.
    
    Now powermem uses field names directly.
    
    Args:
        config: Configuration dictionary (legacy format)
        
    Returns:
        configuration dictionary
    """
    if not config:
        return config

    # First, convert any ConfigObject instances to dicts
    config = convert_config_object_to_dict(config)

    # Check if legacy powermem format (has database or embedding)
    if "database" in config or ("llm" in config and "embedding" in config):
        converted = config.copy()

        # Convert llm
        if "llm" in config:
            converted["llm"] = config["llm"]
        
        # Convert embedding to embedder
        if "embedding" in config:
            converted["embedder"] = config["embedding"]
            converted.pop("embedding", None)

        # Convert database to vector_store
        if "database" in config:
            db_config = config["database"]
            converted["vector_store"] = {
                "provider": db_config.get("provider", "oceanbase"),
                "config": db_config.get("config", {})
            }
            converted.pop("database", None)
        elif "vector_store" not in converted:
            converted["vector_store"] = {
                "provider": "oceanbase",
                "config": {}
            }
        
        logger.info("Converted legacy powermem config format")
        return converted

    return config


class AsyncMemory(MemoryBase):
    """
    Asynchronous memory management implementation.
    
    This class provides the main interface for asynchronous memory operations.
    """
    
    def __init__(
        self,
        config: Optional[Dict[str, Any] | MemoryConfig] = None,
        storage_type: Optional[str] = None,
        llm_provider: Optional[str] = None,
        embedding_provider: Optional[str] = None,
        agent_id: Optional[str] = None,
    ):
        """
        Initialize the async memory manager.

        Compatible with both dict config and MemoryConfig object.

        Args:
            config: Configuration dictionary or MemoryConfig object containing all settings.
                   Dict format supports style (llm, embedder, vector_store)
                   and powermem style (database, llm, embedding)
            storage_type: Type of storage backend to use (overrides config)
            llm_provider: LLM provider to use (overrides config)
            embedding_provider: Embedding provider to use (overrides config)
            agent_id: Agent identifier for multi-agent scenarios
        """
        # Handle MemoryConfig object or dict
        if isinstance(config, MemoryConfig):
            # Use MemoryConfig object directly
            self.memory_config = config
            # For backward compatibility, also store as dict
            self.config = config.to_dict()
        else:
            # Convert dict config
            dict_config = config or {}
            dict_config = _auto_convert_config(dict_config)
            self.config = dict_config
            self.memory_config = None
            logger.debug("Using dict config mode")

        self.agent_id = agent_id
        
        # Extract providers from config with fallbacks
        self.storage_type = storage_type or self._get_provider('vector_store', 'oceanbase')
        self.llm_provider = llm_provider or self._get_provider('llm', 'mock')
        self.embedding_provider = embedding_provider or self._get_provider('embedder', 'mock')
        
        # Initialize components
        vector_store_config = self._get_component_config('vector_store')
        vector_store = VectorStoreFactory.create(self.storage_type, vector_store_config)
        
        # Extract graph_store config
        self.enable_graph = self._get_graph_enabled()
        self.graph_store = None
        if self.enable_graph:
            logger.debug("Graph store enabled")
            graph_store_config = self.config.get("graph_store", {})
            if graph_store_config:
                provider = graph_store_config.get("provider", "oceanbase")
                config_to_pass = self.memory_config if self.memory_config else self.config
                self.graph_store = GraphStoreFactory.create(provider, config_to_pass)

        # Extract LLM config
        llm_config = self._get_component_config('llm')
        self.llm = LLMFactory.create(self.llm_provider, llm_config)
        
        # Extract embedder config
        embedder_config = self._get_component_config('embedder')
        # Pass vector_store_config so factory can extract embedding_model_dims for mock embeddings
        self.embedding = EmbedderFactory.create(self.embedding_provider, embedder_config, vector_store_config)
        
        # Initialize storage adapter with embedding service
        # Automatically select adapter based on sub_stores configuration
        sub_stores_list = self.config.get('sub_stores', [])
        if sub_stores_list and self.storage_type.lower() == 'oceanbase':
            # Use SubStorageAdapter if sub stores are configured and using OceanBase
            self.storage = SubStorageAdapter(vector_store, self.embedding)
            logger.info("Using SubStorageAdapter with sub-store support")
        else:
            # Use basic StorageAdapter for single store operations
            self.storage = StorageAdapter(vector_store, self.embedding)
            logger.info("Using basic StorageAdapter")
        
        self.intelligence = IntelligenceManager(self.config)
        telemetry_config = self.config.get("telemetry")
        if telemetry_config is None:
            telemetry_config = self.config
        self.telemetry = TelemetryManager(telemetry_config)
        audit_config = self.config.get("audit")
        if audit_config is None:
            audit_config = self.config
        self.audit = AuditLogger(audit_config)

        # Save custom prompts from config
        if self.memory_config:
            self.custom_fact_extraction_prompt = self.memory_config.custom_fact_extraction_prompt
            self.custom_update_memory_prompt = self.memory_config.custom_update_memory_prompt
        else:
            self.custom_fact_extraction_prompt = self.config.get('custom_fact_extraction_prompt')
            self.custom_update_memory_prompt = self.config.get('custom_update_memory_prompt')

        # Intelligent memory plugin (pluggable)
        merged_cfg = self._get_intelligent_memory_config()

        plugin_type = merged_cfg.get("plugin", "ebbinghaus")
        self._intelligence_plugin: Optional[IntelligentMemoryPlugin] = None
        if merged_cfg.get("enabled", False):
            try:
                if plugin_type == "ebbinghaus":
                    self._intelligence_plugin = EbbinghausIntelligencePlugin(merged_cfg)
                else:
                    logger.warning(f"Unknown intelligence plugin: {plugin_type}")
            except Exception as e:
                logger.warning(f"Failed to initialize intelligence plugin (async): {e}")
                self._intelligence_plugin = None

        
        # Sub stores configuration (support multiple)
        self.sub_stores_config: List[Dict] = []

        # Initialize sub stores
        self._init_sub_stores()

        logger.info(f"AsyncMemory initialized with storage: {self.storage_type}, LLM: {self.llm_provider}, agent: {self.agent_id or 'default'}")
        self.telemetry.capture_event("async_memory.init", {"storage_type": self.storage_type, "llm_provider": self.llm_provider, "agent_id": self.agent_id})
    
    async def initialize(self):
        """Initialize async components."""
        await self.storage.initialize_async()
    
    def _get_provider(self, component: str, default: str) -> str:
        """
        Helper method to get component provider uniformly.

        Args:
            component: Component name ('vector_store', 'llm', 'embedder')
            default: Default provider name

        Returns:
            Provider name string
        """
        if self.memory_config:
            component_obj = getattr(self.memory_config, component, None)
            return component_obj.provider if component_obj else default
        else:
            return self.config.get(component, {}).get('provider', default)

    def _is_llm_disabled(self) -> bool:
        """Return True when PowerMem is running without LLM-backed features."""
        return self.llm_provider == "noop" or getattr(self.llm, "is_noop", False) is True

    def _get_component_config(self, component: str) -> Dict[str, Any]:
        """
        Helper method to get component configuration uniformly.

        Args:
            component: Component name ('vector_store', 'llm', 'embedder', 'graph_store')

        Returns:
            Component configuration dictionary
        """
        if self.memory_config:
            component_obj = getattr(self.memory_config, component, None)
            return component_obj.config or {} if component_obj else {}
        else:
            return self.config.get(component, {}).get('config', {})

    def _get_graph_enabled(self) -> bool:
        """
        Helper method to get graph store enabled status.

        Returns:
            Boolean indicating whether graph store is enabled
        """
        if self.memory_config:
            return self.memory_config.graph_store.enabled if self.memory_config.graph_store else False
        else:
            graph_store_config = self.config.get('graph_store', {})
            return graph_store_config.get('enabled', False) if graph_store_config else False

    def _get_intelligent_memory_config(self) -> Dict[str, Any]:
        """
        Helper method to get intelligent memory configuration.
        Supports both "intelligence" and "intelligent_memory" config keys for backward compatibility.

        Returns:
            Merged intelligent memory configuration dictionary
        """
        if self.memory_config and self.memory_config.intelligent_memory:
            # Use MemoryConfig's intelligent_memory
            cfg = self.memory_config.intelligent_memory.model_dump()
            # Merge custom_importance_evaluation_prompt from top level if present
            if self.memory_config.custom_importance_evaluation_prompt:
                cfg["custom_importance_evaluation_prompt"] = self.memory_config.custom_importance_evaluation_prompt
            return cfg
        else:
            # Fallback to dict access
            intelligence_cfg = (self.config or {}).get("intelligence", {})
            intelligent_memory_cfg = (self.config or {}).get("intelligent_memory", {})
            merged_cfg = {**intelligence_cfg, **intelligent_memory_cfg}
            # Merge custom_importance_evaluation_prompt from top level if present
            if "custom_importance_evaluation_prompt" in self.config:
                merged_cfg["custom_importance_evaluation_prompt"] = self.config["custom_importance_evaluation_prompt"]
            return merged_cfg
    
    async def _extract_facts(self, messages: Any) -> List[str]:
        """
        Extract facts from messages using LLM asynchronously.
        
        Args:
            messages: Messages (list of dicts, single dict, or str)
            
        Returns:
            List of extracted facts
        """
        if self._is_llm_disabled():
            logger.info("LLM is disabled; skipping fact extraction.")
            return []

        try:
            # Parse messages into conversation format
            conversation = parse_messages_for_facts(messages)
            
            # Use custom prompt if provided, otherwise use default
            if self.custom_fact_extraction_prompt:
                system_prompt = self.custom_fact_extraction_prompt
                user_prompt = f"Input:\n{conversation}"
            else:
                system_prompt = FACT_RETRIEVAL_PROMPT
                user_prompt = f"Input:\n{conversation}"
            
            try:
                response = await asyncio.to_thread(
                    llm_json_text_with_fallback,
                    self.llm,
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt},
                    ],
                )
            except Exception as e:
                logger.error(f"Error in fact extraction: {e}")
                response = ""
            
            try:
                facts = parse_fact_extraction_json(str(response or ""))
                if not facts and str(response or "").strip():
                    logger.debug(
                        "Fact extraction parsed no facts from non-empty body; retrying without response_format"
                    )
                    response_plain = await asyncio.to_thread(
                        self.llm.generate_response,
                        messages=[
                            {"role": "system", "content": system_prompt},
                            {"role": "user", "content": user_prompt},
                        ],
                        response_format=None,
                    )
                    facts = parse_fact_extraction_json(str(response_plain or ""))
                logger.debug(f"Extracted {len(facts)} facts: {facts}")
                return facts
            except Exception as e:
                logger.error(f"Error in new_retrieved_facts: {e}")
                return []
                
        except Exception as e:
            logger.error(f"Error extracting facts: {e}")
            return []
    
    async def _decide_memory_actions(
        self, 
        new_facts: List[str], 
        existing_memories: List[Dict[str, Any]],
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """
        Use LLM to decide memory actions (ADD/UPDATE/DELETE/NONE) asynchronously.
        
        Args:
            new_facts: List of newly extracted facts
            existing_memories: List of existing memories with 'id' and 'text'
            user_id: User identifier
            agent_id: Agent identifier
            
        Returns:
            List of memory action dictionaries
        """
        if self._is_llm_disabled():
            logger.info("LLM is disabled; skipping memory action planning.")
            return []

        try:
            if not new_facts:
                logger.debug("No new facts to process")
                return []
            
            # Format existing memories for prompt
            old_memory = []
            for mem in existing_memories:
                # Support both "memory" and "content" field names for compatibility
                content = mem.get("memory", "") or mem.get("content", "")
                old_memory.append({
                    "id": mem.get("id", "unknown"),
                    "text": content
                })
            
            # Generate update prompt with custom prompt if provided
            custom_prompt = None
            if hasattr(self, 'custom_update_memory_prompt') and self.custom_update_memory_prompt:
                custom_prompt = self.custom_update_memory_prompt
            update_prompt = get_memory_update_prompt(old_memory, new_facts, custom_prompt)
            
            try:
                response = await asyncio.to_thread(
                    llm_json_text_with_fallback,
                    self.llm,
                    messages=[{"role": "user", "content": update_prompt}],
                )
            except Exception as e:
                logger.error(f"Error in new memory actions response: {e}")
                response = ""
            
            try:
                actions = parse_memory_actions_json(str(response or ""))
                if not actions and str(response or "").strip():
                    logger.debug(
                        "Memory actions JSON empty after parse; retrying without response_format"
                    )
                    response_plain = await asyncio.to_thread(
                        self.llm.generate_response,
                        messages=[{"role": "user", "content": update_prompt}],
                        response_format=None,
                    )
                    actions = parse_memory_actions_json(str(response_plain or ""))
                return actions
            except Exception as e:
                logger.error(f"Invalid JSON response: {e}")
                return []
                
        except Exception as e:
            logger.error(f"Error deciding memory actions: {e}")
            return []
    
    async def add(
        self,
        messages,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
        filters: Optional[Dict[str, Any]] = None,
        scope: Optional[str] = None,
        memory_type: Optional[str] = None,
        prompt: Optional[str] = None,
        infer: bool = True,
    ) -> Dict[str, Any]:
        """Add a new memory asynchronously with optional intelligent processing.
        
        Returns:
            Dict[str, Any]: A dictionary containing the add operation results with the following structure:
                - "results" (List[Dict]): List of memory operation results, where each result contains:
                    - "id" (int): Memory ID
                    - "memory" (str): The memory content
                    - "event" (str): Operation event type (e.g., "ADD", "UPDATE", "DELETE")
                    - "user_id" (str, optional): User ID associated with the memory
                    - "agent_id" (str, optional): Agent ID associated with the memory
                    - "run_id" (str, optional): Run ID associated with the memory
                    - "metadata" (Dict, optional): Metadata dictionary
                    - "created_at" (str, optional): Creation timestamp in ISO format
                    - "previous_memory" (str, optional): Previous memory content (for UPDATE events)
                - "relations" (Dict, optional): Graph relations if graph store is enabled, containing:
                    - "deleted_entities" (List): List of deleted graph entities
                    - "added_entities" (List): List of added graph entities
        """
        try:
            # Handle messages parameter
            if messages is None:
                raise ValueError("messages must be provided (str, dict, or list[dict])")
            
            # Normalize input format
            if isinstance(messages, str):
                messages = [{"role": "user", "content": messages}]
            elif isinstance(messages, dict):
                messages = [messages]
            elif not isinstance(messages, list):
                raise ValueError("messages must be str, dict, or list[dict]")
            
            # Vision-aware message processing
            llm_cfg = {}
            try:
                llm_cfg = (self.config or {}).get("llm", {}).get("config", {})
            except Exception:
                llm_cfg = {}
            if llm_cfg.get("enable_vision"):
                messages = parse_vision_messages(messages, self.llm, llm_cfg.get("vision_details"))
            else:
                messages = parse_vision_messages(messages)

            if isinstance(messages, list) and len(messages) == 0:
                return {"results": []}
            
            # Use self.agent_id as fallback if agent_id is not provided
            agent_id = agent_id or self.agent_id
            
            # Check if intelligent memory should be used
            llm_disabled = self._is_llm_disabled()
            use_infer = infer and isinstance(messages, list) and len(messages) > 0 and not llm_disabled
            if infer and llm_disabled:
                logger.info("LLM is disabled; falling back to simple add mode.")
            
            # If not using intelligent memory, fall back to simple mode
            if not use_infer:
                return await self._simple_add_async(messages, user_id, agent_id, run_id, metadata, filters, scope, memory_type, prompt)
            
            # Intelligent memory mode: extract facts, search similar memories, and consolidate
            return await self._intelligent_add_async(messages, user_id, agent_id, run_id, metadata, filters, scope, memory_type, prompt)
            
        except Exception as e:
            logger.error(f"Failed to add memory: {e}")
            self.telemetry.capture_event("memory.add.error", {"error": str(e)})
            raise
    
    async def _simple_add_async(
        self,
        messages,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
        filters: Optional[Dict[str, Any]] = None,
        scope: Optional[str] = None,
        memory_type: Optional[str] = None,
        prompt: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Simple add mode: direct storage without intelligence."""
        # Parse messages into content
        if isinstance(messages, str):
            content = messages
        elif isinstance(messages, dict):
            content = messages.get("content", "")
        elif isinstance(messages, list):
            content = "\n".join([msg.get("content", "") for msg in messages if isinstance(msg, dict) and msg.get("content")])
        else:
            raise ValueError("messages must be str, dict, or list[dict]")
        
        # Validate content is not empty
        if not content or not content.strip():
            logger.error(f"Cannot store empty content. Messages: {messages}")
            raise ValueError(f"Cannot create memory with empty content. Original messages: {messages}")
        
        # Select embedding service based on metadata (for sub-store routing)
        embedding_service = self._get_embedding_service(metadata)

        # Generate embedding asynchronously
        embedding = await asyncio.to_thread(embedding_service.embed, content, memory_action="add")
        
        # Disabled LLM-based importance evaluation to save tokens
        # Process with intelligence manager
        # enhanced_metadata = await self.intelligence.process_metadata_async(content, metadata)
        enhanced_metadata = metadata  # Use original metadata without LLM evaluation

        # Intelligent plugin annotations: merge into metadata for persistence
        if self._intelligence_plugin and self._intelligence_plugin.enabled:
            extra_fields = self._intelligence_plugin.on_add(content=content, metadata=enhanced_metadata)
            if extra_fields:
                if enhanced_metadata is None:
                    enhanced_metadata = {}
                enhanced_metadata = {**enhanced_metadata, **extra_fields}

        # Generate content hash for deduplication
        content_hash = hashlib.md5(content.encode('utf-8')).hexdigest()

        # Extract category from enhanced metadata if present
        category = ""
        if enhanced_metadata and isinstance(enhanced_metadata, dict):
            category = enhanced_metadata.get("category", "")
            # Remove category from metadata to avoid duplication
            enhanced_metadata = {k: v for k, v in enhanced_metadata.items() if k != "category"}
        if memory_type:
            category = memory_type
        if scope:
            if enhanced_metadata is None:
                enhanced_metadata = {"scope": scope}
            elif isinstance(enhanced_metadata, dict):
                enhanced_metadata = {**enhanced_metadata, "scope": scope}
            else:
                enhanced_metadata = {"scope": scope}

        # Final validation before storage
        if not content or not content.strip():
            raise ValueError(f"Refusing to store empty content. Original messages: {messages}")
        
        # Use self.agent_id as fallback if agent_id is not provided
        agent_id = agent_id or self.agent_id
        
        # Store in database asynchronously
        memory_data = {
            "content": content,
            "embedding": embedding,
            "user_id": user_id,
            "agent_id": agent_id,
            "run_id": run_id,
            "hash": content_hash,
            "category": category,
            "metadata": enhanced_metadata or {},
            "filters": filters or {},
            "created_at": get_current_datetime(),
            "updated_at": get_current_datetime(),
        }

        memory_id = await self.storage.add_memory_async(memory_data)
        
        # Log audit event
        await self.audit.log_event_async("memory.add", {
            "memory_id": memory_id,
            "user_id": user_id,
            "agent_id": agent_id,
            "content_length": len(content)
        }, user_id=user_id, agent_id=agent_id)
        
        # Capture telemetry
        self.telemetry.capture_event("memory.add", {
            "memory_id": memory_id,
            "user_id": user_id,
            "agent_id": agent_id
        })
        
        # Add to graph store and get relations (only if graph store is enabled)
        graph_result = None
        if self.enable_graph:
            graph_result = await self._add_to_graph_async(messages, filters, user_id, agent_id, run_id)
        
        result: Dict[str, Any] = {
            "results": [{
                "id": memory_id,
                "memory": content,
                "event": "ADD",
                "user_id": user_id,
                "agent_id": agent_id,
                "run_id": run_id,
                "metadata": metadata,
                "created_at": memory_data["created_at"].isoformat() if isinstance(memory_data["created_at"], datetime) else memory_data["created_at"],
            }]
        }
        if graph_result:
            result["relations"] = graph_result
        return result
    
    async def _intelligent_add_async(
        self,
        messages,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
        filters: Optional[Dict[str, Any]] = None,
        scope: Optional[str] = None,
        memory_type: Optional[str] = None,
        prompt: Optional[str] = None,
    ) -> Dict[str, Any]:
        """Intelligent add mode: extract facts, consolidate with existing memories."""
        # Use self.agent_id as fallback if agent_id is not provided
        agent_id = agent_id or self.agent_id
        
        # Get intelligent memory config to check fallback setting
        intelligent_config = self._get_intelligent_memory_config()
        fallback_to_simple = intelligent_config.get("fallback_to_simple_add", False)
        
        # Step 1: Extract facts from messages
        logger.info("Extracting facts from messages...")
        facts = await self._extract_facts(messages)
        
        if not facts:
            logger.debug("No facts extracted, skip intelligent add")
            if fallback_to_simple:
                logger.warning("No facts extracted from messages, falling back to simple add mode")
                return await self._simple_add_async(messages, user_id, agent_id, run_id, metadata, filters, scope, memory_type, prompt)
            return {"results": []}

        logger.info(f"Extracted {len(facts)} facts: {facts}")
        
        # Step 2: Search for similar memories for each fact
        existing_memories = []
        fact_embeddings = {}
        
        # Select embedding service based on metadata (for sub-store routing)
        embedding_service = self._get_embedding_service(metadata)

        for fact in facts:
            fact_embedding = await asyncio.to_thread(embedding_service.embed, fact, memory_action="add")
            fact_embeddings[fact] = fact_embedding
            
            # Merge metadata into filters for correct routing
            search_filters = filters.copy() if filters else {}
            if metadata:
                # Filter metadata to only include simple values (strings, numbers, booleans, None)
                # This prevents nested dicts like {'agent': {'agent_id': ...}} from causing issues
                # when OceanBase's build_condition tries to parse them as operators
                simple_metadata = {
                    k: v for k, v in metadata.items()
                    if not isinstance(v, (dict, list)) and k not in ['agent_id', 'user_id', 'run_id']
                }
                search_filters.update(simple_metadata)

            # Search for similar memories with reduced limit to reduce noise
            # Pass fact text to enable hybrid search for better results
            similar = await self.storage.search_memories_async(
                query_embedding=fact_embedding,
                user_id=user_id,
                agent_id=agent_id,
                run_id=run_id,
                filters=search_filters,
                limit=5,
                query=fact  # Enable hybrid search
            )
            existing_memories.extend(similar)
        
        # Improved deduplication: prefer memories with better similarity scores
        unique_memories = {}
        for mem in existing_memories:
            mem_id = mem.get("id")
            if mem_id and mem_id not in unique_memories:
                unique_memories[mem_id] = mem
            elif mem_id:
                # If duplicate ID, keep the one with better similarity (lower distance)
                existing = unique_memories.get(mem_id)
                mem_distance = mem.get("distance", float('inf'))
                existing_distance = existing.get("distance", float('inf')) if existing else float('inf')
                if mem_distance < existing_distance:
                    unique_memories[mem_id] = mem
        
        # Limit candidates to avoid LLM prompt overload
        existing_memories = list(unique_memories.values())[:10]  # Max 10 memories
        
        logger.info(f"Found {len(existing_memories)} existing memories to consider (after dedup and limiting)")
        
        # Mapping IDs with integers for handling ID hallucinations
        # Maps temporary string indices to real Snowflake IDs (integers)
        temp_uuid_mapping = {}
        for idx, item in enumerate(existing_memories):
            temp_uuid_mapping[str(idx)] = item["id"]
            existing_memories[idx]["id"] = str(idx)
        
        # Step 3: Let LLM decide memory actions (only if we have new facts)
        actions = []
        if facts:
            actions = await self._decide_memory_actions(facts, existing_memories, user_id, agent_id)
            logger.info(f"LLM decided on {len(actions)} memory actions")
        else:
            logger.debug("No new facts, skipping LLM decision step")
        
        # Step 4: Execute actions
        results = []
        action_counts = {"ADD": 0, "UPDATE": 0, "DELETE": 0, "NONE": 0}
        
        if not actions:
            logger.warning("No actions returned from LLM, skip intelligent add")
            if fallback_to_simple:
                logger.warning("No actions returned from LLM, falling back to simple add mode")
                return await self._simple_add_async(messages, user_id, agent_id, run_id, metadata, filters, scope, memory_type, prompt)
            return {"results": []}

        for action in actions:
            action_text = action.get("text", "") or action.get("memory", "")
            event_type = action.get("event", "NONE")
            action_id = action.get("id", "")
            
            # Skip actions with empty text UNLESS it's a NONE event (duplicates may have empty text)
            if not action_text and event_type != "NONE":
                logger.warning(f"Skipping action with empty text: {action}")
                continue
            
            logger.debug(f"Processing action: {event_type} - '{action_text[:50] if action_text else 'NONE'}...' (id: {action_id})")
            
            try:
                if event_type == "ADD":
                    # Add new memory
                    memory_id = await self._create_memory_async(
                        content=action_text,
                        user_id=user_id,
                        agent_id=agent_id,
                        run_id=run_id,
                        metadata=metadata,
                        filters=filters,
                        existing_embeddings=fact_embeddings,
                        scope=scope,
                        memory_type=memory_type,
                    )
                    results.append({
                        "id": memory_id,
                        "memory": action_text,
                        "event": event_type
                    })
                    action_counts["ADD"] += 1
                    
                elif event_type == "UPDATE":
                    # Use ID mapping to get the real memory ID (Snowflake ID - integer)
                    real_memory_id = temp_uuid_mapping.get(str(action_id))
                    if real_memory_id:
                        await self._update_memory_async(
                            memory_id=real_memory_id,
                            content=action_text,
                            user_id=user_id,
                            agent_id=agent_id,
                            existing_embeddings=fact_embeddings,
                            metadata=metadata
                        )
                        results.append({
                            "id": real_memory_id,
                            "memory": action_text,
                            "event": event_type,
                            "previous_memory": action.get("old_memory")
                        })
                        action_counts["UPDATE"] += 1
                    else:
                        logger.warning(f"Could not find real memory ID for action ID: {action_id}")
                        
                elif event_type == "DELETE":
                    # Use ID mapping to get the real memory ID (Snowflake ID - integer)
                    real_memory_id = temp_uuid_mapping.get(str(action_id))
                    if real_memory_id:
                        await self.delete(real_memory_id, user_id, agent_id)
                        results.append({
                            "id": real_memory_id,
                            "memory": action_text,
                            "event": event_type
                        })
                        action_counts["DELETE"] += 1
                    else:
                        logger.warning(f"Could not find real memory ID for action ID: {action_id}")
                        
                elif event_type == "NONE":
                    logger.debug("No action needed for memory (duplicate detected)")
                    action_counts["NONE"] += 1
                    
            except Exception as e:
                logger.error(f"Error executing memory action {event_type}: {e}")
        
        # Log audit event for intelligent add operation
        await self.audit.log_event_async("memory.intelligent_add", {
            "user_id": user_id,
            "agent_id": agent_id,
            "facts_count": len(facts),
            "action_counts": action_counts,
            "results_count": len(results)
        }, user_id=user_id, agent_id=agent_id)
        
        # Add to graph store and get relations (only if graph store is enabled)
        graph_result = None
        if self.enable_graph:
            graph_result = await self._add_to_graph_async(messages, filters, user_id, agent_id, run_id)
        
        # API format: {"results": [...]}
        if results:
            result: Dict[str, Any] = {"results": results}
            if graph_result:
                result["relations"] = graph_result
            return result
        # If we processed actions but they were all NONE (duplicates detected), return empty results
        elif action_counts.get("NONE", 0) > 0:
            logger.info(f"All actions were NONE (duplicates detected), returning empty results")
            result: Dict[str, Any] = {"results": []}
            if graph_result:
                result["relations"] = graph_result
            return result
        # If we had actions but no results (all failed), check fallback setting
        else:
            logger.warning("Actions were processed but no results were created")
            if fallback_to_simple:
                logger.warning("Falling back to simple add mode")
                return await self._simple_add_async(messages, user_id, agent_id, run_id, metadata, filters, scope, memory_type, prompt)
            return {"results": []}

    async def _add_to_graph_async(
        self,
        messages,
        filters: Optional[Dict[str, Any]] = None,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
    ) -> Optional[Dict[str, Any]]:
        """
        Add messages to graph store and return relations asynchronously.
        
        Returns:
            dict with added_entities and deleted_entities, or None if graph store is disabled
        """
        if not self.enable_graph:
            return None

        if self._is_llm_disabled():
            logger.info("LLM is disabled; skipping graph extraction.")
            return None
        
        # Extract content from messages for graph processing
        if isinstance(messages, str):
            data = messages
        elif isinstance(messages, dict):
            data = messages.get("content", "")
        elif isinstance(messages, list):
            data = "\n".join([
                msg.get("content", "") 
                for msg in messages 
                if isinstance(msg, dict) and msg.get("content") and msg.get("role") != "system"
            ])
        else:
            data = ""
        
        if not data:
            return None
        
        graph_filters = {**(filters or {}), "user_id": user_id, "agent_id": agent_id, "run_id": run_id}
        if graph_filters.get("user_id") is None:
            graph_filters["user_id"] = "user"
        
        return self.graph_store.add(data, graph_filters)
    
    async def _create_memory_async(
        self,
        content: str,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
        filters: Optional[Dict[str, Any]] = None,
        existing_embeddings: Optional[Dict[str, Any]] = None,
        scope: Optional[str] = None,
        memory_type: Optional[str] = None,
    ) -> int:
        """Create a memory asynchronously with optional embeddings."""
        # Validate content is not empty
        if not content or not content.strip():
            raise ValueError(f"Cannot create memory with empty content: '{content}'")
        
        # Select embedding service based on metadata (for sub-store routing)
        embedding_service = self._get_embedding_service(metadata)

        # Generate or use existing embedding
        if existing_embeddings and content in existing_embeddings:
            embedding = existing_embeddings[content]
        else:
            embedding = await asyncio.to_thread(embedding_service.embed, content, memory_action="add")
        
        # Disabled LLM-based importance evaluation to save tokens
        # Process metadata
        # enhanced_metadata = await self.intelligence.process_metadata_async(content, metadata)
        enhanced_metadata = metadata  # Use original metadata without LLM evaluation
        
        # Extract category from metadata; prefer explicit memory_type param
        category = ""
        if enhanced_metadata and isinstance(enhanced_metadata, dict):
            category = enhanced_metadata.get("category", "")
            enhanced_metadata = {k: v for k, v in enhanced_metadata.items() if k != "category"}
        if memory_type:
            category = memory_type
        if scope:
            if enhanced_metadata is None:
                enhanced_metadata = {"scope": scope}
            elif isinstance(enhanced_metadata, dict):
                enhanced_metadata = {**enhanced_metadata, "scope": scope}
            else:
                enhanced_metadata = {"scope": scope}
        
        # Generate content hash
        content_hash = hashlib.md5(content.encode('utf-8')).hexdigest()
        
        # Use self.agent_id as fallback if agent_id is not provided
        agent_id = agent_id or self.agent_id
        
        # Create memory data
        memory_data = {
            "content": content,
            "embedding": embedding,
            "user_id": user_id,
            "agent_id": agent_id,
            "run_id": run_id,
            "hash": content_hash,
            "category": category,
            "metadata": enhanced_metadata or {},
            "filters": filters or {},
            "created_at": get_current_datetime(),
            "updated_at": get_current_datetime(),
        }
        
        memory_id = await self.storage.add_memory_async(memory_data)
        
        return memory_id
    
    async def _update_memory_async(
        self,
        memory_id: int,
        content: str,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        existing_embeddings: Optional[Dict[str, Any]] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ):
        """Update a memory asynchronously with optional embeddings."""
        # Use self.agent_id as fallback if agent_id is not provided
        agent_id = agent_id or self.agent_id
        
        # Validate content is not empty
        if not content or not content.strip():
            raise ValueError(f"Cannot update memory with empty content: '{content}'")
        
        # Generate or use existing embedding
        if existing_embeddings and content in existing_embeddings:
            embedding = existing_embeddings[content]
        else:
            # If no metadata provided, try to get existing memory's metadata
            if metadata is None:
                existing = await self.storage.get_memory_async(memory_id, user_id, agent_id)
                if existing:
                    metadata = existing.get("metadata", {})

            # Select embedding service based on metadata (for sub-store routing)
            embedding_service = self._get_embedding_service(metadata)

            embedding = await asyncio.to_thread(embedding_service.embed, content, memory_action="update")
        
        # Generate content hash
        content_hash = hashlib.md5(content.encode('utf-8')).hexdigest()
        
        update_data = {
            "content": content,
            "embedding": embedding,
            "hash": content_hash,  # Update hash
            "updated_at": get_current_datetime(),
        }
        
        logger.debug(f"Updating memory {memory_id} with content: '{content[:50]}...'")
        
        await self.storage.update_memory_async(memory_id, update_data, user_id, agent_id)
    
    async def search(
        self,
        query: str,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
        filters: Optional[Dict[str, Any]] = None,
        limit: int = 30,
        threshold: Optional[float] = None,
    ) -> Dict[str, Any]:
        """Search for memories asynchronously.
        
        Returns:
            Dict[str, Any]: A dictionary containing search results with the following structure:
                - "results" (List[Dict]): List of memory search results, where each result contains:
                    - "memory" (str): The memory content
                    - "metadata" (Dict): Metadata associated with the memory
                    - "score" (float): Similarity score for the result
                    - "id" (int, optional): Memory ID
                    - "created_at" (datetime, optional): Creation timestamp
                    - "updated_at" (datetime, optional): Update timestamp
                    - "user_id" (str, optional): User ID
                    - "agent_id" (str, optional): Agent ID
                    - "run_id" (str, optional): Run ID
                - "relations" (List, optional): Graph relations if graph store is enabled
        """
        try:
            if not query or not query.strip():
                return {
                    "results": [],
                    "relations": []
                }

            # Select embedding service based on filters (for sub-store routing)
            embedding_service = self._get_embedding_service(filters)

            # Generate query embedding asynchronously
            query_embedding = await asyncio.to_thread(embedding_service.embed, query, memory_action="search")
            

            # Search in storage asynchronously - pass query text to enable hybrid search
            results = await self.storage.search_memories_async(
                query_embedding=query_embedding,
                user_id=user_id,
                agent_id=agent_id,
                run_id=run_id,
                filters=filters,
                limit=limit,
                query=query  # Pass query text for hybrid search (vector + full-text)
            )
            
            # Process results with intelligence manager (only if enabled to avoid unnecessary calls)
            if self.intelligence.enabled:
                processed_results = await self.intelligence.process_search_results_async(results, query)
            else:
                processed_results = results

            # Intelligent plugin lifecycle management on search
            if self._intelligence_plugin and self._intelligence_plugin.enabled:
                updates, deletes = self._intelligence_plugin.on_search(processed_results)
                for mem_id, upd in updates:
                    try:
                        await self.storage.update_memory_async(mem_id, {**upd}, user_id, agent_id)
                    except Exception:
                        continue
                # The plugin's "deletes" are memories that should be
                # forgotten by marking, not physically removed from storage.
                for mem_id in deletes:
                    try:
                        await self.storage.update_memory_async(
                            mem_id,
                            _forget_marker_updates(),
                            user_id,
                            agent_id,
                        )
                    except Exception:
                        continue
            
            # Transform results to match benchmark expected format
            # Benchmark expects: {"results": [{"memory": ..., "metadata": {...}, "score": ...}], "relations": [...]}
            transformed_results = []
            for result in processed_results:
                score = result.get("score", 0.0)
                
                # Get quality score for threshold filtering
                # Quality score represents absolute similarity quality (0-1 range)
                # It's calculated from weighted average of all search paths' similarity scores
                metadata = result.get("metadata", {})
                quality_score = metadata.get("_quality_score")
                
                # If quality_score is not available (e.g., from older data or non-hybrid search),
                # fall back to using the ranking score
                if quality_score is None:
                    quality_score = score
                
                # Apply threshold filtering using quality score
                # Only include results if threshold is None or quality_score >= threshold
                if threshold is not None and quality_score < threshold:
                    continue
                
                transformed_result = {
                    "memory": result.get("memory", ""), 
                    "metadata": metadata,  # Keep metadata as-is from storage (includes debug info like _quality_score)
                    "score": score,
                }
                # Preserve other fields if needed
                for key in ["id", "created_at", "updated_at", "user_id", "agent_id", "run_id"]:
                    if key in result:
                        transformed_result[key] = result[key]
                transformed_results.append(transformed_result)
            
            # Log audit event
            await self.audit.log_event_async("memory.search", {
                "query": query,
                "user_id": user_id,
                "agent_id": agent_id,
                "results_count": len(transformed_results)
            }, user_id=user_id, agent_id=agent_id)
            
            # Capture telemetry
            self.telemetry.capture_event("memory.search", {
                "user_id": user_id,
                "agent_id": agent_id,
                "results_count": len(transformed_results),
                "threshold": threshold
            })

            # Search in graph store
            if self.enable_graph and not self._is_llm_disabled():
                filters = {**(filters or {}), "user_id": user_id, "agent_id": agent_id, "run_id": run_id}
                graph_results = await asyncio.to_thread(self.graph_store.search, query, filters, limit)
                return {"results": transformed_results, "relations": graph_results}

            # Return in benchmark expected format
            return {"results": transformed_results}
            
        except Exception as e:
            logger.error(f"Failed to search memories: {e}")
            self.telemetry.capture_event("memory.search.error", {"error": str(e)})
            raise
    
    async def get(
        self,
        memory_id: int,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
    ) -> Optional[Dict[str, Any]]:
        """Get a specific memory by ID asynchronously.
        
        Returns:
            Optional[Dict[str, Any]]: A dictionary containing the memory data if found, None otherwise.
                The dictionary contains the following fields:
                    - "id" (int): Memory ID
                    - "content" (str): The memory content
                    - "user_id" (str, optional): User ID associated with the memory
                    - "agent_id" (str, optional): Agent ID associated with the memory
                    - "run_id" (str, optional): Run ID associated with the memory
                    - "metadata" (Dict): Metadata dictionary associated with the memory
                    - "created_at" (datetime, optional): Creation timestamp
                    - "updated_at" (datetime, optional): Update timestamp
                Returns None if the memory is not found or access is denied.
        """
        try:
            result = await self.storage.get_memory_async(memory_id, user_id, agent_id)
            
            if result:
                if self._intelligence_plugin and self._intelligence_plugin.enabled:
                    updates, delete_flag = self._intelligence_plugin.on_get(result)
                    try:
                        if delete_flag:
                            if updates is None:
                                updates = {}
                            updates.update(_forget_marker_updates())
                        if updates:
                            await self.storage.update_memory_async(memory_id, {**updates}, user_id, agent_id)
                    except Exception:
                        pass
                await self.audit.log_event_async("memory.get", {
                    "memory_id": memory_id,
                    "user_id": user_id,
                    "agent_id": agent_id
                }, user_id=user_id, agent_id=agent_id)
            
            return result
            
        except Exception as e:
            logger.error(f"Failed to get memory {memory_id}: {e}")
            raise
    
    async def update(
        self,
        memory_id: int,
        content: str,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Update an existing memory asynchronously.
        
        Returns:
            Dict[str, Any]: A dictionary containing the updated memory data if successful, None if memory not found or access denied.
                The dictionary contains the following fields:
                    - "id" (int): Memory ID
                    - "content" (str): The updated memory content (stored as "data" in payload)
                    - "user_id" (str, optional): User ID associated with the memory
                    - "agent_id" (str, optional): Agent ID associated with the memory
                    - "run_id" (str, optional): Run ID associated with the memory
                    - "metadata" (Dict): Metadata dictionary associated with the memory
                    - "created_at" (str, optional): Creation timestamp in ISO format
                    - "updated_at" (str): Update timestamp in ISO format
                    - "hash" (str): Content hash for deduplication
                    - "category" (str, optional): Category of the memory
                Returns None if the memory is not found or access is denied.
        """
        try:
            # Validate content is not empty
            if not content or not content.strip():
                raise ValueError(f"Cannot update memory with empty content: '{content}'")

            # If no metadata provided, try to get existing memory's metadata
            if metadata is None:
                existing = await self.storage.get_memory_async(memory_id, user_id, agent_id)
                if existing:
                    metadata = existing.get("metadata", {})

            # Select embedding service based on metadata (for sub-store routing)
            embedding_service = self._get_embedding_service(metadata)

            # Generate new embedding asynchronously
            embedding = await asyncio.to_thread(embedding_service.embed, content, memory_action="update")
            
            # Process metadata with intelligence manager (if enabled)
            # Disabled LLM-based importance evaluation to save tokens (consistent with add method)
            # enhanced_metadata = await self.intelligence.process_metadata_async(content, metadata)
            enhanced_metadata = metadata  # Use original metadata without LLM evaluation
            
            # Intelligent plugin annotations
            extra_fields = {}
            if self._intelligence_plugin and self._intelligence_plugin.enabled:
                # Get existing memory for context
                existing_memory = await self.get(memory_id, user_id=user_id)
                if existing_memory:
                    # Plugin can process update event
                    extra_fields = self._intelligence_plugin.on_add(content=content, metadata=enhanced_metadata)
            
            # Generate content hash for deduplication
            content_hash = hashlib.md5(content.encode('utf-8')).hexdigest()
            
            # Extract category from enhanced metadata if present
            category = ""
            if enhanced_metadata and isinstance(enhanced_metadata, dict):
                category = enhanced_metadata.get("category", "")
                # Remove category from metadata to avoid duplication
                enhanced_metadata = {k: v for k, v in enhanced_metadata.items() if k != "category"}
            
            # Merge extra fields from intelligence plugin
            if extra_fields and isinstance(extra_fields, dict):
                enhanced_metadata = {**(enhanced_metadata or {}), **extra_fields}

            # Update in storage asynchronously
            update_data = {
                "content": content,
                "embedding": embedding,
                "metadata": enhanced_metadata,
                "hash": content_hash,  # Update hash
                "category": category,
                "updated_at": get_current_datetime(),
            }
            
            result = await self.storage.update_memory_async(memory_id, update_data, user_id, agent_id)
            
            # Log audit event
            await self.audit.log_event_async("memory.update", {
                "memory_id": memory_id,
                "user_id": user_id,
                "agent_id": agent_id
            }, user_id=user_id, agent_id=agent_id)
            
            return result
            
        except Exception as e:
            logger.error(f"Failed to update memory {memory_id}: {e}")
            raise
    
    async def delete(
        self,
        memory_id: int,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
    ) -> bool:
        """Delete a memory asynchronously."""
        try:
            result = await self.storage.delete_memory_async(memory_id, user_id, agent_id)
            
            if result:
                await self.audit.log_event_async("memory.delete", {
                    "memory_id": memory_id,
                    "user_id": user_id,
                    "agent_id": agent_id
                }, user_id=user_id, agent_id=agent_id)
            
            return result
            
        except Exception as e:
            logger.error(f"Failed to delete memory {memory_id}: {e}")
            raise
    
    async def get_all(
        self,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
        limit: int = 100,
        offset: int = 0,
        filters: Optional[Dict[str, Any]] = None,
        sort_by: Optional[str] = None,
        order: str = "desc",
    ) -> Dict[str, List[Dict[str, Any]]]:
        """Get all memories with optional filtering and sorting asynchronously.
        
        Args:
            user_id: Optional user ID filter
            agent_id: Optional agent ID filter
            run_id: Optional run ID filter
            limit: Maximum number of results to return (default: 100)
            offset: Number of results to skip (default: 0)
            filters: Optional additional filters dictionary
            sort_by: Optional field to sort results by. Options: "created_at" (creation time),
                     "updated_at" (update time), "id" (memory ID). If None, results are returned
                     in their original order (typically by ID).
            order: Sort order. "desc" for descending (default), "asc" for ascending
        
        Returns:
            Dict[str, List[Dict[str, Any]]]: A dictionary containing all memories with the following structure:
                - "results" (List[Dict]): List of memory dictionaries, where each memory contains:
                    - "id" (int): Memory ID
                    - "content" (str): The memory content
                    - "user_id" (str, optional): User ID associated with the memory
                    - "agent_id" (str, optional): Agent ID associated with the memory
                    - "run_id" (str, optional): Run ID associated with the memory
                    - "metadata" (Dict): Metadata dictionary associated with the memory
                    - "created_at" (datetime or str, optional): Creation timestamp
                    - "updated_at" (datetime or str, optional): Update timestamp
                - "relations" (List[Dict], optional): Graph relations if graph store is enabled
        """
        try:
            results = await self.storage.get_all_memories_async(
                user_id, agent_id, run_id, limit, offset,
                sort_by=sort_by, order=order, filters=filters
            )
            
            await self.audit.log_event_async("memory.get_all", {
                "user_id": user_id,
                "agent_id": agent_id,
                "run_id": run_id,
                "limit": limit,
                "offset": offset,
                "results_count": len(results)
            }, user_id=user_id, agent_id=agent_id)

            # get from graph store
            if self.enable_graph:
                filters = {**(filters or {}), "user_id": user_id, "agent_id": agent_id, "run_id": run_id}
                graph_results = await asyncio.to_thread(self.graph_store.get_all, filters, limit + offset)
                return {"results": results, "relations": graph_results}

            return {"results": results}
            
        except Exception as e:
            logger.error(f"Failed to get all memories: {e}")
            raise
    
    async def delete_all(
        self,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None,
    ) -> bool:
        """Delete all memories for given identifiers asynchronously."""
        try:
            result = await self.storage.clear_memories_async(user_id, agent_id, run_id)
            
            if result:
                await self.audit.log_event_async("memory.delete_all", {
                    "user_id": user_id,
                    "agent_id": agent_id,
                    "run_id": run_id
                }, user_id=user_id, agent_id=agent_id)
                
                self.telemetry.capture_event("memory.delete_all", {
                    "user_id": user_id,
                    "agent_id": agent_id,
                    "run_id": run_id
                })

            if self.enable_graph:
                filters = {"user_id": user_id, "agent_id": agent_id, "run_id": run_id}
                await asyncio.to_thread(self.graph_store.delete_all, filters)
            
            return result
            
        except Exception as e:
            logger.error(f"Failed to delete all memories: {e}")
            raise

    async def reset(self):
        """
        Reset the memory store asynchronously by:
            Deletes the vector store collection
            Resets the database
            Recreates the vector store with a new client
        """
        logger.warning("Resetting all memories")
        
        try:
            # Reset vector store asynchronously
            if hasattr(self.storage.vector_store, "reset"):
                await asyncio.to_thread(self.storage.vector_store.reset)
            else:
                logger.warning("Vector store does not support reset. Skipping.")
                await asyncio.to_thread(self.storage.vector_store.delete_col)
                # Recreate vector store
                from ..storage.factory import VectorStoreFactory
                vector_store_config = self._get_component_config('vector_store')
                self.storage.vector_store = VectorStoreFactory.create(self.storage_type, vector_store_config)
                # Update storage adapter
                self.storage = StorageAdapter(self.storage.vector_store, self.embedding)
            
            # Reset graph store if enabled
            if self.enable_graph and hasattr(self.graph_store, "reset"):
                await asyncio.to_thread(self.graph_store.reset)
            
            # Log telemetry event
            self.telemetry.capture_event("memory.reset", {"sync_type": "async"})
            
            logger.info("Memory store reset completed successfully")
            
        except Exception as e:
            logger.error(f"Failed to reset memory store: {e}")
            raise

    def _init_sub_stores(self):
        """Initialize multiple sub stores configuration"""
        if self.sub_stores_config:
            logger.info(f"Sub stores enabled: {len(self.sub_stores_config)} stores")

        sub_stores_list = self.config.get('sub_stores', [])

        if not sub_stores_list:
            logger.info("No sub stores configured")
            return

        # Sub store feature only supports OceanBase storage
        if self.storage_type.lower() != 'oceanbase':
            logger.warning(f"Sub store feature only supports OceanBase storage, current storage: {self.storage_type}")
            logger.warning("Sub stores configuration will be ignored")
            return

        # Get main table information
        main_collection_name = self.config.get('vector_store', {}).get('config', {}).get('collection_name', 'memories')
        main_embedding_dims = self.config.get('vector_store', {}).get('config', {}).get('embedding_model_dims', 1536)

        # Iterate through configs and initialize each sub store
        for index, sub_config in enumerate(sub_stores_list):
            try:
                self._init_single_sub_store(index, sub_config, main_collection_name, main_embedding_dims)
            except Exception as e:
                logger.error(f"Failed to initialize sub store {index}: {e}")
                continue

    def _init_single_sub_store(
        self,
        index: int,
        sub_config: Dict,
        main_collection_name: str,
        main_embedding_dims: int
    ):
        """Initialize a single sub store"""

        # 1. Determine sub store name (default: {main_table_name}_sub_{index})
        sub_store_name = sub_config.get(
            'collection_name',
            f"{main_collection_name}_sub_{index}"
        )

        # 2. Get routing rules (required)
        routing_filter = sub_config.get('routing_filter')
        if not routing_filter:
            logger.warning(f"Sub store {index} has no routing_filter, skipping")
            return

        # 3. Determine vector dimension (default: same as main table)
        embedding_model_dims = sub_config.get('embedding_model_dims', main_embedding_dims)

        # 4. Initialize sub store's embedding service
        sub_embedding_config = sub_config.get('embedding', {})

        if sub_embedding_config:
            # Has independent embedding configuration
            sub_embedding_provider = sub_embedding_config.get('provider', self.embedding_provider)
            sub_embedding_params = sub_embedding_config.get('config', {})

            # Inherit api_key and other configs from main table
            main_embedding_config = self.config.get('embedder', {}).get('config', {})
            for key in ['api_key', 'openai_base_url', 'timeout']:
                if key not in sub_embedding_params and key in main_embedding_config:
                    sub_embedding_params[key] = main_embedding_config[key]

            # Create a config dict with embedding_model_dims for mock embeddings
            sub_vector_config = {'embedding_model_dims': embedding_model_dims}
            sub_embedding = EmbedderFactory.create(
                sub_embedding_provider,
                sub_embedding_params,
                sub_vector_config
            )
            logger.info(f"Created sub embedding service for store {index}: {sub_embedding_provider}")
        else:
            # Reuse main table's embedding service
            sub_embedding = self.embedding
            logger.info(f"Sub store {index} using main embedding service")

        # 5. Create sub store storage instance
        db_config = self.config.get('vector_store', {}).get('config', {}).copy()
        db_config['collection_name'] = sub_store_name
        db_config['embedding_model_dims'] = embedding_model_dims

        sub_vector_store = VectorStoreFactory.create(self.storage_type, db_config)

        # 6. Register sub store in Adapter (with embedding service for migration)
        self.storage.register_sub_store(
            store_name=sub_store_name,
            routing_filter=routing_filter,
            vector_store=sub_vector_store,
            embedding_service=sub_embedding,
        )

        # 7. Save sub store configuration
        self.sub_stores_config.append({
            'name': sub_store_name,
            'routing_filter': routing_filter,
            'embedding_service': sub_embedding,
            'embedding_dims': embedding_model_dims,
        })

        logger.info(f"Registered sub store {index}: {sub_store_name} (dims={embedding_model_dims})")

    def _get_embedding_service(self, filters_or_metadata: Optional[Dict] = None):
        """
        Select appropriate embedding service based on filters or metadata

        Args:
            filters_or_metadata: Query filters (for search) or memory metadata (for add)

        Returns:
            Corresponding embedding service instance
        """
        if not filters_or_metadata or not self.sub_stores_config:
            return self.embedding

        # Iterate through all sub stores to find a match
        if isinstance(self.storage, SubStorageAdapter):
            for sub_config in self.sub_stores_config:
                # Check if sub store is ready
                if not self.storage.is_sub_store_ready(sub_config['name']):
                    continue

                # Check if filters_or_metadata matches routing rules
                routing_filter = sub_config['routing_filter']
                if all(
                    key in filters_or_metadata and filters_or_metadata[key] == value
                    for key, value in routing_filter.items()
                ):
                    logger.debug(f"Using sub embedding for store: {sub_config['name']}")
                    return sub_config['embedding_service']

        logger.debug("Using main embedding service")
        return self.embedding

    async def migrate_to_sub_store(self, sub_store_index: int = 0, delete_source: bool = False) -> int:
        """
        Migrate data to specified sub store

        Args:
            sub_store_index: Sub store index (default 0, i.e., first sub store)
            delete_source: Whether to delete source data

        Returns:
            Number of migrated records
        """
        if not self.sub_stores_config:
            raise ValueError("No sub stores configured.")

        if sub_store_index >= len(self.sub_stores_config):
            raise ValueError(f"Sub store index {sub_store_index} out of range")

        sub_config = self.sub_stores_config[sub_store_index]

        logger.info(f"Starting migration to sub store: {sub_config['name']}")

        # Call adapter's migration method
        if isinstance(self.storage, SubStorageAdapter):
            migrated_count = await asyncio.to_thread(
                self.storage.migrate_to_sub_store,
                store_name=sub_config['name'],
                delete_source=delete_source
            )

            logger.info(f"Migration completed: {migrated_count} records migrated")
            return migrated_count
        else:
            raise ValueError("Storage adapter does not support migration")

    async def migrate_all_sub_stores(self, delete_source: bool = True) -> Dict[str, int]:
        """
        Migrate all sub stores

        Args:
            delete_source: Whether to delete source data

        Returns:
            Dict[str, int]: A dictionary mapping sub store names to the number of migrated records.
                Each key is a sub store name (str), and each value is the count of migrated records (int).
                If migration fails for a sub store, its count will be 0.
        """
        results = {}
        for index, sub_config in enumerate(self.sub_stores_config):
            try:
                count = await self.migrate_to_sub_store(index, delete_source)
                results[sub_config['name']] = count
            except Exception as e:
                logger.error(f"Failed to migrate sub store {index}: {e}")
                results[sub_config['name']] = 0

        return results

    @classmethod
    async def from_config(cls, config: Optional[Dict[str, Any]] = None, **kwargs):
        """
        Create AsyncMemory instance from configuration.

        Deprecated: prefer `create_memory()` or `auto_config()`.
        
        Args:
            config: Configuration dictionary
            **kwargs: Additional parameters
        
        Returns:
            AsyncMemory instance
            
        Example:
            ```python
            memory = await AsyncMemory.from_config({
                "llm": {"provider": "openai", "config": {"api_key": "..."}},
                "embedder": {"provider": "openai", "config": {"api_key": "..."}},
                "vector_store": {"provider": "oceanbase", "config": {...}},
            })
            ```
        """
        warnings.warn(
            "AsyncMemory.from_config is deprecated; prefer create_memory() or auto_config().",
            DeprecationWarning,
            stacklevel=2,
        )
        if config is None:
            # Use auto config from environment
            from ..config_loader import auto_config
            config = auto_config()

        converted_config = _auto_convert_config(config)
        
        return cls(config=converted_config, **kwargs)
