import keyword
from typing import List, Dict, Optional, Any, Tuple
import uuid
from datetime import datetime
from .llm_controller import LLMController
from .retrievers import ChromaRetriever
from .thread_safe_cache import ThreadSafeMemoryCache
import json
import logging
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import os
from abc import ABC, abstractmethod
from transformers import AutoModel, AutoTokenizer
from nltk.tokenize import word_tokenize
import pickle
from pathlib import Path
from litellm import completion
import time
from collections import OrderedDict

logger = logging.getLogger(__name__)

class MemoryNote:
    """A memory note that represents a single unit of information in the memory system.
    
    This class encapsulates all metadata associated with a memory, including:
    - Core content and identifiers
    - Temporal information (creation and access times)
    - Semantic metadata (keywords, context, tags)
    - Relationship data (links to other memories)
    - Usage statistics (retrieval count)
    - Evolution tracking (history of changes)
    """
    
    def __init__(self, 
                 content: str,
                 id: Optional[str] = None,
                 keywords: Optional[List[str]] = None,
                 links: Optional[Dict] = None,
                 retrieval_count: Optional[int] = None,
                 timestamp: Optional[str] = None,
                 last_accessed: Optional[str] = None,
                 context: Optional[str] = None,
                 evolution_history: Optional[List] = None,
                 category: Optional[str] = None,
                 tags: Optional[List[str]] = None):
        """Initialize a new memory note with its associated metadata.
        
        Args:
            content (str): The main text content of the memory
            id (Optional[str]): Unique identifier for the memory. If None, a UUID will be generated
            keywords (Optional[List[str]]): Key terms extracted from the content
            links (Optional[Dict]): References to related memories
            retrieval_count (Optional[int]): Number of times this memory has been accessed
            timestamp (Optional[str]): Creation time in format YYYYMMDDHHMM
            last_accessed (Optional[str]): Last access time in format YYYYMMDDHHMM
            context (Optional[str]): The broader context or domain of the memory
            evolution_history (Optional[List]): Record of how the memory has evolved
            category (Optional[str]): Classification category
            tags (Optional[List[str]]): Additional classification tags
        """
        # Core content and ID
        self.content = content
        self.id = id or str(uuid.uuid4())
        
        # Semantic metadata
        self.keywords = keywords or []
        self.links = links or []
        self.context = context or "General"
        self.category = category or "Uncategorized"
        self.tags = tags or []
        
        # Temporal information
        current_time = datetime.now().strftime("%Y%m%d%H%M")
        self.timestamp = timestamp or current_time
        self.last_accessed = last_accessed or current_time
        
        # Usage and evolution data
        self.retrieval_count = retrieval_count or 0
        self.evolution_history = evolution_history or []

class MemoryCache:
    """LRU cache for MemoryNote objects.

    Provides O(1) access to recently used memories with automatic eviction
    of least recently used items when capacity is reached.
    """

    def __init__(self, max_size: int = 1000):
        """Initialize the cache with specified capacity.

        Args:
            max_size: Maximum number of memories to keep in cache
        """
        self.max_size = max_size
        self.cache = OrderedDict()
        self.hits = 0
        self.misses = 0
        self.evictions = 0

    def get(self, memory_id: str) -> Optional[MemoryNote]:
        """Get memory from cache, moving to end (most recent).

        Args:
            memory_id: ID of memory to retrieve

        Returns:
            MemoryNote if found in cache, None otherwise
        """
        if memory_id in self.cache:
            self.hits += 1
            # Move to end to mark as recently used
            self.cache.move_to_end(memory_id)
            return self.cache[memory_id]
        self.misses += 1
        return None

    def put(self, memory_id: str, note: MemoryNote):
        """Add/update memory in cache with LRU eviction.

        Args:
            memory_id: ID of memory
            note: MemoryNote object to cache
        """
        if memory_id in self.cache:
            # Update existing entry and move to end
            self.cache.move_to_end(memory_id)
        else:
            # Check if at capacity
            if len(self.cache) >= self.max_size:
                # Evict least recently used (first item)
                self.cache.popitem(last=False)
                self.evictions += 1
        self.cache[memory_id] = note

    def remove(self, memory_id: str):
        """Remove memory from cache.

        Args:
            memory_id: ID of memory to remove
        """
        if memory_id in self.cache:
            del self.cache[memory_id]

    def clear(self):
        """Clear entire cache and reset statistics."""
        self.cache.clear()
        self.hits = 0
        self.misses = 0
        self.evictions = 0

    def get_stats(self) -> Dict[str, Any]:
        """Get cache statistics.

        Returns:
            Dict with cache metrics (size, hits, misses, evictions, hit_rate)
        """
        total_requests = self.hits + self.misses
        hit_rate = self.hits / total_requests if total_requests > 0 else 0
        return {
            'size': len(self.cache),
            'max_size': self.max_size,
            'hits': self.hits,
            'misses': self.misses,
            'evictions': self.evictions,
            'hit_rate': hit_rate
        }

