from ast import Str
from typing import List, Dict, Optional, Literal, Any, Union
import json
import re
from datetime import datetime
import uuid
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 requests
import json as json_lib
import time
import torch


def simple_tokenize(text):
    return word_tokenize(text)


class BaseLLMController(ABC):
    @abstractmethod
    def get_completion(self, prompt: str) -> str:
        """Get completion from LLM"""
        pass


class OpenAIController(BaseLLMController):
    def __init__(self, model: str = "gpt-4", api_key: Optional[str] = None):
        try:
            from openai import OpenAI

            self.model = model
            if api_key is None:
                api_key = os.getenv("OPENAI_API_KEY")
            if api_key is None:
                raise ValueError(
                    "OpenAI API key not found. Set OPENAI_API_KEY environment variable."
                )
            self.client = OpenAI(api_key=api_key)
        except ImportError:
            raise ImportError(
                "OpenAI package not found. Install it with: pip install openai"
            )

    def get_completion(
        self, prompt: str, response_format: dict, temperature: float = 0.7
    ) -> str:
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You must respond with a JSON object."},
                {"role": "user", "content": prompt},
            ],
            response_format=response_format,
            temperature=temperature,
            max_tokens=1000,
        )
        return response.choices[0].message.content


class OllamaController(BaseLLMController):
    def __init__(self, model: str = "llama2"):
        from ollama import chat

        self.model = model

    def _generate_empty_value(self, schema_type: str, schema_items: dict = None) -> Any:
        if schema_type == "array":
            return []
        elif schema_type == "string":
            return ""
        elif schema_type == "object":
            return {}
        elif schema_type == "number":
            return 0
        elif schema_type == "boolean":
            return False
        return None

    def _generate_empty_response(self, response_format: dict) -> dict:
        if "json_schema" not in response_format:
            return {}

        schema = response_format["json_schema"]["schema"]
        result = {}

        if "properties" in schema:
            for prop_name, prop_schema in schema["properties"].items():
                result[prop_name] = self._generate_empty_value(
                    prop_schema["type"], prop_schema.get("items")
                )

        return result

    def get_completion(
        self, prompt: str, response_format: dict, temperature: float = 0.7
    ) -> str:
        try:
            response = completion(
                model="ollama_chat/{}".format(self.model),
                messages=[
                    {
                        "role": "system",
                        "content": "You must respond with a JSON object.",
                    },
                    {"role": "user", "content": prompt},
                ],
                response_format=response_format,
            )
            return response.choices[0].message.content
        except Exception as e:
            empty_response = self._generate_empty_response(response_format)
            return json.dumps(empty_response)


class SGLangController(BaseLLMController):
    def __init__(
        self,
        model: str = "llama2",
        sglang_host: str = "http://localhost",
        sglang_port: int = 30000,
    ):
        self.model = model
        self.sglang_host = sglang_host
        self.sglang_port = sglang_port
        self.base_url = f"{sglang_host}:{sglang_port}"

    def _generate_empty_value(self, schema_type: str, schema_items: dict = None) -> Any:
        if schema_type == "array":
            return []
        elif schema_type == "string":
            return ""
        elif schema_type == "object":
            return {}
        elif schema_type == "number" or schema_type == "integer":
            return 0
        elif schema_type == "boolean":
            return False
        return None

    def _generate_empty_response(self, response_format: dict) -> dict:
        if "json_schema" not in response_format:
            return {}

        schema = response_format["json_schema"]["schema"]
        result = {}

        if "properties" in schema:
            for prop_name, prop_schema in schema["properties"].items():
                result[prop_name] = self._generate_empty_value(
                    prop_schema["type"], prop_schema.get("items")
                )

        return result

    def get_completion(
        self, prompt: str, response_format: dict, temperature: float = 0.7
    ) -> str:
        try:
            # Extract JSON schema from response_format and convert to string format
            json_schema = response_format.get("json_schema", {}).get("schema", {})
            json_schema_str = json.dumps(json_schema)

            # Prepare SGLang request with correct format
            payload = {
                "text": prompt,
                "sampling_params": {
                    "temperature": temperature,
                    "max_new_tokens": 1000,
                    "json_schema": json_schema_str,  # SGLang expects JSON schema as string
                },
            }

            # Make request to SGLang server
            response = requests.post(
                f"{self.base_url}/generate",
                headers={"Content-Type": "application/json"},
                json=payload,
                timeout=60,
            )

            if response.status_code == 200:
                result = response.json()
                # SGLang returns the generated text in 'text' field
                generated_text = result.get("text", "")
                return generated_text
            else:
                print(
                    f"SGLang server returned status {response.status_code}: {response.text}"
                )
                raise Exception(f"SGLang server error: {response.status_code}")

        except Exception as e:
            print(f"SGLang completion error: {e}")
            empty_response = self._generate_empty_response(response_format)
            return json.dumps(empty_response)


