import json
import os
import time

from datetime import datetime, timezone
from pathlib import Path
from threading import Lock
from typing import Any, Literal

from memos.configs.mem_os import MOSConfig
from memos.context.context import ContextThreadPoolExecutor
from memos.llms.factory import LLMFactory
from memos.log import get_logger
from memos.mem_cube.general import GeneralMemCube
from memos.mem_reader.factory import MemReaderFactory
from memos.mem_scheduler.general_scheduler import GeneralScheduler
from memos.mem_scheduler.scheduler_factory import SchedulerFactory
from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
from memos.mem_scheduler.schemas.task_schemas import (
    ADD_TASK_LABEL,
    ANSWER_TASK_LABEL,
    MEM_READ_TASK_LABEL,
    PREF_ADD_TASK_LABEL,
    QUERY_TASK_LABEL,
)
from memos.mem_user.user_manager import UserManager, UserRole
from memos.memories.activation.item import ActivationMemoryItem
from memos.memories.parametric.item import ParametricMemoryItem
from memos.memories.textual.item import TextualMemoryItem, TextualMemoryMetadata
from memos.memos_tools.thread_safe_dict_segment import OptimizedThreadSafeDict
from memos.templates.mos_prompts import QUERY_REWRITING_PROMPT
from memos.types import ChatHistory, MessageList, MOSSearchResult


logger = get_logger(__name__)