class AgenticMemorySystem:
    """Core memory system that manages memory notes and their evolution.
    
    This system provides:
    - Memory creation, retrieval, update, and deletion
    - Content analysis and metadata extraction
    - Memory evolution and relationship management
    - Hybrid search capabilities
    """
    
    def __init__(self,
                 model_name: str = 'all-MiniLM-L6-v2',
                 llm_backend: str = "openai",
                 llm_model: str = "gpt-4o-mini",
                 evo_threshold: int = 100,
                 api_key: Optional[str] = None,
                 sglang_host: str = "http://localhost",
                 sglang_port: int = 30000,
                 storage_path: str = "./chroma_db",
                 cache_size: int = 1000,
                 enable_cache: bool = True):
        """Initialize the memory system.

        Args:
            model_name: Name of the sentence transformer model
            llm_backend: LLM backend to use (openai/ollama/sglang)
            llm_model: Name of the LLM model
            evo_threshold: Number of memories before triggering evolution
            api_key: API key for the LLM service
            sglang_host: Host URL for SGLang server (default: http://localhost)
            sglang_port: Port for SGLang server (default: 30000)
            storage_path: Directory path for persistent ChromaDB storage (default: ./chroma_db)
            cache_size: Maximum number of memories to keep in LRU cache (default: 1000)
            enable_cache: Whether to enable memory caching (default: True)
        """
        # Initialize thread-safe LRU cache instead of self.memories dict
        self.cache = ThreadSafeMemoryCache(max_size=cache_size)
        self.cache_enabled = enable_cache

        self.model_name = model_name
        self.storage_path = storage_path

        # Initialize ChromaDB retriever WITHOUT resetting (preserves existing data)
        self.retriever = ChromaRetriever(
            collection_name="memories",
            model_name=self.model_name,
            persist_directory=self.storage_path
        )

        # Initialize LLM controller
        self.llm_controller = LLMController(llm_backend, llm_model, api_key, sglang_host, sglang_port)
        self.evo_cnt = 0
        self.evo_threshold = evo_threshold

        # Log initialization info
        existing_count = self.retriever.count()
        logger.info(f"AgenticMemorySystem initialized with {existing_count} existing memories in ChromaDB")

        # Evolution system prompt
        self._evolution_system_prompt = '''
                                You are an AI memory evolution agent responsible for managing and evolving a knowledge base.
                                Analyze the the new memory note according to keywords and context, also with their several nearest neighbors memory.
                                Make decisions about its evolution.

                                The new memory context:
                                {context}
                                content: {content}
                                keywords: {keywords}

                                The nearest neighbors memories (each line starts with memory_id):
                                {nearest_neighbors_memories}

                                Based on this information, determine:
                                1. Should this memory be evolved? Consider its relationships with other memories.
                                2. What specific actions should be taken (strengthen, update_neighbor)?
                                   2.1 If choose to strengthen the connection, which memory should it be connected to? Use the memory_id from the neighbors above. Can you give the updated tags of this memory?
                                   2.2 If choose to update_neighbor, you can update the context and tags of these memories based on the understanding of these memories. If the context and the tags are not updated, the new context and tags should be the same as the original ones. Generate the new context and tags in the sequential order of the input neighbors.
                                Tags should be determined by the content of these characteristic of these memories, which can be used to retrieve them later and categorize them.
                                Note that the length of new_tags_neighborhood must equal the number of input neighbors, and the length of new_context_neighborhood must equal the number of input neighbors.
                                The number of neighbors is {neighbor_number}.
                                Return your decision in JSON format with the following structure:
                                {{
                                    "should_evolve": True or False,
                                    "actions": ["strengthen", "update_neighbor"],
                                    "suggested_connections": ["memory_id_1", "memory_id_2", ...],
                                    "tags_to_update": ["tag_1",..."tag_n"],
                                    "new_context_neighborhood": ["new context",...,"new context"],
                                    "new_tags_neighborhood": [["tag_1",...,"tag_n"],...["tag_1",...,"tag_n"]],
                                }}
                                '''

    def _metadata_to_memory_note(self, metadata: Dict) -> MemoryNote:
        """Convert ChromaDB metadata dict to MemoryNote object.

        Args:
            metadata: Deserialized metadata dict from ChromaDB

        Returns:
            MemoryNote object reconstructed from metadata
        """
        return MemoryNote(
            content=metadata.get('content', ''),
            id=metadata.get('id'),
            keywords=metadata.get('keywords', []),
            links=metadata.get('links', []),
            retrieval_count=metadata.get('retrieval_count', 0),
            timestamp=metadata.get('timestamp'),
            last_accessed=metadata.get('last_accessed'),
            context=metadata.get('context', 'General'),
            evolution_history=metadata.get('evolution_history', []),
            category=metadata.get('category', 'Uncategorized'),
            tags=metadata.get('tags', [])
        )

    def _memory_note_to_metadata(self, note: MemoryNote) -> Dict:
        """Convert MemoryNote object to ChromaDB metadata dict.

        Args:
            note: MemoryNote object to serialize

        Returns:
            Dict suitable for ChromaDB storage
        """
        return {
            "id": note.id,
            "content": note.content,
            "keywords": note.keywords,
            "links": note.links,
            "retrieval_count": note.retrieval_count,
            "timestamp": note.timestamp,
            "last_accessed": note.last_accessed,
            "context": note.context,
            "evolution_history": note.evolution_history,
            "category": note.category,
            "tags": note.tags
        }

    def analyze_content(self, content: str) -> Dict:            
        """Analyze content using LLM to extract semantic metadata.
        
        Uses a language model to understand the content and extract:
        - Keywords: Important terms and concepts
        - Context: Overall domain or theme
        - Tags: Classification categories
        
        Args:
            content (str): The text content to analyze
            
        Returns:
            Dict: Contains extracted metadata with keys:
                - keywords: List[str]
                - context: str
                - tags: List[str]
        """
        prompt = """Generate a structured analysis of the following content by:
            1. Identifying the most salient keywords (focus on nouns, verbs, and key concepts)
            2. Extracting core themes and contextual elements
            3. Creating relevant categorical tags

            Format the response as a JSON object:
            {
                "keywords": [
                    // several specific, distinct keywords that capture key concepts and terminology
                    // Order from most to least important
                    // Don't include keywords that are the name of the speaker or time
                    // At least three keywords, but don't be too redundant.
                ],
                "context": 
                    // one sentence summarizing:
                    // - Main topic/domain
                    // - Key arguments/points
                    // - Intended audience/purpose
                ,
                "tags": [
                    // several broad categories/themes for classification
                    // Include domain, format, and type tags
                    // At least three tags, but don't be too redundant.
                ]
            }

            Content for analysis:
            """ + content
        try:
            response = self.llm_controller.llm.get_completion(prompt, response_format={"type": "json_schema", "json_schema": {
                        "name": "response",
                        "schema": {
                            "type": "object",
                            "properties": {
                                "keywords": {
                                    "type": "array",
                                    "items": {
                                        "type": "string"
                                    }
                                },
                                "context": {
                                    "type": "string",
                                },
                                "tags": {
                                    "type": "array",
                                    "items": {
                                        "type": "string"
                                    }
                                }
                            }
                        }
                    }})
            return json.loads(response)
        except Exception as e:
            print(f"Error analyzing content: {e}")
            return {"keywords": [], "context": "General", "tags": []}

    def add_note(self, content: str, time: str = None, **kwargs) -> str:
        """Add a new memory note"""
        # Create MemoryNote without llm_controller
        if time is not None:
            kwargs['timestamp'] = time
        note = MemoryNote(content=content, **kwargs)
        
        # 🔧 LLM Analysis Enhancement: Auto-generate attributes using LLM if they are empty or default values
        needs_analysis = (
            not note.keywords or  # keywords is empty list
            note.context == "General" or  # context is default value
            not note.tags  # tags is empty list
        )
        
        if needs_analysis:
            # try:
            analysis = self.analyze_content(content)
            
            # Only update attributes that are not provided or have default values
            if not note.keywords:
                note.keywords = analysis.get("keywords", [])
            if note.context == "General":
                note.context = analysis.get("context", "General") 
            if not note.tags:
                note.tags = analysis.get("tags", [])
                    
            # except Exception as e:
            #     print(f"Warning: LLM analysis failed, using default values: {e}")
        
        # Process memory evolution
        evo_label, note = self.process_memory(note)

        # Cache the new memory (write-through caching)
        if self.cache_enabled:
            self.cache.put(note.id, note)

        # Add to ChromaDB with complete metadata (persistent storage)
        metadata = {
            "id": note.id,
            "content": note.content,
            "keywords": note.keywords,
            "links": note.links,
            "retrieval_count": note.retrieval_count,
            "timestamp": note.timestamp,
            "last_accessed": note.last_accessed,
            "context": note.context,
            "evolution_history": note.evolution_history,
            "category": note.category,
            "tags": note.tags
        }
        self.retriever.add_document(note.content, metadata, note.id)

        # Track evolution count (could be used for metrics/logging)
        if evo_label == True:
            self.evo_cnt += 1
            # Note: consolidate_memories() removed - all changes now sync immediately

        return note.id
    
    def find_related_memories(self, query: str, k: int = 5) -> Tuple[str, List[str]]:
        """Find related memories using ChromaDB retrieval

        Returns:
            Tuple[str, List[str]]: (formatted_memory_string, list_of_memory_ids)
        """
        if self.retriever.count() == 0:
            return "", []

        try:
            # Get results from ChromaDB
            results = self.retriever.search(query, k)

            # Convert to list of memories
            memory_str = ""
            memory_ids = []

            if 'ids' in results and results['ids'] and len(results['ids']) > 0 and len(results['ids'][0]) > 0:
                for i, doc_id in enumerate(results['ids'][0]):
                    # Get metadata from ChromaDB results
                    if i < len(results['metadatas'][0]):
                        metadata = results['metadatas'][0][i]
                        # Format memory string with actual memory ID
                        memory_str += f"memory_id:{doc_id}\ttalk start time:{metadata.get('timestamp', '')}\tmemory content: {metadata.get('content', '')}\tmemory context: {metadata.get('context', '')}\tmemory keywords: {str(metadata.get('keywords', []))}\tmemory tags: {str(metadata.get('tags', []))}\n"
                        memory_ids.append(doc_id)

            return memory_str, memory_ids
        except Exception as e:
            logger.error(f"Error in find_related_memories: {str(e)}")
            return "", []

    def find_related_memories_raw(self, query: str, k: int = 5) -> str:
        """Find related memories using ChromaDB retrieval in raw format"""
        if self.retriever.count() == 0:
            return ""
            
        # Get results from ChromaDB
        results = self.retriever.search(query, k)
        
        # Convert to list of memories
        memory_str = ""
        
        if 'ids' in results and results['ids'] and len(results['ids']) > 0:
            for i, doc_id in enumerate(results['ids'][0][:k]):
                if i < len(results['metadatas'][0]):
                    # Get metadata from ChromaDB results
                    metadata = results['metadatas'][0][i]
                    
                    # Add main memory info
                    memory_str += f"talk start time:{metadata.get('timestamp', '')}\tmemory content: {metadata.get('content', '')}\tmemory context: {metadata.get('context', '')}\tmemory keywords: {str(metadata.get('keywords', []))}\tmemory tags: {str(metadata.get('tags', []))}\n"
                    
                    # Add linked memories if available
                    links = metadata.get('links', [])
                    j = 0
                    for link_id in links:
                        neighbor = self.read(link_id)  # Cache-aware lazy loading
                        if neighbor and j < k:
                            memory_str += f"talk start time:{neighbor.timestamp}\tmemory content: {neighbor.content}\tmemory context: {neighbor.context}\tmemory keywords: {str(neighbor.keywords)}\tmemory tags: {str(neighbor.tags)}\n"
                            j += 1
                            
        return memory_str

    def read(self, memory_id: str) -> Optional[MemoryNote]:
        """Retrieve a memory note by its ID (cache-aware lazy loading).

        Implements cache-first pattern:
        1. Check cache for memory (fast O(1) lookup)
        2. If miss, load from ChromaDB (persistent storage)
        3. Cache the loaded memory for future reads

        Args:
            memory_id (str): ID of the memory to retrieve

        Returns:
            MemoryNote if found, None otherwise
        """
        # Try cache first
        if self.cache_enabled:
            cached = self.cache.get(memory_id)
            if cached:
                # Update last_accessed timestamp
                cached.last_accessed = datetime.now().strftime("%Y%m%d%H%M")
                cached.retrieval_count += 1
                return cached

        # Cache miss - lazy load from ChromaDB
        metadata = self.retriever.get_by_id(memory_id)
        if metadata:
            note = self._metadata_to_memory_note(metadata)
            # Update access metadata
            note.last_accessed = datetime.now().strftime("%Y%m%d%H%M")
            note.retrieval_count += 1

            # Cache for future access
            if self.cache_enabled:
                self.cache.put(memory_id, note)

            return note

        return None
    
    def update(self, memory_id: str, **kwargs) -> bool:
        """Update a memory note with write-through caching.

        Implements write-through pattern:
        1. Load memory from cache or ChromaDB
        2. Update fields in MemoryNote object
        3. Sync to ChromaDB immediately (persistent)
        4. Update cache (if enabled)

        Args:
            memory_id: ID of memory to update
            **kwargs: Fields to update

        Returns:
            bool: True if update successful, False if memory not found
        """
        # Load memory (cache-aware)
        note = self.read(memory_id)
        if not note:
            return False

        # Update fields
        for key, value in kwargs.items():
            if hasattr(note, key):
                setattr(note, key, value)

        # Sync to ChromaDB immediately (write-through)
        metadata = self._memory_note_to_metadata(note)
        self.retriever.update_document(memory_id, metadata, note.content)

        # Update cache with new version
        if self.cache_enabled:
            self.cache.put(memory_id, note)

        return True
    
    def delete(self, memory_id: str) -> bool:
        """Delete a memory note by its ID with cache invalidation.

        Deletes from both persistent storage (ChromaDB) and cache.

        Args:
            memory_id (str): ID of the memory to delete

        Returns:
            bool: True if memory was deleted, False if not found
        """
        try:
            # Delete from ChromaDB (persistent storage)
            self.retriever.delete_document(memory_id)

            # Invalidate cache entry
            if self.cache_enabled:
                self.cache.remove(memory_id)

            return True
        except Exception as e:
            logger.error(f"Error deleting memory {memory_id}: {e}")
            return False
    
    def _search_raw(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        """Internal search method that returns raw results from ChromaDB.
        
        This is used internally by the memory evolution system to find
        related memories for potential evolution.
        
        Args:
            query (str): The search query text
            k (int): Maximum number of results to return
            
        Returns:
            List[Dict[str, Any]]: Raw search results from ChromaDB
        """
        results = self.retriever.search(query, k)
        return [{'id': doc_id, 'score': score} 
                for doc_id, score in zip(results['ids'][0], results['distances'][0])]
                
    def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        """Search for memories using ChromaDB vector search with cache-aware loading."""
        # Get results from ChromaDB
        search_results = self.retriever.search(query, k)
        memories = []

        # Process ChromaDB results - load via cache-aware read()
        for i, doc_id in enumerate(search_results['ids'][0]):
            memory = self.read(doc_id)  # Cache-aware lazy loading
            if memory:
                memories.append({
                    'id': doc_id,
                    'content': memory.content,
                    'context': memory.context,
                    'keywords': memory.keywords,
                    'tags': memory.tags,
                    'score': search_results['distances'][0][i]
                })

        return memories[:k]
    
    def _search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        """Search for memories using a hybrid retrieval approach.

        This method combines results from both:
        1. ChromaDB vector store (semantic similarity)
        2. Embedding-based retrieval (dense vectors)

        The results are deduplicated and ranked by relevance.

        Args:
            query (str): The search query text
            k (int): Maximum number of results to return

        Returns:
            List[Dict[str, Any]]: List of search results, each containing:
                - id: Memory ID
                - content: Memory content
                - context: Memory context
                - keywords: Memory keywords
                - tags: Memory tags
                - score: Similarity score
        """
        # Get results from ChromaDB
        chroma_results = self.retriever.search(query, k)
        memories = []

        # Process ChromaDB results
        for i, doc_id in enumerate(chroma_results['ids'][0]):
            memory = self.read(doc_id)
            if memory:
                memories.append({
                    'id': doc_id,
                    'content': memory.content,
                    'context': memory.context,
                    'keywords': memory.keywords,
                    'tags': memory.tags,
                    'score': chroma_results['distances'][0][i]
                })

        # Get results from embedding retriever
        embedding_results = self.retriever.search(query, k)

        # Combine results with deduplication
        seen_ids = set(m['id'] for m in memories)
        for result in embedding_results:
            memory_id = result.get('id')
            if memory_id and memory_id not in seen_ids:
                memory = self.read(memory_id)
                if memory:
                    memories.append({
                        'id': memory_id,
                        'content': memory.content,
                        'context': memory.context,
                        'keywords': memory.keywords,
                        'tags': memory.tags,
                        'score': result.get('score', 0.0)
                    })
                    seen_ids.add(memory_id)

        return memories[:k]

    def search_agentic(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        """Search for memories using ChromaDB retrieval with linked neighbors."""
        # No need to check self.memories - ChromaDB is source of truth
        if self.retriever.count() == 0:
            return []
            
        try:
            # Get results from ChromaDB
            results = self.retriever.search(query, k)
            
            # Process results
            memories = []
            seen_ids = set()
            
            # Check if we have valid results
            if ('ids' not in results or not results['ids'] or 
                len(results['ids']) == 0 or len(results['ids'][0]) == 0):
                return []
                
            # Process ChromaDB results
            for i, doc_id in enumerate(results['ids'][0][:k]):
                if doc_id in seen_ids:
                    continue
                    
                if i < len(results['metadatas'][0]):
                    metadata = results['metadatas'][0][i]
                    
                    # Create result dictionary with all metadata fields
                    memory_dict = {
                        'id': doc_id,
                        'content': metadata.get('content', ''),
                        'context': metadata.get('context', ''),
                        'keywords': metadata.get('keywords', []),
                        'tags': metadata.get('tags', []),
                        'timestamp': metadata.get('timestamp', ''),
                        'category': metadata.get('category', 'Uncategorized'),
                        'is_neighbor': False
                    }
                    
                    # Add score if available
                    if 'distances' in results and len(results['distances']) > 0 and i < len(results['distances'][0]):
                        memory_dict['score'] = results['distances'][0][i]
                        
                    memories.append(memory_dict)
                    seen_ids.add(doc_id)
            
            # Add linked memories (neighbors)
            neighbor_count = 0
            for memory in list(memories):  # Use a copy to avoid modification during iteration
                if neighbor_count >= k:
                    break
                    
                # Get links from metadata
                links = memory.get('links', [])
                if not links and 'id' in memory:
                    # Try to get links from memory object (cache-aware)
                    mem_obj = self.read(memory['id'])
                    if mem_obj:
                        links = mem_obj.links

                for link_id in links:
                    if link_id not in seen_ids and neighbor_count < k:
                        neighbor = self.read(link_id)  # Cache-aware lazy loading
                        if neighbor:
                            memories.append({
                                'id': link_id,
                                'content': neighbor.content,
                                'context': neighbor.context,
                                'keywords': neighbor.keywords,
                                'tags': neighbor.tags,
                                'timestamp': neighbor.timestamp,
                                'category': neighbor.category,
                                'is_neighbor': True
                            })
                            seen_ids.add(link_id)
                            neighbor_count += 1
            
            return memories[:k]
        except Exception as e:
            logger.error(f"Error in search_agentic: {str(e)}")
            return []

    def process_memory(self, note: MemoryNote) -> Tuple[bool, MemoryNote]:
        """Process a memory note and determine if it should evolve.

        Args:
            note: The memory note to process

        Returns:
            Tuple[bool, MemoryNote]: (should_evolve, processed_note)
        """
        # For first memory, just return the note without evolution
        if self.retriever.count() == 0:
            return False, note

        try:
            # Get nearest neighbors - now returns actual memory IDs
            neighbors_text, memory_ids = self.find_related_memories(note.content, k=5)
            if not neighbors_text or not memory_ids:
                return False, note

            # Format neighbors for LLM - in this case, neighbors_text is already formatted

            # Query LLM for evolution decision
            prompt = self._evolution_system_prompt.format(
                content=note.content,
                context=note.context,
                keywords=note.keywords,
                nearest_neighbors_memories=neighbors_text,
                neighbor_number=len(memory_ids)
            )
            
            try:
                response = self.llm_controller.llm.get_completion(
                    prompt,
                    response_format={"type": "json_schema", "json_schema": {
                        "name": "response",
                        "schema": {
                            "type": "object",
                            "properties": {
                                "should_evolve": {
                                    "type": "boolean"
                                },
                                "actions": {
                                    "type": "array",
                                    "items": {
                                        "type": "string"
                                    }
                                },
                                "suggested_connections": {
                                    "type": "array",
                                    "items": {
                                        "type": "string"
                                    }
                                },
                                "new_context_neighborhood": {
                                    "type": "array",
                                    "items": {
                                        "type": "string"
                                    }
                                },
                                "tags_to_update": {
                                    "type": "array",
                                    "items": {
                                        "type": "string"
                                    }
                                },
                                "new_tags_neighborhood": {
                                    "type": "array",
                                    "items": {
                                        "type": "array",
                                        "items": {
                                            "type": "string"
                                        }
                                    }
                                }
                            },
                            "required": ["should_evolve", "actions", "suggested_connections", 
                                      "tags_to_update", "new_context_neighborhood", "new_tags_neighborhood"],
                            "additionalProperties": False
                        },
                        "strict": True
                    }}
                )
                
                response_json = json.loads(response)
                should_evolve = response_json["should_evolve"]
                
                if should_evolve:
                    actions = response_json["actions"]
                    for action in actions:
                        if action == "strengthen":
                            suggest_connections = response_json["suggested_connections"]
                            new_tags = response_json["tags_to_update"]
                            note.links.extend(suggest_connections)
                            note.tags = new_tags
                        elif action == "update_neighbor":
                            new_context_neighborhood = response_json["new_context_neighborhood"]
                            new_tags_neighborhood = response_json["new_tags_neighborhood"]

                            # Update each neighbor memory using its actual ID
                            for i in range(min(len(memory_ids), len(new_tags_neighborhood))):
                                memory_id = memory_ids[i]

                                # Load neighbor from ChromaDB/cache (cache-aware)
                                neighbor_memory = self.read(memory_id)
                                if not neighbor_memory:
                                    logger.warning(f"Neighbor memory {memory_id} not found during evolution")
                                    continue

                                # Prepare update kwargs
                                update_kwargs = {}

                                # Update tags
                                if i < len(new_tags_neighborhood):
                                    update_kwargs['tags'] = new_tags_neighborhood[i]

                                # Update context
                                if i < len(new_context_neighborhood):
                                    update_kwargs['context'] = new_context_neighborhood[i]

                                # CRITICAL: Sync to ChromaDB immediately (write-through)
                                if update_kwargs:
                                    success = self.update(memory_id, **update_kwargs)
                                    if not success:
                                        logger.error(f"Failed to update neighbor {memory_id} during evolution")

                return should_evolve, note
                
            except (json.JSONDecodeError, KeyError, Exception) as e:
                logger.error(f"Error in memory evolution: {str(e)}")
                return False, note
                
        except Exception as e:
            # For testing purposes, catch all exceptions and return the original note
            logger.error(f"Error in process_memory: {str(e)}")
            return False, note