class LiteLLMController(BaseLLMController):
    """LiteLLM controller for universal LLM access including Ollama and SGLang"""

    def __init__(
        self, model: str, api_base: Optional[str] = None, api_key: Optional[str] = None
    ):
        self.model = model
        self.api_base = api_base
        self.api_key = api_key or "EMPTY"

    def _generate_empty_value(self, schema_type: str, schema_items: dict = None) -> Any:
        if schema_type == "array":
            return []
        elif schema_type == "string":
            return ""
        elif schema_type == "object":
            return {}
        elif schema_type == "number":
            return 0
        elif schema_type == "boolean":
            return False
        return None

    def _generate_empty_response(self, response_format: dict) -> dict:
        if "json_schema" not in response_format:
            return {}

        schema = response_format["json_schema"]["schema"]
        result = {}

        if "properties" in schema:
            for prop_name, prop_schema in schema["properties"].items():
                result[prop_name] = self._generate_empty_value(
                    prop_schema["type"], prop_schema.get("items")
                )

        return result

    def get_completion(
        self, prompt: str, response_format: dict, temperature: float = 0.7
    ) -> str:
        try:
            # Prepare completion arguments
            completion_args = {
                "model": self.model,
                "messages": [
                    {
                        "role": "system",
                        "content": "You must respond with a JSON object.",
                    },
                    {"role": "user", "content": prompt},
                ],
                "response_format": response_format,
                "temperature": temperature,
            }

            # Add API base and key if provided
            if self.api_base:
                completion_args["api_base"] = self.api_base
            if self.api_key:
                completion_args["api_key"] = self.api_key

            response = completion(**completion_args)
            return response.choices[0].message.content

        except Exception as e:
            print(f"LiteLLM completion error: {e}")
            empty_response = self._generate_empty_response(response_format)
            return json.dumps(empty_response)


class LLMController:
    """LLM-based controller for memory metadata generation"""

    def __init__(
        self,
        backend: Literal["openai", "ollama", "sglang"] = "sglang",
        model: str = "gpt-4",
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
        sglang_host: str = "http://localhost",
        sglang_port: int = 30000,
    ):
        if backend == "openai":
            self.llm = OpenAIController(model, api_key)
        elif backend == "ollama":
            # Use LiteLLM to control Ollama with JSON output
            ollama_model = (
                f"ollama/{model}" if not model.startswith("ollama/") else model
            )
            self.llm = LiteLLMController(
                model=ollama_model, api_base="http://localhost:11434", api_key="EMPTY"
            )
        elif backend == "sglang":
            # Direct SGLang API calls (better performance, no proxy)
            self.llm = SGLangController(model, sglang_host, sglang_port)
        else:
            raise ValueError("Backend must be 'openai', 'ollama', or 'sglang'")