class MOSCore:
    """
    The MOSCore (Memory Operating System Core) class manages multiple MemCube objects and their operations.
    It provides methods for creating, searching, updating, and deleting MemCubes, supporting multi-user scenarios.
    MOSCore acts as an operating system layer for handling and orchestrating MemCube instances.
    """

    def __init__(self, config: MOSConfig, user_manager: UserManager | None = None):
        self.config = config
        self.user_id = config.user_id
        self.session_id = config.session_id
        self.chat_llm = LLMFactory.from_config(config.chat_model)
        self.mem_reader = MemReaderFactory.from_config(config.mem_reader)
        self.chat_history_manager: dict[str, ChatHistory] = {}
        # use thread safe dict for multi-user product-server scenario
        self.mem_cubes: OptimizedThreadSafeDict[str, GeneralMemCube] = (
            OptimizedThreadSafeDict() if user_manager is not None else {}
        )
        self._register_chat_history()

        # Use provided user_manager or create a new one
        if user_manager is not None:
            self.user_manager = user_manager
        else:
            self.user_manager = UserManager(user_id=self.user_id if self.user_id else "root")

        # Validate user exists
        if not self.user_manager.validate_user(self.user_id):
            raise ValueError(
                f"User '{self.user_id}' does not exist or is inactive. Please create user first."
            )

        # Initialize mem_scheduler
        self._mem_scheduler_lock = Lock()
        self.enable_mem_scheduler = self.config.get("enable_mem_scheduler", False)
        if self.enable_mem_scheduler:
            self._mem_scheduler = self._initialize_mem_scheduler()
            self._mem_scheduler.mem_cubes = self.mem_cubes
            self._mem_scheduler.mem_reader = self.mem_reader
        else:
            self._mem_scheduler: GeneralScheduler = None

        logger.info(f"MOS initialized for user: {self.user_id}")

    @property
    def mem_scheduler(self) -> GeneralScheduler:
        """Lazy-loaded property for memory scheduler."""
        if self.enable_mem_scheduler and self._mem_scheduler is None:
            self._initialize_mem_scheduler()
        self._mem_scheduler.mem_cubes = self.mem_cubes
        return self._mem_scheduler

    @mem_scheduler.setter
    def mem_scheduler(self, value: GeneralScheduler | None) -> None:
        """Setter for memory scheduler with validation.

        Args:
            value: GeneralScheduler instance or None to disable
        Raises:
            TypeError: If value is neither GeneralScheduler nor None
        """
        with self._mem_scheduler_lock:
            if value is not None and not isinstance(value, GeneralScheduler):
                raise TypeError(f"Expected GeneralScheduler or None, got {type(value)}")

            self._mem_scheduler = value
            self._mem_scheduler.mem_cubes = self.mem_cubes

            if value:
                logger.info("Memory scheduler manually set")
            else:
                logger.debug("Memory scheduler cleared")

    def _initialize_mem_scheduler(self) -> GeneralScheduler:
        """Initialize the memory scheduler on first access."""
        if not self.config.enable_mem_scheduler:
            logger.debug("Memory scheduler is disabled in config")
            self._mem_scheduler = None
            return self._mem_scheduler
        elif not hasattr(self.config, "mem_scheduler"):
            logger.error("Config of Memory scheduler is not available")
            self._mem_scheduler = None
            return self._mem_scheduler
        else:
            logger.info("Initializing memory scheduler...")
            scheduler_config = self.config.mem_scheduler
            self._mem_scheduler = SchedulerFactory.from_config(scheduler_config)
            # Validate required components
            if not hasattr(self.mem_reader, "llm"):
                raise AttributeError(
                    f"Memory reader of type {type(self.mem_reader).__name__} "
                    "missing required 'llm' attribute"
                )
            else:
                # Configure scheduler general_modules
                self._mem_scheduler.initialize_modules(
                    chat_llm=self.chat_llm,
                    process_llm=self.mem_reader.general_llm,
                    db_engine=self.user_manager.engine,
                )
            self._mem_scheduler.start()
            return self._mem_scheduler

    def mem_scheduler_on(self) -> bool:
        if not self.config.enable_mem_scheduler or self._mem_scheduler is None:
            logger.error("Cannot start scheduler: disabled in configuration")

        try:
            self._mem_scheduler.start()
            logger.info("Memory scheduler service started")
            return True
        except Exception as e:
            logger.error(f"Failed to start scheduler: {e!s}")
            return False

    def mem_scheduler_off(self) -> bool:
        if not self.config.enable_mem_scheduler:
            logger.error("Cannot stop scheduler: disabled in configuration")

        if self._mem_scheduler is None:
            logger.warning("No scheduler instance to stop")
            return False

        try:
            self._mem_scheduler.stop()
            logger.info("Memory scheduler service stopped")
            return True
        except Exception as e:
            logger.error(f"Failed to stop scheduler: {e!s}")
            return False

    def mem_reorganizer_on(self) -> bool:
        pass

    def mem_reorganizer_off(self) -> bool:
        """temporally implement"""
        for mem_cube in self.mem_cubes.values():
            logger.info(f"try to close reorganizer for {mem_cube.text_mem.config.cube_id}")
            if mem_cube.text_mem and mem_cube.text_mem.is_reorganize:
                logger.info(f"close reorganizer for {mem_cube.text_mem.config.cube_id}")
                mem_cube.text_mem.memory_manager.close()
                mem_cube.text_mem.memory_manager.wait_reorganizer()

    def mem_reorganizer_wait(self) -> bool:
        for mem_cube in self.mem_cubes.values():
            logger.info(f"try to close reorganizer for {mem_cube.text_mem.config.cube_id}")
            if mem_cube.text_mem and mem_cube.text_mem.is_reorganize:
                logger.info(f"close reorganizer for {mem_cube.text_mem.config.cube_id}")
                mem_cube.text_mem.memory_manager.wait_reorganizer()

    def _register_chat_history(
        self, user_id: str | None = None, session_id: str | None = None
    ) -> None:
        """Initialize chat history with user ID."""
        self.chat_history_manager[user_id] = ChatHistory(
            user_id=user_id if user_id is not None else self.user_id,
            session_id=session_id if session_id is not None else self.session_id,
            created_at=datetime.now(timezone.utc),
            total_messages=0,
            chat_history=[],
        )

    def _validate_user_exists(self, user_id: str) -> None:
        """Validate user exists and is active.

        Args:
            user_id (str): The user ID to validate.

        Raises:
            ValueError: If user doesn't exist or is inactive.
        """
        if not self.user_manager.validate_user(user_id):
            raise ValueError(
                f"User '{user_id}' does not exist or is inactive. Please register the user first."
            )

    def _validate_cube_access(self, user_id: str, cube_id: str) -> None:
        """Validate user has access to the cube.

        Args:
            user_id (str): The user ID to validate.
            cube_id (str): The cube ID to validate.

        Raises:
            ValueError: If user doesn't have access to the cube.
        """
        # First validate user exists
        self._validate_user_exists(user_id)

        # Then validate cube access
        if not self.user_manager.validate_user_cube_access(user_id, cube_id):
            raise ValueError(
                f"User '{user_id}' does not have access to cube '{cube_id}'. Please register the cube first or request access."
            )

    def _get_all_documents(self, path: str) -> list[str]:
        """Get all documents from path.

        Args:
            path (str): The path to get documents.

        Returns:
            list[str]: The list of documents.
        """
        documents = []

        path_obj = Path(path)
        doc_extensions = {".txt", ".pdf", ".json", ".md", ".ppt", ".pptx"}
        for file_path in path_obj.rglob("*"):
            if file_path.is_file() and (file_path.suffix.lower() in doc_extensions):
                documents.append(str(file_path))
        return documents

    def chat(self, query: str, user_id: str | None = None, base_prompt: str | None = None) -> str:
        """
        Chat with the MOS.

        Args:
            query (str): The user's query.
            user_id (str, optional): The user ID for the chat session. Defaults to the user ID from the config.
            base_prompt (str, optional): A custom base prompt to use for the chat.
                It can be a template string with a `{memories}` placeholder.
                If not provided, a default prompt is used.

        Returns:
            str: The response from the MOS.
        """
        target_user_id = user_id if user_id is not None else self.user_id
        accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
        user_cube_ids = [cube.cube_id for cube in accessible_cubes]
        if target_user_id not in self.chat_history_manager:
            self._register_chat_history(target_user_id)

        chat_history = self.chat_history_manager[target_user_id]

        if self.config.enable_textual_memory and self.mem_cubes:
            memories_all = []
            for mem_cube_id, mem_cube in self.mem_cubes.items():
                if mem_cube_id not in user_cube_ids:
                    continue
                if not mem_cube.text_mem:
                    continue

                # submit message to scheduler
                if self.enable_mem_scheduler and self.mem_scheduler is not None:
                    message_item = ScheduleMessageItem(
                        user_id=target_user_id,
                        mem_cube_id=mem_cube_id,
                        label=QUERY_TASK_LABEL,
                        content=query,
                        timestamp=datetime.utcnow(),
                    )
                    self.mem_scheduler.submit_messages(messages=[message_item])

                memories = mem_cube.text_mem.search(
                    query,
                    top_k=self.config.top_k,
                    info={
                        "user_id": target_user_id,
                        "session_id": self.session_id,
                        "chat_history": chat_history.chat_history,
                    },
                )
                memories_all.extend(memories)
            logger.info(f"🧠 [Memory] Searched memories:\n{self._str_memories(memories_all)}\n")
            system_prompt = self._build_system_prompt(memories_all, base_prompt=base_prompt)
        else:
            system_prompt = self._build_system_prompt(base_prompt=base_prompt)
        current_messages = [
            {"role": "system", "content": system_prompt},
            *chat_history.chat_history,
            {"role": "user", "content": query},
        ]
        past_key_values = None

        if self.config.enable_activation_memory:
            if self.config.chat_model.backend not in ["huggingface", "huggingface_singleton"]:
                logger.error(
                    "Activation memory only used for huggingface backend. Skipping activation memory."
                )
            else:
                # TODO this only one cubes
                for mem_cube_id, mem_cube in self.mem_cubes.items():
                    if mem_cube_id not in user_cube_ids:
                        continue
                    if mem_cube.act_mem:
                        kv_cache = next(iter(mem_cube.act_mem.get_all()), None)
                        past_key_values = (
                            kv_cache.memory if (kv_cache and hasattr(kv_cache, "memory")) else None
                        )
                    break
            # Generate response
            response = self.chat_llm.generate(current_messages, past_key_values=past_key_values)
        else:
            response = self.chat_llm.generate(current_messages)
        logger.info(f"🤖 [Assistant] {response}\n")
        chat_history.chat_history.append({"role": "user", "content": query})
        chat_history.chat_history.append({"role": "assistant", "content": response})
        self.chat_history_manager[user_id] = chat_history

        # submit message to scheduler
        for accessible_mem_cube in accessible_cubes:
            mem_cube_id = accessible_mem_cube.cube_id
            mem_cube = self.mem_cubes[mem_cube_id]
            if self.enable_mem_scheduler and self.mem_scheduler is not None:
                message_item = ScheduleMessageItem(
                    user_id=target_user_id,
                    mem_cube_id=mem_cube_id,
                    label=ANSWER_TASK_LABEL,
                    content=response,
                    timestamp=datetime.utcnow(),
                )
                self.mem_scheduler.submit_messages(messages=[message_item])

        return response

    def _build_system_prompt(
        self,
        memories: list[TextualMemoryItem] | list[str] | None = None,
        base_prompt: str | None = None,
        **kwargs,
    ) -> str:
        """Build system prompt with optional memories context."""
        if base_prompt is None:
            base_prompt = (
                "You are a knowledgeable and helpful AI assistant. "
                "You have access to conversation memories that help you provide more personalized responses. "
                "Use the memories to understand the user's context, preferences, and past interactions. "
                "If memories are provided, reference them naturally when relevant, but don't explicitly mention having memories."
            )

        memory_context = ""
        if memories:
            memory_list = []
            for i, memory in enumerate(memories, 1):
                if isinstance(memory, TextualMemoryItem):
                    text_memory = memory.memory
                else:
                    if not isinstance(memory, str):
                        logger.error("Unexpected memory type.")
                    text_memory = memory
                memory_list.append(f"{i}. {text_memory}")
            memory_context = "\n".join(memory_list)

        if "{memories}" in base_prompt:
            return base_prompt.format(memories=memory_context)
        elif memories:
            # For backward compatibility, append memories if no placeholder is found
            memory_context_with_header = "\n\n## Memories:\n" + memory_context
            return base_prompt + memory_context_with_header
        return base_prompt

    def _str_memories(
        self, memories: list[TextualMemoryItem], mode: Literal["concise", "full"] = "full"
    ) -> str:
        """Format memories for display."""
        if not memories:
            return "No memories."
        if mode == "concise":
            return "\n".join(f"{i + 1}. {memory.memory}" for i, memory in enumerate(memories))
        elif mode == "full":
            return "\n".join(f"{i + 1}. {memory}" for i, memory in enumerate(memories))

    def clear_messages(self, user_id: str | None = None) -> None:
        """Clear chat history."""
        user_id = user_id if user_id is not None else self.user_id
        self._register_chat_history(user_id)

    def create_user(
        self, user_id: str, role: UserRole = UserRole.USER, user_name: str | None = None
    ) -> str:
        """Create a new user.

        Args:
            user_name (str): Name of the user.
            role (UserRole): Role of the user.
            user_id (str, optional): Custom user ID.

        Returns:
            str: The created user ID.
        """
        if not user_name:
            user_name = user_id
        return self.user_manager.create_user(user_name, role, user_id)

    def list_users(self) -> list:
        """List all active users.

        Returns:
            list: List of user information dictionaries.
        """
        users = self.user_manager.list_users()
        return [
            {
                "user_id": user.user_id,
                "user_name": user.user_name,
                "role": user.role.value,
                "created_at": user.created_at.isoformat(),
                "is_active": user.is_active,
            }
            for user in users
        ]

    def create_cube_for_user(
        self,
        cube_name: str,
        owner_id: str,
        cube_path: str | None = None,
        cube_id: str | None = None,
    ) -> str:
        """Create a new cube for the current user.

        Args:
            cube_name (str): Name of the cube.
            cube_path (str, optional): Path to the cube.
            cube_id (str, optional): Custom cube ID.

        Returns:
            str: The created cube ID.
        """
        return self.user_manager.create_cube(cube_name, owner_id, cube_path, cube_id)

    def register_mem_cube(
        self,
        mem_cube_name_or_path: str | GeneralMemCube,
        mem_cube_id: str | None = None,
        user_id: str | None = None,
    ) -> None:
        """
        Register a MemCube with the MOS.

        Args:
            mem_cube_name_or_path (str): The name or path of the MemCube to register.
            mem_cube_id (str, optional): The identifier for the MemCube. If not provided, a default ID is used.
        """
        target_user_id = user_id if user_id is not None else self.user_id
        self._validate_user_exists(target_user_id)

        if mem_cube_id is None:
            if isinstance(mem_cube_name_or_path, GeneralMemCube):
                mem_cube_id = f"cube_{target_user_id}"
            else:
                mem_cube_id = mem_cube_name_or_path

        if mem_cube_id in self.mem_cubes:
            logger.info(f"MemCube with ID {mem_cube_id} already in MOS, skip install.")
        else:
            if isinstance(mem_cube_name_or_path, GeneralMemCube):
                self.mem_cubes[mem_cube_id] = mem_cube_name_or_path
                logger.info(f"register new cube {mem_cube_id} for user {target_user_id}")
            elif os.path.exists(mem_cube_name_or_path):
                mem_cube_obj = GeneralMemCube.init_from_dir(mem_cube_name_or_path)
                self.mem_cubes[mem_cube_id] = mem_cube_obj
            else:
                logger.warning(
                    f"MemCube {mem_cube_name_or_path} does not exist, try to init from remote repo."
                )
                mem_cube_obj = GeneralMemCube.init_from_remote_repo(mem_cube_name_or_path)
                self.mem_cubes[mem_cube_id] = mem_cube_obj
        # Check if cube already exists in database
        existing_cube = self.user_manager.get_cube(mem_cube_id)

        # check the embedder is it consistent with MOSConfig
        if hasattr(
            self.mem_cubes[mem_cube_id].text_mem.config, "embedder"
        ) and self.config.mem_reader.config.embedder != (
            cube_embedder := self.mem_cubes[mem_cube_id].text_mem.config.embedder
        ):
            logger.warning(
                f"Cube Embedder is not consistent with MOSConfig for cube: {mem_cube_id}, will use Cube Embedder: {cube_embedder}"
            )

        if existing_cube:
            # Cube exists, just add user to cube if not already associated
            if not self.user_manager.validate_user_cube_access(target_user_id, mem_cube_id):
                success = self.user_manager.add_user_to_cube(target_user_id, mem_cube_id)
                if success:
                    logger.info(f"User {target_user_id} added to existing cube {mem_cube_id}")
                else:
                    logger.error(f"Failed to add user {target_user_id} to cube {mem_cube_id}")
            else:
                logger.info(f"User {target_user_id} already has access to cube {mem_cube_id}")
        else:
            # Cube doesn't exist, create it
            self.create_cube_for_user(
                cube_name=mem_cube_name_or_path
                if not isinstance(mem_cube_name_or_path, GeneralMemCube)
                else mem_cube_id,
                owner_id=target_user_id,
                cube_id=mem_cube_id,
                cube_path=mem_cube_name_or_path
                if not isinstance(mem_cube_name_or_path, GeneralMemCube)
                else "init",
            )
            logger.info(f"register new cube {mem_cube_id} for user {target_user_id}")

    def unregister_mem_cube(self, mem_cube_id: str, user_id: str | None = None) -> None:
        """
        Unregister a MemCube by its identifier.

        Args:
            mem_cube_id (str): The identifier of the MemCube to unregister.
        """
        if mem_cube_id in self.mem_cubes:
            del self.mem_cubes[mem_cube_id]
        else:
            raise ValueError(f"MemCube with ID {mem_cube_id} does not exist.")

    def search(
        self,
        query: str,
        user_id: str | None = None,
        install_cube_ids: list[str] | None = None,
        top_k: int | None = None,
        mode: Literal["fast", "fine"] = "fast",
        internet_search: bool = False,
        moscube: bool = False,
        session_id: str | None = None,
        **kwargs,
    ) -> MOSSearchResult:
        """
        Search for textual memories across all registered MemCubes.

        Args:
            query (str): The search query.
            user_id (str, optional): The identifier of the user to search for.
                If None, the default user is used.
            install_cube_ids (list[str], optional): The list of MemCube IDs to install.
                If None, all MemCube for the user is used.

        Returns:
            MemoryResult: A dictionary containing the search results.
        """
        target_session_id = session_id if session_id is not None else self.session_id
        target_user_id = user_id if user_id is not None else self.user_id

        self._validate_user_exists(target_user_id)
        # Get all cubes accessible by the target user
        accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
        user_cube_ids = [cube.cube_id for cube in accessible_cubes]

        logger.info(
            f"User {target_user_id} has access to {len(user_cube_ids)} cubes: {user_cube_ids}"
        )
        if target_user_id not in self.chat_history_manager:
            self._register_chat_history(target_user_id)
        chat_history = self.chat_history_manager[target_user_id]

        # Create search filter if session_id is provided
        search_filter = None
        if session_id is not None:
            search_filter = {"session_id": session_id}

        result: MOSSearchResult = {
            "text_mem": [],
            "act_mem": [],
            "para_mem": [],
            "pref_mem": [],
        }
        if install_cube_ids is None:
            install_cube_ids = user_cube_ids
        # create exist dict in mem_cubes and avoid  one search slow
        tmp_mem_cubes = {}
        time_start_cube_get = time.time()
        for mem_cube_id in install_cube_ids:
            if mem_cube_id in self.mem_cubes:
                tmp_mem_cubes[mem_cube_id] = self.mem_cubes.get(mem_cube_id)
        logger.info(
            f"time search: transform cube time user_id: {target_user_id} time is: {time.time() - time_start_cube_get}"
        )

        for mem_cube_id, mem_cube in tmp_mem_cubes.items():
            # Define internal functions for parallel search execution
            def search_textual_memory(cube_id, cube):
                if (
                    (cube_id in install_cube_ids)
                    and (cube.text_mem is not None)
                    and self.config.enable_textual_memory
                ):
                    time_start = time.time()
                    memories = cube.text_mem.search(
                        query,
                        top_k=top_k if top_k else self.config.top_k,
                        mode=mode,
                        manual_close_internet=not internet_search,
                        info={
                            "user_id": target_user_id,
                            "session_id": target_session_id,
                            "chat_history": chat_history.chat_history,
                        },
                        moscube=moscube,
                        search_filter=search_filter,
                    )
                    search_time_end = time.time()
                    logger.info(
                        f"🧠 [Memory] Searched memories from {cube_id}:\n{self._str_memories(memories)}\n"
                    )
                    logger.info(
                        f"time search graph: search graph time user_id: {target_user_id} time is: {search_time_end - time_start}"
                    )
                    return {"cube_id": cube_id, "memories": memories}
                return None

            def search_preference_memory(cube_id, cube):
                if (
                    (cube_id in install_cube_ids)
                    and (cube.pref_mem is not None)
                    and self.config.enable_preference_memory
                ):
                    time_start = time.time()
                    memories = cube.pref_mem.search(
                        query,
                        top_k=top_k if top_k else self.config.top_k,
                        info={
                            "user_id": target_user_id,
                            "session_id": self.session_id,
                            "chat_history": chat_history.chat_history,
                        },
                    )
                    search_time_end = time.time()
                    logger.info(
                        f"🧠 [Memory] Searched preferences from {cube_id}:\n{self._str_memories(memories)}\n"
                    )
                    logger.info(
                        f"time search pref: search pref time user_id: {target_user_id} time is: {search_time_end - time_start}"
                    )
                    return {"cube_id": cube_id, "memories": memories}
                return None

            # Execute both search functions in parallel
            with ContextThreadPoolExecutor(max_workers=2) as executor:
                text_future = executor.submit(search_textual_memory, mem_cube_id, mem_cube)
                pref_future = executor.submit(search_preference_memory, mem_cube_id, mem_cube)

                # Wait for both tasks to complete and collect results
                text_result = text_future.result()
                pref_result = pref_future.result()

                # Add results to the main result dictionary
                if text_result is not None:
                    result["text_mem"].append(text_result)
                if pref_result is not None:
                    result["pref_mem"].append(pref_result)

        return result

    def add(
        self,
        messages: MessageList | None = None,
        memory_content: str | None = None,
        doc_path: str | None = None,
        mem_cube_id: str | None = None,
        user_id: str | None = None,
        session_id: str | None = None,
        task_id: str | None = None,  # New: Add task_id parameter
        **kwargs,
    ) -> None:
        """
        Add textual memories to a MemCube.

        Args:
            messages (Union[MessageList, str]): The path to a document or a list of messages.
            memory_content (str, optional): The content of the memory to add.
            doc_path (str, optional): The path to the document associated with the memory.
            mem_cube_id (str, optional): The identifier of the MemCube to add the memories to.
                If None, the default MemCube for the user is used.
            user_id (str, optional): The identifier of the user to add the memories to.
                If None, the default user is used.
            session_id (str, optional): session_id
        """
        # user input messages
        assert (messages is not None) or (memory_content is not None) or (doc_path is not None), (
            "messages_or_doc_path or memory_content or doc_path must be provided."
        )
        # TODO: asure that session_id is a valid string
        time_start = time.time()

        target_session_id = session_id if session_id else self.session_id
        target_user_id = user_id if user_id is not None else self.user_id
        if mem_cube_id is None:
            # Try to find a default cube for the user
            accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
            if not accessible_cubes:
                raise ValueError(
                    f"No accessible cubes found for user '{target_user_id}'. Please register a cube first."
                )
            mem_cube_id = accessible_cubes[0].cube_id  # TODO not only first
        else:
            self._validate_cube_access(target_user_id, mem_cube_id)
        logger.info(
            f"time add: get mem_cube_id time user_id: {target_user_id} time is: {time.time() - time_start}"
        )

        if mem_cube_id not in self.mem_cubes:
            raise ValueError(f"MemCube '{mem_cube_id}' is not loaded. Please register.")

        sync_mode = self.mem_cubes[mem_cube_id].text_mem.mode
        if sync_mode == "async":
            assert self.mem_scheduler is not None, (
                "Mem-Scheduler must be working when use asynchronous memory adding."
            )
        logger.debug(f"Mem-reader mode is: {sync_mode}")

        def process_textual_memory():
            if (
                (messages is not None)
                and self.config.enable_textual_memory
                and self.mem_cubes[mem_cube_id].text_mem
            ):
                if self.mem_cubes[mem_cube_id].config.text_mem.backend != "tree_text":
                    add_memory = []
                    metadata = TextualMemoryMetadata(
                        user_id=target_user_id, session_id=target_session_id, source="conversation"
                    )
                    for message in messages:
                        add_memory.append(
                            TextualMemoryItem(memory=message["content"], metadata=metadata)
                        )
                    self.mem_cubes[mem_cube_id].text_mem.add(add_memory)
                else:
                    messages_list = [messages]
                    memories = self.mem_reader.get_memory(
                        messages_list,
                        type="chat",
                        info={"user_id": target_user_id, "session_id": target_session_id},
                        mode="fast" if sync_mode == "async" else "fine",
                    )
                    memories_flatten = [m for m_list in memories for m in m_list]
                    mem_ids: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(memories_flatten)
                    logger.info(
                        f"Added memory user {target_user_id} to memcube {mem_cube_id}: {mem_ids}"
                    )
                    # submit messages for scheduler
                    if self.enable_mem_scheduler and self.mem_scheduler is not None:
                        if sync_mode == "async":
                            message_item = ScheduleMessageItem(
                                user_id=target_user_id,
                                mem_cube_id=mem_cube_id,
                                label=MEM_READ_TASK_LABEL,
                                content=json.dumps(mem_ids),
                                timestamp=datetime.utcnow(),
                                task_id=task_id,
                            )
                            self.mem_scheduler.submit_messages(messages=[message_item])
                        else:
                            message_item = ScheduleMessageItem(
                                user_id=target_user_id,
                                mem_cube_id=mem_cube_id,
                                label=ADD_TASK_LABEL,
                                content=json.dumps(mem_ids),
                                timestamp=datetime.utcnow(),
                                task_id=task_id,
                            )
                            logger.info(
                                f"[DIAGNOSTIC] core.add: Submitting message to scheduler: {message_item.model_dump_json(indent=2)}"
                            )
                            self.mem_scheduler.submit_messages(messages=[message_item])

        def process_preference_memory():
            if (
                (messages is not None)
                and self.config.enable_preference_memory
                and self.mem_cubes[mem_cube_id].pref_mem
            ):
                messages_list = [messages]
                if sync_mode == "sync":
                    pref_memories = self.mem_cubes[mem_cube_id].pref_mem.get_memory(
                        messages_list,
                        type="chat",
                        info={
                            "user_id": target_user_id,
                            "session_id": self.session_id,
                            "mem_cube_id": mem_cube_id,
                        },
                    )
                    pref_ids = self.mem_cubes[mem_cube_id].pref_mem.add(pref_memories)
                    logger.info(
                        f"Added preferences user {target_user_id} to memcube {mem_cube_id}: {pref_ids}"
                    )
                elif sync_mode == "async":
                    assert self.mem_scheduler is not None, (
                        "Mem-Scheduler must be working when use asynchronous memory adding."
                    )
                    message_item = ScheduleMessageItem(
                        user_id=target_user_id,
                        session_id=target_session_id,
                        mem_cube_id=mem_cube_id,
                        label=PREF_ADD_TASK_LABEL,
                        content=json.dumps(messages_list),
                        timestamp=datetime.utcnow(),
                    )
                    self.mem_scheduler.submit_messages(messages=[message_item])

        # Execute both memory processing functions in parallel
        with ContextThreadPoolExecutor(max_workers=2) as executor:
            text_future = executor.submit(process_textual_memory)
            pref_future = executor.submit(process_preference_memory)

            # Wait for both tasks to complete
            text_future.result()
            pref_future.result()

        # user profile
        if (
            (memory_content is not None)
            and self.config.enable_textual_memory
            and self.mem_cubes[mem_cube_id].text_mem
        ):
            if self.mem_cubes[mem_cube_id].config.text_mem.backend != "tree_text":
                metadata = TextualMemoryMetadata(
                    user_id=target_user_id, session_id=target_session_id, source="conversation"
                )
                self.mem_cubes[mem_cube_id].text_mem.add(
                    [TextualMemoryItem(memory=memory_content, metadata=metadata)]
                )
            else:
                messages_list = [
                    [{"role": "user", "content": memory_content}]
                ]  # for only user-str input and convert message

                memories = self.mem_reader.get_memory(
                    messages_list,
                    type="chat",
                    info={"user_id": target_user_id, "session_id": target_session_id},
                    mode="fast" if sync_mode == "async" else "fine",
                )

                mem_ids = []
                for mem in memories:
                    mem_id_list: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(mem)
                    logger.info(
                        f"Added memory user {target_user_id} to memcube {mem_cube_id}: {mem_id_list}"
                    )
                    mem_ids.extend(mem_id_list)

                # submit messages for scheduler
                if self.enable_mem_scheduler and self.mem_scheduler is not None:
                    if sync_mode == "async":
                        message_item = ScheduleMessageItem(
                            user_id=target_user_id,
                            mem_cube_id=mem_cube_id,
                            label=MEM_READ_TASK_LABEL,
                            content=json.dumps(mem_ids),
                            timestamp=datetime.utcnow(),
                        )
                        self.mem_scheduler.submit_messages(messages=[message_item])
                    else:
                        message_item = ScheduleMessageItem(
                            user_id=target_user_id,
                            mem_cube_id=mem_cube_id,
                            label=ADD_TASK_LABEL,
                            content=json.dumps(mem_ids),
                            timestamp=datetime.utcnow(),
                        )
                        self.mem_scheduler.submit_messages(messages=[message_item])

        # user doc input
        if (
            (doc_path is not None)
            and self.config.enable_textual_memory
            and self.mem_cubes[mem_cube_id].text_mem
        ):
            documents = self._get_all_documents(doc_path)
            doc_memories = self.mem_reader.get_memory(
                documents,
                type="doc",
                info={"user_id": target_user_id, "session_id": target_session_id},
            )

            mem_ids = []
            for mem in doc_memories:
                mem_id_list: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(mem)
                mem_ids.extend(mem_id_list)

            # submit messages for scheduler
            if self.enable_mem_scheduler and self.mem_scheduler is not None:
                message_item = ScheduleMessageItem(
                    user_id=target_user_id,
                    mem_cube_id=mem_cube_id,
                    label=ADD_TASK_LABEL,
                    content=json.dumps(mem_ids),
                    timestamp=datetime.utcnow(),
                )
                self.mem_scheduler.submit_messages(messages=[message_item])

        logger.info(f"Add memory to {mem_cube_id} successfully")

    def get(
        self, mem_cube_id: str, memory_id: str, user_id: str | None = None
    ) -> TextualMemoryItem | ActivationMemoryItem | ParametricMemoryItem:
        """
        Get a textual memory from a MemCube.

        Args:
            mem_cube_id (str): The identifier of the MemCube to get the memory from.
            memory_id (str): The identifier of the  memory to get.
            user_id (str, optional): The identifier of the user to get the memory from.
                If None, the default user is used.

        Returns:
            Union[TextualMemoryItem, ActivationMemoryItem, ParametricMemoryItem]: The requested memory item.
        """
        target_user_id = user_id if user_id is not None else self.user_id
        # Validate user has access to this cube
        self._validate_cube_access(target_user_id, mem_cube_id)
        if mem_cube_id is None:
            # Try to find a default cube for the user
            accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
            if not accessible_cubes:
                raise ValueError(
                    f"No accessible cubes found for user '{target_user_id}'. Please register a cube first."
                )
            mem_cube_id = accessible_cubes[0].cube_id  # TODO not only first
        else:
            self._validate_cube_access(target_user_id, mem_cube_id)

        assert mem_cube_id in self.mem_cubes, (
            f"MemCube with ID {mem_cube_id} does not exist. please regiester"
        )
        return self.mem_cubes[mem_cube_id].text_mem.get(memory_id)

    def get_all(
        self, mem_cube_id: str | None = None, user_id: str | None = None
    ) -> MOSSearchResult:
        """
        Get all textual memories from a MemCube.

        Args:
            mem_cube_id (str, optional): The identifier of the MemCube to get the memories from.
                If None, all MemCube for the user is used.
            user_id (str, optional): The identifier of the user to get the memories from.
                If None, the default user is used.

        Returns:
            MemoryResult: A dictionary containing the search results.
        """
        result: MOSSearchResult = {"para_mem": [], "act_mem": [], "text_mem": []}
        target_user_id = user_id if user_id is not None else self.user_id
        # Validate user has access to this cube
        if mem_cube_id is None:
            # Try to find a default cube for the user
            accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
            if not accessible_cubes:
                raise ValueError(
                    f"No accessible cubes found for user '{target_user_id}'. Please register a cube first."
                )
            mem_cube_id = accessible_cubes[0].cube_id  # TODO not only first
        else:
            self._validate_cube_access(target_user_id, mem_cube_id)
        if self.config.enable_textual_memory and self.mem_cubes[mem_cube_id].text_mem:
            result["text_mem"].append(
                {"cube_id": mem_cube_id, "memories": self.mem_cubes[mem_cube_id].text_mem.get_all()}
            )
        if self.config.enable_activation_memory and self.mem_cubes[mem_cube_id].act_mem:
            result["act_mem"].append(
                {"cube_id": mem_cube_id, "memories": self.mem_cubes[mem_cube_id].act_mem.get_all()}
            )
        return result

    def update(
        self,
        mem_cube_id: str,
        memory_id: str,
        text_memory_item: TextualMemoryItem | dict[str, Any],
        user_id: str | None = None,
    ) -> None:
        """
        Update a textual memory in a MemCube by text_memory_id and text_memory_id.

        Args:
            mem_cube_id (str): The identifier of the MemCube to update the memory in.
            memory_id (str): The identifier of the textual memory to update.
            text_memory_item (TextualMemoryItem | dict[str, Any]): The updated textual memory item.
        """
        assert mem_cube_id in self.mem_cubes, (
            f"MemCube with ID {mem_cube_id} does not exist. please regiester"
        )
        target_user_id = user_id if user_id is not None else self.user_id
        # Validate user has access to this cube
        self._validate_cube_access(target_user_id, mem_cube_id)
        if mem_cube_id is None:
            # Try to find a default cube for the user
            accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
            if not accessible_cubes:
                raise ValueError(
                    f"No accessible cubes found for user '{target_user_id}'. Please register a cube first."
                )
            mem_cube_id = accessible_cubes[0].cube_id  # TODO not only first
        else:
            self._validate_cube_access(target_user_id, mem_cube_id)
        if self.mem_cubes[mem_cube_id].config.text_mem.backend != "tree_text":
            self.mem_cubes[mem_cube_id].text_mem.update(memory_id, memories=text_memory_item)
            logger.info(f"MemCube {mem_cube_id} updated memory {memory_id}")
        else:
            logger.warning(
                f" {self.mem_cubes[mem_cube_id].config.text_mem.backend} does not support update memory"
            )

    def delete(self, mem_cube_id: str, memory_id: str, user_id: str | None = None) -> None:
        """
        Delete a textual memory from a MemCube by memory_id.

        Args:
            mem_cube_id (str): The identifier of the MemCube to delete the memory from.
            memory_id (str): The identifier of the  memory to delete.
        """
        assert mem_cube_id in self.mem_cubes, (
            f"MemCube with ID {mem_cube_id} does not exist. please regiester"
        )
        target_user_id = user_id if user_id is not None else self.user_id
        # Validate user has access to this cube
        self._validate_cube_access(target_user_id, mem_cube_id)
        if mem_cube_id is None:
            # Try to find a default cube for the user
            accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
            if not accessible_cubes:
                raise ValueError(
                    f"No accessible cubes found for user '{target_user_id}'. Please register a cube first."
                )
            mem_cube_id = accessible_cubes[0].cube_id  # TODO not only first
        else:
            self._validate_cube_access(target_user_id, mem_cube_id)
        self.mem_cubes[mem_cube_id].text_mem.delete(memory_id)
        logger.info(f"MemCube {mem_cube_id} deleted memory {memory_id}")

    def delete_all(self, mem_cube_id: str | None = None, user_id: str | None = None) -> None:
        """
        Delete all textual memories from a MemCube for user.

        Args:
            mem_cube_id (str): The identifier of the MemCube to delete the memories from.
        """
        assert mem_cube_id in self.mem_cubes, (
            f"MemCube with ID {mem_cube_id} does not exist. please regiester"
        )
        target_user_id = user_id if user_id is not None else self.user_id
        # Validate user has access to this cube
        self._validate_cube_access(target_user_id, mem_cube_id)
        if mem_cube_id is None:
            # Try to find a default cube for the user
            accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
            if not accessible_cubes:
                raise ValueError(
                    f"No accessible cubes found for user '{target_user_id}'. Please register a cube first."
                )
            mem_cube_id = accessible_cubes[0].cube_id  # TODO not only first
        else:
            self._validate_cube_access(target_user_id, mem_cube_id)
        self.mem_cubes[mem_cube_id].text_mem.delete_all()
        logger.info(f"MemCube {mem_cube_id} deleted all memories")

    def dump(
        self, dump_dir: str, user_id: str | None = None, mem_cube_id: str | None = None
    ) -> None:
        """Dump the MemCube to a dictionary.
        Args:
            dump_dir (str): The directory to dump the MemCube to.
            user_id (str, optional): The identifier of the user to dump the MemCube from.
                If None, the default user is used.
            mem_cube_id (str, optional): The identifier of the MemCube to dump.
                If None, the default MemCube for the user is used.
        """
        target_user_id = user_id if user_id is not None else self.user_id
        accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
        if not mem_cube_id:
            mem_cube_id = accessible_cubes[0].cube_id
        if mem_cube_id not in self.mem_cubes:
            raise ValueError(f"MemCube with ID {mem_cube_id} does not exist. please regiester")
        self.mem_cubes[mem_cube_id].dump(dump_dir)
        logger.info(f"MemCube {mem_cube_id} dumped to {dump_dir}")

    def load(
        self,
        load_dir: str,
        user_id: str | None = None,
        mem_cube_id: str | None = None,
        memory_types: list[Literal["text_mem", "act_mem", "para_mem", "pref_mem"]] | None = None,
    ) -> None:
        """Dump the MemCube to a dictionary.
        Args:
            load_dir (str): The directory to load the MemCube from.
            user_id (str, optional): The identifier of the user to load the MemCube from.
                If None, the default user is used.
            mem_cube_id (str, optional): The identifier of the MemCube to load.
                If None, the default MemCube for the user is used.
        """
        target_user_id = user_id if user_id is not None else self.user_id
        accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
        if not mem_cube_id:
            mem_cube_id = accessible_cubes[0].cube_id
        if mem_cube_id not in self.mem_cubes:
            raise ValueError(f"MemCube with ID {mem_cube_id} does not exist. please regiester")
        self.mem_cubes[mem_cube_id].load(load_dir, memory_types=memory_types)
        logger.info(f"MemCube {mem_cube_id} loaded from {load_dir}")

    def get_user_info(self) -> dict[str, Any]:
        """Get current user information including accessible cubes.
        TODO: maybe input user_id
        Returns:
            dict: User information and accessible cubes.
        """
        user = self.user_manager.get_user(self.user_id)
        if not user:
            return {}

        accessible_cubes = self.user_manager.get_user_cubes(self.user_id)

        return {
            "user_id": user.user_id,
            "user_name": user.user_name,
            "role": user.role.value if hasattr(user.role, "value") else user.role,
            "created_at": user.created_at.isoformat(),
            "accessible_cubes": [
                {
                    "cube_id": cube.cube_id,
                    "cube_name": cube.cube_name,
                    "cube_path": cube.cube_path,
                    "owner_id": cube.owner_id,
                    "is_loaded": cube.cube_id in self.mem_cubes,
                }
                for cube in accessible_cubes
            ],
        }

    def share_cube_with_user(self, cube_id: str, target_user_id: str) -> bool:
        """Share a cube with another user.

        Args:
            cube_id (str): The cube ID to share.
            target_user_id (str): The user ID to share with.

        Returns:
            bool: True if successful, False otherwise.
        """
        # Validate current user has access to this cube
        self._validate_cube_access(self.user_id, cube_id)

        # Validate target user exists
        if not self.user_manager.validate_user(target_user_id):
            raise ValueError(f"Target user '{target_user_id}' does not exist or is inactive.")

        return self.user_manager.add_user_to_cube(target_user_id, cube_id)

    def get_query_rewrite(self, query: str, user_id: str | None = None):
        """
        Rewrite user's query according the context.
        Args:
            query (str): The search query that needs rewriting.
            user_id(str, optional): The identifier of the user that the query belongs to.
                If None, the default user is used.

        Returns:
            str: query after rewriting process.
        """
        target_user_id = user_id if user_id is not None else self.user_id
        chat_history = self.chat_history_manager[target_user_id]

        dialogue = "————{}".format("\n————".join(chat_history.chat_history))
        user_prompt = QUERY_REWRITING_PROMPT.format(dialogue=dialogue, query=query)
        messages = {"role": "user", "content": user_prompt}
        rewritten_result = self.chat_llm.generate(messages=messages)
        rewritten_result = json.loads(rewritten_result)
        if rewritten_result.get("former_dialogue_related", False):
            rewritten_query = rewritten_result.get("rewritten_question")
            return rewritten_query if len(rewritten_query) > 0 else query
        return query