class MemoryNote:
    """Basic memory unit with metadata"""

    def __init__(
        self,
        content: str,
        id: Optional[str] = None,
        keywords: Optional[List[str]] = None,
        links: Optional[Dict] = None,
        importance_score: Optional[float] = 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,
        llm_controller: Optional[LLMController] = None,
    ):
        self.content = content

        # Generate metadata using LLM if not provided and controller is available
        if llm_controller and any(
            param is None for param in [keywords, context, category, tags]
        ):
            analysis = self.analyze_content(content, llm_controller)
            print("analysis", analysis)
            keywords = keywords or analysis["keywords"]
            context = context or analysis["context"]
            tags = tags or analysis["tags"]

        # Set default values for optional parameters
        self.id = id or str(uuid.uuid4())
        self.keywords = keywords or []
        self.links = links or []
        self.importance_score = importance_score or 1.0
        self.retrieval_count = retrieval_count or 0
        current_time = datetime.now().strftime("%Y%m%d%H%M")
        self.timestamp = timestamp or current_time
        self.last_accessed = last_accessed or current_time

        # Handle context that can be either string or list
        self.context = context or "General"
        if isinstance(self.context, list):
            self.context = " ".join(self.context)  # Convert list to string by joining

        self.evolution_history = evolution_history or []
        self.category = category or "Uncategorized"
        self.tags = tags or []

    @staticmethod
    def analyze_content(content: str, llm_controller: LLMController) -> Dict:
        """Analyze content to extract keywords, context, and other metadata"""
        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 = 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"}},
                            },
                            "required": ["keywords", "context", "tags"],
                            "additionalProperties": False,
                        },
                        "strict": True,
                    },
                },
            )

            # try:
            #     # Clean the response in case there's extra text
            #     response_cleaned = response.strip()
            #     # Try to find JSON content if wrapped in other text
            #     if not response_cleaned.startswith('{'):
            #         start_idx = response_cleaned.find('{')
            #         if start_idx != -1:
            #             response_cleaned = response_cleaned[start_idx:]
            #     if not response_cleaned.endswith('}'):
            #         end_idx = response_cleaned.rfind('}')
            #         if end_idx != -1:
            #             response_cleaned = response_cleaned[:end_idx+1]
            try:
                response = re.sub(
                    r"^```json\s*|\s*```$", "", response, flags=re.MULTILINE
                ).strip()
                analysis = json.loads(response)
            except Exception as e:
                print(f"JSON parsing error in analyze_content: {e}")
                print(f"Raw response: {response}")
                analysis = {"keywords": [], "context": "General", "tags": []}

            return analysis

        except Exception as e:
            print(f"Error analyzing content: {str(e)}")
            print(f"Raw response: {response}")
            return {
                "keywords": [],
                "context": "General",
                "category": "Uncategorized",
                "tags": [],
            }


class HybridRetriever:
    """Hybrid retrieval system combining BM25 and semantic search."""

    def __init__(self, model_name: str = "all-MiniLM-L6-v2", alpha: float = 0.5):
        """Initialize the hybrid retriever.

        Args:
            model_name: Name of the SentenceTransformer model to use
            alpha: Weight for combining BM25 and semantic scores (0 = only BM25, 1 = only semantic)
        """
        self.model = SentenceTransformer(model_name)
        self.alpha = alpha
        self.bm25 = None
        self.corpus = []
        self.embeddings = None
        self.document_ids = {}  # Map document content to its index

    def save(self, retriever_cache_file: str, retriever_cache_embeddings_file: str):
        """Save retriever state to disk"""

        # Save embeddings using numpy
        if self.embeddings is not None:
            np.save(retriever_cache_embeddings_file, self.embeddings)

        # Save everything else using pickle
        state = {
            "alpha": self.alpha,
            "bm25": self.bm25,
            "corpus": self.corpus,
            "document_ids": self.document_ids,
            "model_name": "all-MiniLM-L6-v2",  # Default value for model name
        }

        # Try to get the actual model name if possible
        try:
            state["model_name"] = self.model.get_config_dict()["model_name"]
        except (AttributeError, KeyError):
            pass

        with open(retriever_cache_file, "wb") as f:
            pickle.dump(state, f)

    @classmethod
    def load(cls, retriever_cache_file: str, retriever_cache_embeddings_file: str):
        """Load retriever state from disk"""
        # Load the pickled state
        with open(retriever_cache_file, "rb") as f:
            state = pickle.load(f)

        # Create new instance
        retriever = cls(model_name=state["model_name"], alpha=state["alpha"])
        retriever.bm25 = state["bm25"]
        retriever.corpus = state["corpus"]
        retriever.document_ids = state.get("document_ids", {})

        # Load embeddings from numpy file if it exists
        if retriever_cache_embeddings_file.exists():
            retriever.embeddings = np.load(retriever_cache_embeddings_file)

        return retriever

    @classmethod
    def load_from_local_memory(
        cls, memories: Dict, model_name: str, alpha: float
    ) -> bool:
        """Load retriever state from memory"""
        all_docs = [
            ", ".join(m.keywords) for m in memories.values()
        ]  # [m.content for m in memories.values()]
        retriever = cls(model_name, alpha)
        retriever.add_documents(all_docs)
        return retriever

    def add_documents(self, documents: List[str]) -> bool:
        """One-time Add documents to both BM25 and semantic index"""
        if not documents:
            return

        # Tokenize for BM25
        tokenized_docs = [doc.lower().split() for doc in documents]
        self.bm25 = BM25Okapi(tokenized_docs)

        # Create embeddings
        self.embeddings = self.model.encode(documents)
        self.corpus = documents
        doc_idx = 0
        for document in documents:
            self.document_ids[document] = doc_idx
            doc_idx += 1

        return True

    def add_document(self, document: str) -> bool:
        """Add a single document to the retriever.

        Args:
            document: Text content to add

        Returns:
            bool: True if document was added, False if it was already present
        """
        # Check if document already exists
        if document in self.document_ids:
            return False

        # Add to corpus and get index
        doc_idx = len(self.corpus)
        self.corpus.append(document)
        self.document_ids[document] = doc_idx

        # Update BM25
        if self.bm25 is None:
            # First document, initialize BM25
            tokenized_corpus = [simple_tokenize(document)]
            self.bm25 = BM25Okapi(tokenized_corpus)
        else:
            # Add to existing BM25
            tokenized_doc = simple_tokenize(document)
            self.bm25.add_document(tokenized_doc)

        # Update embeddings
        doc_embedding = self.model.encode([document], convert_to_tensor=True)
        if self.embeddings is None:
            self.embeddings = doc_embedding
        else:
            self.embeddings = torch.cat([self.embeddings, doc_embedding])

        return True

    def retrieve(self, query: str, k: int = 5) -> List[int]:
        """Retrieve documents using hybrid scoring"""
        if not self.corpus:
            return []

        # Get BM25 scores
        tokenized_query = query.lower().split()
        bm25_scores = np.array(self.bm25.get_scores(tokenized_query))

        # Normalize BM25 scores if they exist
        if len(bm25_scores) > 0:
            bm25_scores = (bm25_scores - bm25_scores.min()) / (
                bm25_scores.max() - bm25_scores.min() + 1e-6
            )

        # Get semantic scores
        query_embedding = self.model.encode([query])[0]
        semantic_scores = cosine_similarity([query_embedding], self.embeddings)[0]

        # Combine scores
        hybrid_scores = self.alpha * bm25_scores + (1 - self.alpha) * semantic_scores

        # Get top k indices
        k = min(k, len(self.corpus))
        top_k_indices = np.argsort(hybrid_scores)[-k:][::-1]
        return top_k_indices.tolist()


class SimpleEmbeddingRetriever:
    """Simple retrieval system using only text embeddings."""

    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        """Initialize the simple embedding retriever.

        Args:
            model_name: Name of the SentenceTransformer model to use
        """
        self.model = SentenceTransformer(model_name)
        self.corpus = []
        self.embeddings = None
        self.document_ids = {}  # Map document content to its index

    def add_documents(self, documents: List[str]):
        """Add documents to the retriever."""
        # Reset if no existing documents
        if not self.corpus:
            self.corpus = documents
            # print("documents", documents, len(documents))
            self.embeddings = self.model.encode(documents)
            self.document_ids = {doc: idx for idx, doc in enumerate(documents)}
        else:
            # Append new documents
            start_idx = len(self.corpus)
            self.corpus.extend(documents)
            new_embeddings = self.model.encode(documents)
            if self.embeddings is None:
                self.embeddings = new_embeddings
            else:
                self.embeddings = np.vstack([self.embeddings, new_embeddings])
            for idx, doc in enumerate(documents):
                self.document_ids[doc] = start_idx + idx

    def search(self, query: str, k: int = 5) -> List[Dict[str, float]]:
        """Search for similar documents using cosine similarity.

        Args:
            query: Query text
            k: Number of results to return

        Returns:
            List of dicts with document text and score
        """
        if not self.corpus:
            return []
        # print("corpus", len(self.corpus), self.corpus)
        # Encode query
        query_embedding = self.model.encode([query])[0]

        # Calculate cosine similarities
        similarities = cosine_similarity([query_embedding], self.embeddings)[0]
        # Get top k results
        top_k_indices = np.argsort(similarities)[-k:][::-1]

        return top_k_indices

    def save(self, retriever_cache_file: str, retriever_cache_embeddings_file: str):
        """Save retriever state to disk"""
        # Save embeddings using numpy
        if self.embeddings is not None:
            np.save(retriever_cache_embeddings_file, self.embeddings)

        # Save other attributes
        state = {"corpus": self.corpus, "document_ids": self.document_ids}
        with open(retriever_cache_file, "wb") as f:
            pickle.dump(state, f)

    def load(self, retriever_cache_file: str, retriever_cache_embeddings_file: str):
        """Load retriever state from disk"""
        print(
            f"Loading retriever from {retriever_cache_file} and {retriever_cache_embeddings_file}"
        )

        # Load embeddings
        if os.path.exists(retriever_cache_embeddings_file):
            print(f"Loading embeddings from {retriever_cache_embeddings_file}")
            self.embeddings = np.load(retriever_cache_embeddings_file)
            print(f"Embeddings shape: {self.embeddings.shape}")
        else:
            print(f"Embeddings file not found: {retriever_cache_embeddings_file}")

        # Load other attributes
        if os.path.exists(retriever_cache_file):
            print(f"Loading corpus from {retriever_cache_file}")
            with open(retriever_cache_file, "rb") as f:
                state = pickle.load(f)
                self.corpus = state["corpus"]
                self.document_ids = state["document_ids"]
                print(f"Loaded corpus with {len(self.corpus)} documents")
        else:
            print(f"Corpus file not found: {retriever_cache_file}")

        return self

    @classmethod
    def load_from_local_memory(
        cls, memories: Dict, model_name: str
    ) -> "SimpleEmbeddingRetriever":
        """Load retriever state from memory"""
        # Create documents combining content and metadata for each memory
        all_docs = []
        for m in memories.values():
            metadata_text = f"{m.context} {' '.join(m.keywords)} {' '.join(m.tags)}"
            doc = f"{m.content} , {metadata_text}"
            all_docs.append(doc)

        # Create and initialize retriever
        retriever = cls(model_name)
        retriever.add_documents(all_docs)
        return retriever


class AgenticMemorySystem:
    """Memory management system with embedding-based retrieval"""

    def __init__(
        self,
        model_name: str = "all-MiniLM-L6-v2",
        llm_backend: str = "sglang",
        llm_model: str = "gpt-4o-mini",
        evo_threshold: int = 100,
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
        sglang_host: str = "http://localhost",
        sglang_port: int = 30000,
    ):
        self.memories = {}  # id -> MemoryNote
        self.retriever = SimpleEmbeddingRetriever(model_name)
        self.llm_controller = LLMController(
            llm_backend, llm_model, api_key, api_base, sglang_host, sglang_port
        )
        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:
                                {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? 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": ["neighbor_memory_ids"],
                                    "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"]],
                                }}
                                """
        self.evo_cnt = 0
        self.evo_threshold = evo_threshold

    def add_note(self, content: str, time: str = None, **kwargs) -> str:
        """Add a new memory note"""
        note = MemoryNote(
            content=content,
            llm_controller=self.llm_controller,
            timestamp=time,
            **kwargs,
        )

        # Update retriever with all documents
        # all_docs = [m.content for m in self.memories.values()]
        evo_label, note = self.process_memory(note)
        self.memories[note.id] = note
        self.retriever.add_documents(
            [
                "content:"
                + note.content
                + " context:"
                + note.context
                + " keywords: "
                + ", ".join(note.keywords)
                + " tags: "
                + ", ".join(note.tags)
            ]
        )
        if evo_label == True:
            self.evo_cnt += 1
            if self.evo_cnt % self.evo_threshold == 0:
                self.consolidate_memories()
        return note.id

    def consolidate_memories(self):
        """Consolidate memories: update retriever with new documents

        This function re-initializes the retriever and updates it with all memory documents,
        including their context, keywords, and tags to ensure the retrieval system has the
        latest state of all memories.
        """
        # Reset the retriever with the same model
        try:
            # Try to get model name through get_config_dict if available
            model_name = self.retriever.model.get_config_dict()["model_name"]
        except (AttributeError, KeyError):
            # Fallback: use the model name from the class initialization
            model_name = "all-MiniLM-L6-v2"

        self.retriever = SimpleEmbeddingRetriever(model_name)

        # Re-add all memory documents with their metadata
        for memory in self.memories.values():
            # Combine memory metadata into a single searchable document
            metadata_text = (
                f"{memory.context} {' '.join(memory.keywords)} {' '.join(memory.tags)}"
            )
            # Add both the content and metadata as separate documents for better retrieval
            self.retriever.add_documents([memory.content + " , " + metadata_text])

    def process_memory(self, note: MemoryNote) -> bool:
        """Process a memory note and return an evolution label"""
        neighbor_memory, indices = self.find_related_memories(note.content, k=5)
        prompt_memory = self.evolution_system_prompt.format(
            context=note.context,
            content=note.content,
            keywords=note.keywords,
            nearest_neighbors_memories=neighbor_memory,
            neighbor_number=len(indices),
        )
        print("prompt_memory", prompt_memory)
        response = self.llm_controller.llm.get_completion(
            prompt_memory,
            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": "integer"},
                            },
                            "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,
                },
            },
        )
        try:
            print("response", response, type(response))
            # Clean the response in case there's extra text
            response_cleaned = response.strip()
            # Try to find JSON content if wrapped in other text
            if not response_cleaned.startswith("{"):
                start_idx = response_cleaned.find("{")
                if start_idx != -1:
                    response_cleaned = response_cleaned[start_idx:]
            if not response_cleaned.endswith("}"):
                end_idx = response_cleaned.rfind("}")
                if end_idx != -1:
                    response_cleaned = response_cleaned[: end_idx + 1]

            response_json = json.loads(response_cleaned)
            print("response_json", response_json, type(response_json))
        except json.JSONDecodeError as e:
            print(f"JSON parsing error: {e}")
            print(f"Raw response: {response}")
            # Return default values for failed parsing
            return False, note
        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"]
                    noteslist = list(self.memories.values())
                    notes_id = list(self.memories.keys())
                    print("indices", indices)
                    # if slms output less than the number of neighbors, use the sequential order of new tags and context.
                    for i in range(min(len(indices), len(new_tags_neighborhood))):
                        # find some memory
                        tag = new_tags_neighborhood[i]
                        if i < len(new_context_neighborhood):
                            context = new_context_neighborhood[i]
                        else:
                            context = noteslist[indices[i]].context
                        memorytmp_idx = indices[i]
                        notetmp = noteslist[memorytmp_idx]
                        # add tag to memory
                        notetmp.tags = tag
                        notetmp.context = context
                        self.memories[notes_id[memorytmp_idx]] = notetmp
        return should_evolve, note

    def find_related_memories(self, query: str, k: int = 5) -> List[MemoryNote]:
        """Find related memories using hybrid retrieval"""
        if not self.memories:
            return "", []

        # Get indices of related memories
        # indices = self.retriever.retrieve(query_note.content, k)
        indices = self.retriever.search(query, k)

        # Convert to list of memories
        all_memories = list(self.memories.values())
        memory_str = ""
        # print("indices", indices)
        # print("all_memories", all_memories)
        for i in indices:
            memory_str += (
                "memory index:"
                + str(i)
                + "\t talk start time:"
                + all_memories[i].timestamp
                + "\t memory content: "
                + all_memories[i].content
                + "\t memory context: "
                + all_memories[i].context
                + "\t memory keywords: "
                + str(all_memories[i].keywords)
                + "\t memory tags: "
                + str(all_memories[i].tags)
                + "\n"
            )
        return memory_str, indices

    def find_related_memories_raw(self, query: str, k: int = 5) -> List[MemoryNote]:
        """Find related memories using hybrid retrieval"""
        if not self.memories:
            return []

        # Get indices of related memories
        # indices = self.retriever.retrieve(query_note.content, k)
        indices = self.retriever.search(query, k)

        # Convert to list of memories
        all_memories = list(self.memories.values())
        memory_str = ""
        for i in indices:
            j = 0
            memory_str += (
                "talk start time:"
                + all_memories[i].timestamp
                + "memory content: "
                + all_memories[i].content
                + "memory context: "
                + all_memories[i].context
                + "memory keywords: "
                + str(all_memories[i].keywords)
                + "memory tags: "
                + str(all_memories[i].tags)
                + "\n"
            )
            neighborhood = all_memories[i].links
            for neighbor in neighborhood:
                memory_str += (
                    "talk start time:"
                    + all_memories[neighbor].timestamp
                    + "memory content: "
                    + all_memories[neighbor].content
                    + "memory context: "
                    + all_memories[neighbor].context
                    + "memory keywords: "
                    + str(all_memories[neighbor].keywords)
                    + "memory tags: "
                    + str(all_memories[neighbor].tags)
                    + "\n"
                )
                if j >= k:
                    break
                j += 1
        return memory_str


def run_tests():
    """Run system tests"""
    print("Starting Memory System Tests...")

    # Initialize memory system with OpenAI backend
    memory_system = AgenticMemorySystem(
        model_name="all-MiniLM-L6-v2", llm_backend="openai", llm_model="gpt-4o-mini"
    )

    print("\nAdding test memories...")

    # Add test memories - only content is required
    memory_ids = []
    memory_ids.append(
        memory_system.add_note(
            "Neural networks are composed of layers of neurons that process information."
        )
    )

    memory_ids.append(
        memory_system.add_note(
            "Data preprocessing involves cleaning and transforming raw data for model training."
        )
    )

    print("\nQuerying for related memories...")
    query = MemoryNote(
        content="How do neural networks process data?",
        llm_controller=memory_system.llm_controller,
    )

    related = memory_system.find_related_memories(query.content, k=2)
    print("related", related)
    print("\nResults:")
    for i, memory in enumerate(related, 1):
        print(f"\n{i}. Memory:")
        print(f"Content: {memory.content}")
        print(f"Category: {memory.category}")
        print(f"Keywords: {memory.keywords}")
        print(f"Tags: {memory.tags}")
        print(f"Context: {memory.context}")
        print("-" * 50)


if __name__ == "__main__":
    run_tests()
