import concurrent.futures
import json
import re
import traceback
import uuid

from typing import TYPE_CHECKING, Any

from memos import log
from memos.configs.mem_reader import MultiModalStructMemReaderConfig
from memos.context.context import ContextThreadPoolExecutor
from memos.mem_reader.read_multi_modal import MultiModalParser, detect_lang
from memos.mem_reader.read_multi_modal.base import _derive_key
from memos.mem_reader.read_pref_memory.process_preference_memory import process_preference_fine
from memos.mem_reader.read_skill_memory.process_skill_memory import process_skill_memory_fine
from memos.mem_reader.simple_struct import PROMPT_DICT, SimpleStructMemReader
from memos.mem_reader.utils import parse_json_result
from memos.memories.textual.item import TextualMemoryItem, TreeNodeTextualMemoryMetadata
from memos.plugins.hook_defs import H
from memos.plugins.hooks import trigger_hook, trigger_single_hook
from memos.templates.mem_reader_prompts import MEMORY_MERGE_PROMPT_EN, MEMORY_MERGE_PROMPT_ZH
from memos.templates.tool_mem_prompts import TOOL_TRAJECTORY_PROMPT_EN, TOOL_TRAJECTORY_PROMPT_ZH
from memos.types import MessagesType
from memos.utils import timed, timed_stage


if TYPE_CHECKING:
    from memos.types.general_types import UserContext


logger = log.get_logger(__name__)


class MultiModalStructMemReader(SimpleStructMemReader):
    """Multimodal implementation of MemReader that inherits from
    SimpleStructMemReader."""

    def __init__(self, config: MultiModalStructMemReaderConfig):
        """
        Initialize the MultiModalStructMemReader with configuration.

        Args:
            config: Configuration object for the reader
        """
        from memos.configs.mem_reader import SimpleStructMemReaderConfig
        from memos.llms.factory import LLMFactory

        # Extract direct_markdown_hostnames before converting to SimpleStructMemReaderConfig
        direct_markdown_hostnames = getattr(config, "direct_markdown_hostnames", None)

        # oss
        self.oss_config = getattr(config, "oss_config", None)

        # skills_dir
        self.skills_dir_config = getattr(config, "skills_dir_config", None)

        # Create config_dict excluding direct_markdown_hostnames for SimpleStructMemReaderConfig
        config_dict = config.model_dump(exclude_none=True)
        config_dict.pop("direct_markdown_hostnames", None)

        simple_config = SimpleStructMemReaderConfig(**config_dict)
        super().__init__(simple_config)

        self.memory_version_switch = getattr(config, "memory_version_switch", "off")

        # Image parser LLM (requires vision model)
        # Falls back to general_llm if not configured (general_llm itself falls back to main llm)
        self.image_parser_llm = (
            LLMFactory.from_config(config.image_parser_llm)
            if config.image_parser_llm is not None
            else self.general_llm
        )
        # Initialize MultiModalParser for routing to different parsers
        # Pass image_parser_llm for image parsing
        self.multi_modal_parser = MultiModalParser(
            embedder=self.embedder,
            llm=self.llm,
            image_parser_llm=self.image_parser_llm,
            parser=None,
            direct_markdown_hostnames=direct_markdown_hostnames,
        )

    def _embed_memory_items(self, items: list[TextualMemoryItem]) -> None:
        """Compute embeddings for a list of memory items in-place.

        Attempts a single batch call first; falls back to per-item calls if the
        batch fails.  Errors are logged but never raised so callers always
        continue normally.
        """
        valid = [w for w in items if w and w.memory]
        if not valid:
            return
        texts = [w.memory for w in valid]
        try:
            embeddings = self.embedder.embed(texts)
            for w, emb in zip(valid, embeddings, strict=True):
                w.metadata.embedding = emb
        except Exception as e:
            logger.error(f"[MultiModalStruct] Error batch computing embeddings: {e}")
            logger.warning("[EMBED_FALLBACK] batch_size=%d", len(texts))
            for w in valid:
                try:
                    w.metadata.embedding = self.embedder.embed([w.memory])[0]
                except Exception as e2:
                    logger.error(f"[MultiModalStruct] Error computing embedding for item: {e2}")

    def _split_large_memory_item(
        self, item: TextualMemoryItem, max_tokens: int
    ) -> list[TextualMemoryItem]:
        """
        Split a single memory item that exceeds max_tokens into multiple chunks.

        Args:
            item: TextualMemoryItem to split
            max_tokens: Maximum tokens per chunk

        Returns:
            List of TextualMemoryItem chunks
        """
        item_text = item.memory or ""
        if not item_text:
            return [item]

        item_tokens = self._count_tokens(item_text)
        if item_tokens <= max_tokens:
            return [item]

        # Use chunker to split the text
        try:
            chunks = self.chunker.chunk(item_text)
            split_items = []
            source_info = dict(item.metadata.info or {})
            source_internal_info = dict(item.metadata.internal_info or {})
            ingest_batch_id = str(source_internal_info.get("ingest_batch_id") or uuid.uuid4())
            chunk_total = len(chunks)

            def _create_chunk_item(chunk_idx: int, chunk):
                # Different chunkers are not fully consistent:
                # some return Chunk-like objects with `.text`, while others return raw strings.
                chunk_text = chunk.text if hasattr(chunk, "text") else chunk
                if not chunk_text or not chunk_text.strip():
                    return None
                chunk_info = {
                    "user_id": item.metadata.user_id,
                    "session_id": item.metadata.session_id,
                    **source_info,
                }
                chunk_internal_info = {
                    **source_internal_info,
                    "ingest_batch_id": ingest_batch_id,
                    "chunk_index": chunk_idx,
                    "chunk_total": chunk_total,
                }
                # Create a new memory item for each chunk, preserving original metadata
                split_item = self._make_memory_item(
                    value=chunk_text,
                    info=chunk_info,
                    memory_type=item.metadata.memory_type,
                    tags=item.metadata.tags or [],
                    key=item.metadata.key,
                    sources=item.metadata.sources or [],
                    background=item.metadata.background or "",
                    need_embed=False,
                )
                split_item.metadata.internal_info = chunk_internal_info
                return split_item

            # Use thread pool to parallel process chunks, but keep the original order
            with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
                futures = [
                    executor.submit(_create_chunk_item, chunk_idx, chunk)
                    for chunk_idx, chunk in enumerate(chunks)
                ]
                for future in futures:
                    split_item = future.result()
                    if split_item is not None:
                        split_items.append(split_item)

            return split_items if split_items else [item]
        except Exception as e:
            logger.warning(
                f"[MultiModalStruct] Failed to split large memory item: {e}. Returning original item."
            )
            return [item]

    def _concat_multi_modal_memories(
        self, all_memory_items: list[TextualMemoryItem], max_tokens=None, overlap=200
    ) -> list[TextualMemoryItem]:
        """
        Aggregates memory items using sliding window logic similar to
        `_iter_chat_windows` in simple_struct:
        1. Groups items into windows based on token count (max_tokens)
        2. Each window has overlap tokens for context continuity
        3. Aggregates items within each window into a single memory item
        4. Determines memory_type based on roles in each window
        5. Splits single large memory items that exceed max_tokens
        """
        if not all_memory_items:
            return []

        max_tokens = max_tokens or self.chat_window_max_tokens

        # Split large memory items before processing
        processed_items = []
        # control whether to parallel chunk large memory items
        parallel_chunking = True

        if parallel_chunking:
            # parallel chunk large memory items, but keep the original order
            with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
                # Create a list to hold futures with their original index
                futures = []
                for idx, item in enumerate(all_memory_items):
                    if (item.memory or "") and self._count_tokens(item.memory) > max_tokens:
                        future = executor.submit(self._split_large_memory_item, item, max_tokens)
                        futures.append(
                            (idx, future, True)
                        )  # True indicates this item needs splitting
                    else:
                        futures.append((idx, item, False))  # False indicates no splitting needed

                # Process results in original order
                temp_results = [None] * len(all_memory_items)
                for idx, future_or_item, needs_splitting in futures:
                    if needs_splitting:
                        # Wait for the future to complete and get the split items
                        split_items = future_or_item.result()
                        temp_results[idx] = split_items
                    else:
                        # No splitting needed, use the original item
                        temp_results[idx] = [future_or_item]

                # Flatten the results while preserving order
                for items in temp_results:
                    processed_items.extend(items)
        else:
            # serial chunk large memory items
            for item in all_memory_items:
                item_text = item.memory or ""
                item_tokens = self._count_tokens(item_text)
                if item_tokens > max_tokens:
                    # Split the large item into multiple chunks
                    split_items = self._split_large_memory_item(item, max_tokens)
                    processed_items.extend(split_items)
                else:
                    processed_items.append(item)

        # If only one item after processing, compute embedding and return
        if len(processed_items) == 1:
            single_item = processed_items[0]
            with timed_stage("add", "embedding", window_count=1):
                self._embed_memory_items([single_item])
            return processed_items

        windows = []
        buf_items = []
        cur_text = ""

        # Extract info from first item (all items should have same user_id, session_id)
        first_item = processed_items[0]
        info = {
            "user_id": first_item.metadata.user_id,
            "session_id": first_item.metadata.session_id,
            **(first_item.metadata.info or {}),
        }

        for _idx, item in enumerate(processed_items):
            item_text = item.memory or ""
            # Ensure line ends with newline (same format as simple_struct)
            line = item_text if item_text.endswith("\n") else f"{item_text}\n"

            # Check if adding this item would exceed max_tokens (same logic as _iter_chat_windows)
            # Note: After splitting large items, each item should be <= max_tokens,
            # but we still check to handle edge cases
            if self._count_tokens(cur_text + line) > max_tokens and cur_text:
                # Yield current window
                window = self._build_window_from_items(buf_items, info)
                if window:
                    windows.append(window)

                # Keep overlap: remove items until remaining tokens <= overlap
                # (same logic as _iter_chat_windows)
                while (
                    buf_items
                    and self._count_tokens("".join([it.memory or "" for it in buf_items])) > overlap
                ):
                    buf_items.pop(0)
                # Recalculate cur_text from remaining items
                cur_text = "".join([it.memory or "" for it in buf_items])

            # Add item to current window
            buf_items.append(item)
            # Recalculate cur_text from all items in buffer (same as _iter_chat_windows)
            cur_text = "".join([it.memory or "" for it in buf_items])

        # Yield final window if any items remain
        if buf_items:
            window = self._build_window_from_items(buf_items, info)
            if window:
                windows.append(window)

        # Batch compute embeddings for all windows
        with timed_stage("add", "embedding", window_count=len(windows)):
            self._embed_memory_items(windows)

        return windows

    def _build_window_from_items(
        self, items: list[TextualMemoryItem], info: dict[str, Any]
    ) -> TextualMemoryItem | None:
        """
        Build a single memory item from a window of items (similar to _build_fast_node).

        Args:
            items: List of TextualMemoryItem objects in the window
            info: Dictionary containing user_id and session_id

        Returns:
            Aggregated TextualMemoryItem or None if no valid content
        """
        if not items:
            return None

        # Collect all memory texts and sources
        memory_texts = []
        all_sources = []
        roles = set()
        aggregated_file_ids: list[str] = []
        ingest_batch_ids: set[str] = set()

        for item in items:
            if item.memory:
                memory_texts.append(item.memory)

            # Collect sources and extract roles
            item_sources = item.metadata.sources or []
            if not isinstance(item_sources, list):
                item_sources = [item_sources]

            for source in item_sources:
                # Add source to all_sources
                all_sources.append(source)

                # Extract role from source
                if hasattr(source, "role") and source.role:
                    roles.add(source.role)
                elif isinstance(source, dict) and source.get("role"):
                    roles.add(source.get("role"))

            # Aggregate file_ids from metadata
            metadata = getattr(item, "metadata", None)
            if metadata is not None:
                item_file_ids = getattr(metadata, "file_ids", None)
                if isinstance(item_file_ids, list):
                    for fid in item_file_ids:
                        if fid and fid not in aggregated_file_ids:
                            aggregated_file_ids.append(fid)
                item_internal_info = getattr(metadata, "internal_info", None)
                if isinstance(item_internal_info, dict):
                    batch_id = item_internal_info.get("ingest_batch_id")
                    if batch_id:
                        ingest_batch_ids.add(str(batch_id))

        # Determine memory_type based on roles (same logic as simple_struct)
        # UserMemory if only user role, else LongTermMemory
        memory_type = "UserMemory" if roles == {"user"} else "LongTermMemory"

        # Merge all memory texts (preserve the format from parser)
        merged_text = "".join(memory_texts) if memory_texts else ""

        if not merged_text.strip():
            # If no text content, return None
            return None

        # Create aggregated memory item without embedding (will be computed in batch later)
        extra_kwargs: dict[str, Any] = {}
        if aggregated_file_ids:
            extra_kwargs["file_ids"] = aggregated_file_ids

        # Propagate manager_user_id and project_id from constituent items
        for item in items:
            metadata = getattr(item, "metadata", None)
            if metadata is not None:
                if not extra_kwargs.get("manager_user_id"):
                    mid = getattr(metadata, "manager_user_id", None)
                    if mid:
                        extra_kwargs["manager_user_id"] = mid
                if not extra_kwargs.get("project_id"):
                    pid = getattr(metadata, "project_id", None)
                    if pid:
                        extra_kwargs["project_id"] = pid

        # Extract info fields
        info_ = info.copy()
        user_id = info_.pop("user_id", "")
        session_id = info_.pop("session_id", "")
        # Create memory item without embedding (set to None, will be filled in batch)
        aggregated_item = TextualMemoryItem(
            memory=merged_text,
            metadata=TreeNodeTextualMemoryMetadata(
                user_id=user_id,
                session_id=session_id,
                memory_type=memory_type,
                status="activated",
                tags=["mode:fast"],
                key=_derive_key(merged_text),
                embedding=None,  # Will be computed in batch
                usage=[],
                sources=all_sources,
                background="",
                confidence=0.99,
                type="fact",
                info=info_,
                **extra_kwargs,
            ),
        )
        if len(ingest_batch_ids) == 1:
            aggregated_item.metadata.internal_info = {
                "ingest_batch_id": next(iter(ingest_batch_ids))
            }

        return aggregated_item

    def _get_llm_response(
        self,
        mem_str: str,
        custom_tags: list[str] | None = None,
        sources: list | None = None,
        prompt_type: str = "chat",
    ) -> dict:
        """
        Override parent method to improve language detection by using actual text content
        from sources instead of JSON-structured memory string.

        Args:
            mem_str: Memory string (may contain JSON structures)
            custom_tags: Optional custom tags
            sources: Optional list of SourceMessage objects to extract text content from
            prompt_type: Type of prompt to use ("chat" or "doc")

        Returns:
            LLM response dictionary
        """
        # Determine language: prioritize lang from sources (set in fast mode),
        # fallback to detecting from mem_str if sources don't have lang
        lang = None

        # First, try to get lang from sources (fast mode already set this)
        if sources:
            for source in sources:
                if hasattr(source, "lang") and source.lang:
                    lang = source.lang
                    break
                elif isinstance(source, dict) and source.get("lang"):
                    lang = source.get("lang")
                    break

        # Fallback: detect language from mem_str if no lang from sources
        if lang is None:
            lang = detect_lang(mem_str)

        # Select prompt template based on prompt_type
        if prompt_type == "doc":
            template = PROMPT_DICT["doc"][lang]
            examples = ""  # doc prompts don't have examples
            prompt = template.replace("{chunk_text}", mem_str)
        elif prompt_type == "general_string":
            template = PROMPT_DICT["general_string"][lang]
            examples = ""
            prompt = template.replace("{chunk_text}", mem_str)
        else:
            template = PROMPT_DICT["chat"][lang]
            examples = PROMPT_DICT["chat"][f"{lang}_example"]
            prompt = template.replace("${conversation}", mem_str)

        custom_tags_prompt = (
            PROMPT_DICT["custom_tags"][lang].replace("{custom_tags}", str(custom_tags))
            if custom_tags
            else ""
        )

        # Replace custom_tags_prompt placeholder (different for doc vs chat)
        if prompt_type in ["doc", "general_string"]:
            prompt = prompt.replace("{custom_tags_prompt}", custom_tags_prompt)
        else:
            prompt = prompt.replace("${custom_tags_prompt}", custom_tags_prompt)

        if self.config.remove_prompt_example and examples:
            prompt = prompt.replace(examples, "")

        logger.info(f"[MultiModalParser] Process String Fine Prompt: {prompt}")

        messages = [{"role": "user", "content": prompt}]
        try:
            response_text = self.llm.generate(messages)
            response_json = parse_json_result(response_text)
        except Exception as e:
            logger.error(f"[LLM] Exception during chat generation: {e}")
            response_json = {
                "memory list": [
                    {
                        "key": mem_str[:10],
                        "memory_type": "UserMemory",
                        "value": mem_str,
                        "tags": [],
                    }
                ],
                "summary": mem_str,
            }
        logger.info(f"[MultiModalFine] Task {messages}, Result {response_json}")
        return response_json

    def _determine_prompt_type(self, sources: list) -> str:
        """
        Determine prompt type based on sources.
        """
        if not sources:
            return "chat"
        prompt_type = "general_string"
        for source in sources:
            source_role = None
            if hasattr(source, "role"):
                source_role = source.role
            elif isinstance(source, dict):
                source_role = source.get("role")
            if source_role in {"user", "assistant", "system", "tool"}:
                prompt_type = "chat"
                if hasattr(source, "type"):
                    source_type = source.type
                    if source_type == "file":
                        prompt_type = "doc"
        return prompt_type

    def _get_maybe_merged_memory(
        self,
        extracted_memory_dict: dict,
        mem_text: str,
        sources: list,
        **kwargs,
    ) -> dict:
        # TODO: delete this function
        """
        Check if extracted memory should be merged with similar existing memories.
        If merge is needed, return merged memory dict with merged_from field.
        Otherwise, return original memory dict.

        Args:
            extracted_memory_dict: The extracted memory dict from LLM response
            mem_text: The memory text content
            sources: Source messages for language detection
            **kwargs: Additional parameters (merge_similarity_threshold, etc.)

        Returns:
            Memory dict (possibly merged) with merged_from field if merged
        """
        return extracted_memory_dict

    def _merge_memories_with_llm(
        self,
        new_memory: TextualMemoryItem,
        similar_memories: list[dict],
        lang: str = "en",
    ) -> dict | None:
        """
        Use LLM to merge new memory with similar existing memories.

        Args:
            new_memory: The newly extracted memory item
            similar_memories: List of similar memories from graph_db (with id and memory fields)
            lang: Language code ("en" or "zh")

        Returns:
            Merged memory dict with merged_from field, or None if no merge needed
        """
        if not similar_memories:
            return None

        # Build merge prompt using template
        similar_memories_text = "\n".join(
            [f"[{mem['id']}]: {mem['memory']}" for mem in similar_memories]
        )

        merge_prompt_template = MEMORY_MERGE_PROMPT_ZH if lang == "zh" else MEMORY_MERGE_PROMPT_EN
        merge_prompt = merge_prompt_template.format(
            new_memory=new_memory.memory,
            similar_memories=similar_memories_text,
        )

        try:
            # Use general_llm for memory merge (not fine-tuned for this task)
            response_text = self.general_llm.generate([{"role": "user", "content": merge_prompt}])
            merge_result = parse_json_result(response_text)

            if merge_result.get("should_merge", False):
                return {
                    "value": merge_result.get("value", new_memory.memory),
                    "merged_from": merge_result.get(
                        "merged_from", [mem["id"] for mem in similar_memories]
                    ),
                }
        except Exception as e:
            logger.error(f"[MultiModalFine] Error in merge LLM call: {e}")

        return None

    @timed
    def _process_string_fine(
        self,
        fast_memory_items: list[TextualMemoryItem],
        info: dict[str, Any],
        custom_tags: list[str] | None = None,
        **kwargs,
    ) -> list[TextualMemoryItem]:
        """
        Process fast mode memory items through LLM to generate fine mode memories.
        Where fast_memory_items are raw chunk memory items, not the final memory items.
        """
        if not fast_memory_items:
            return []

        def _process_one_item(
            fast_item: TextualMemoryItem, chunk_idx: int, total_chunks: int
        ) -> list[TextualMemoryItem]:
            """Process a single fast memory item and return a list of fine items."""
            fine_items: list[TextualMemoryItem] = []

            # Extract memory text (string content)
            mem_str = fast_item.memory or ""
            if not mem_str.strip():
                return fine_items

            sources = fast_item.metadata.sources or []
            if not isinstance(sources, list):
                sources = [sources]

            # Extract file_ids from fast item metadata for propagation
            metadata = getattr(fast_item, "metadata", None)
            file_ids = getattr(metadata, "file_ids", None) if metadata is not None else None
            file_ids = [fid for fid in file_ids if fid] if isinstance(file_ids, list) else []

            # Build per-item info copy and kwargs for _make_memory_item
            info_per_item = info.copy()
            if file_ids and "file_id" not in info_per_item:
                info_per_item["file_id"] = file_ids[0]
            extra_kwargs: dict[str, Any] = {}
            if file_ids:
                extra_kwargs["file_ids"] = file_ids

            # Extract manager_user_id and project_id from user_context
            user_context: UserContext | None = kwargs.get("user_context")
            if user_context:
                extra_kwargs["manager_user_id"] = user_context.manager_user_id
                extra_kwargs["project_id"] = user_context.project_id

            # Determine prompt type based on sources
            prompt_type = self._determine_prompt_type(sources)

            # ========== Stage 0: Memory version async extraction/update pipeline ==========
            if getattr(self, "memory_version_switch", "off") == "on":
                try:
                    user_name = kwargs.get("user_name")
                    should_use_version_pipeline = trigger_single_hook(
                        H.MEMORY_VERSION_PREPARE_UPDATES,
                        item=fast_item,
                        user_name=user_name,
                        judge_llm=self.general_llm,
                    )
                    if should_use_version_pipeline:
                        lang = detect_lang(kwargs.get("chat_history") or mem_str)
                        custom_tags_prompt_template = PROMPT_DICT["custom_tags"][lang]
                        new_items = trigger_single_hook(
                            H.MEMORY_VERSION_APPLY_UPDATES,
                            item=fast_item,
                            user_name=user_name,
                            version_llm=self.qwen_llm,
                            merge_llm=self.general_llm,
                            custom_tags=custom_tags,
                            custom_tags_prompt_template=custom_tags_prompt_template,
                            timeout_sec=30,
                        )
                        return new_items
                except RuntimeError as ex:
                    logger.warning(f"[MultiModalFine] Memory version hook unavailable: {ex}")
                except Exception as ex:
                    logger.warning(f"[MultiModalFine] Fine memory version pipeline failed: {ex}")

            # ========== Stage 1: Normal extraction (without reference) ==========
            try:
                resp = self._get_llm_response(mem_str, custom_tags, sources, prompt_type)
            except Exception as e:
                logger.error(f"[MultiModalFine] Error calling LLM: {e}")
                return fine_items

            if resp.get("memory list", []):
                for m in resp.get("memory list", []):
                    try:
                        m_maybe_merged = m
                        if getattr(self, "memory_version_switch", "off") != "on":
                            m_maybe_merged = self._get_maybe_merged_memory(
                                extracted_memory_dict=m,
                                mem_text=m.get("value", ""),
                                sources=sources,
                                original_query=mem_str,
                                **kwargs,
                            )
                        # Normalize memory_type (same as simple_struct)
                        memory_type = (
                            m_maybe_merged.get("memory_type", "LongTermMemory")
                            .replace("长期记忆", "LongTermMemory")
                            .replace("用户记忆", "UserMemory")
                            .replace("pref", "UserMemory")
                        )
                        node = self._make_memory_item(
                            value=m_maybe_merged.get("value", ""),
                            info=info_per_item,
                            memory_type=memory_type,
                            tags=m_maybe_merged.get("tags", []),
                            key=m_maybe_merged.get("key", ""),
                            sources=sources,  # Preserve sources from fast item
                            background=resp.get("summary", ""),
                            **extra_kwargs,
                        )
                        if (
                            getattr(self, "memory_version_switch", "off") != "on"
                            and "merged_from" in m_maybe_merged
                        ):
                            node.metadata.info = node.metadata.info or {}
                            node.metadata.info["merged_from"] = m_maybe_merged["merged_from"]
                        fine_items.append(node)
                    except Exception as e:
                        logger.error(f"[MultiModalFine] parse error: {e}")
            elif resp.get("value") and resp.get("key"):
                try:
                    # Check and merge with similar memories if needed
                    resp_maybe_merged = resp
                    if getattr(self, "memory_version_switch", "off") != "on":
                        resp_maybe_merged = self._get_maybe_merged_memory(
                            extracted_memory_dict=resp,
                            mem_text=resp.get("value", "").strip(),
                            sources=sources,
                            original_query=mem_str,
                            **kwargs,
                        )
                    node = self._make_memory_item(
                        value=resp_maybe_merged.get("value", "").strip(),
                        info=info_per_item,
                        memory_type="LongTermMemory",
                        tags=resp_maybe_merged.get("tags", []),
                        key=resp_maybe_merged.get("key", None),
                        sources=sources,  # Preserve sources from fast item
                        background=resp.get("summary", ""),
                        **extra_kwargs,
                    )
                    if (
                        getattr(self, "memory_version_switch", "off") != "on"
                        and "merged_from" in resp_maybe_merged
                    ):
                        node.metadata.info = node.metadata.info or {}
                        node.metadata.info["merged_from"] = resp_maybe_merged["merged_from"]
                    fine_items.append(node)
                except Exception as e:
                    logger.error(f"[MultiModalFine] parse error: {e}")

            # save rawfile node
            if self.save_rawfile and prompt_type == "doc" and len(fine_items) > 0:
                rawfile_chunk = mem_str
                file_info = fine_items[0].metadata.sources[0].file_info
                source = self.multi_modal_parser.file_content_parser.create_source(
                    message={"file": file_info},
                    info=info_per_item,
                    chunk_index=chunk_idx,
                    chunk_total=total_chunks,
                    chunk_content="",
                )
                rawfile_node = self._make_memory_item(
                    value=rawfile_chunk,
                    info=info_per_item,
                    memory_type="RawFileMemory",
                    tags=[
                        "mode:fine",
                        "multimodal:file",
                        f"chunk:{chunk_idx + 1}/{total_chunks}",
                    ],
                    sources=[source],
                )
                rawfile_node.metadata.summary_ids = [mem_node.id for mem_node in fine_items]
                fine_items.append(rawfile_node)
            enriched_items = trigger_hook(
                H.MEMORY_ITEMS_AFTER_FINE_EXTRACT,
                items=fine_items,
                user_context=kwargs.get("user_context"),
                mem_reader=self,
                extract_mode="fine",
            )
            fine_items = enriched_items if enriched_items is not None else fine_items
            return fine_items

        fine_memory_items: list[TextualMemoryItem] = []
        total_chunks_len = len(fast_memory_items)

        with ContextThreadPoolExecutor(max_workers=30) as executor:
            futures = [
                executor.submit(_process_one_item, item, idx, total_chunks_len)
                for idx, item in enumerate[TextualMemoryItem](fast_memory_items)
            ]

            for future in concurrent.futures.as_completed(futures):
                try:
                    result = future.result()
                    if result:
                        fine_memory_items.extend(result)
                except Exception as e:
                    logger.error(f"[MultiModalFine] worker error: {e} {traceback.format_exc()}")

        # related preceding and following rawfilememories
        fine_memory_items = self._relate_preceding_following_rawfile_memories(fine_memory_items)
        return fine_memory_items

    def _relate_preceding_following_rawfile_memories(
        self, fine_memory_items: list[TextualMemoryItem]
    ) -> list[TextualMemoryItem]:
        """
        Relate RawFileMemory items to each other by setting preceding_id and following_id.
        """
        # Filter RawFileMemory items and track their original positions
        rawfile_items_with_pos = []
        for idx, item in enumerate[TextualMemoryItem](fine_memory_items):
            if (
                hasattr(item.metadata, "memory_type")
                and item.metadata.memory_type == "RawFileMemory"
            ):
                rawfile_items_with_pos.append((idx, item))

        if len(rawfile_items_with_pos) <= 1:
            return fine_memory_items

        def get_chunk_idx(item_with_pos) -> int:
            """Extract chunk_idx from item's source metadata."""
            _, item = item_with_pos
            if item.metadata.sources and len(item.metadata.sources) > 0:
                source = item.metadata.sources[0]
                # Handle both SourceMessage object and dict
                if isinstance(source, dict):
                    file_info = source.get("file_info")
                    if file_info and isinstance(file_info, dict):
                        chunk_idx = file_info.get("chunk_index")
                        if chunk_idx is not None:
                            return chunk_idx
                else:
                    # SourceMessage object
                    file_info = getattr(source, "file_info", None)
                    if file_info and isinstance(file_info, dict):
                        chunk_idx = file_info.get("chunk_index")
                        if chunk_idx is not None:
                            return chunk_idx
            return float("inf")

        # Sort items by chunk_index
        sorted_rawfile_items_with_pos = sorted(rawfile_items_with_pos, key=get_chunk_idx)

        # Relate adjacent items
        for i in range(len(sorted_rawfile_items_with_pos) - 1):
            _, current_item = sorted_rawfile_items_with_pos[i]
            _, next_item = sorted_rawfile_items_with_pos[i + 1]
            current_item.metadata.following_id = next_item.id
            next_item.metadata.preceding_id = current_item.id

        # Replace sorted items back to original positions in fine_memory_items
        for orig_idx, item in sorted_rawfile_items_with_pos:
            fine_memory_items[orig_idx] = item

        return fine_memory_items

    def _get_llm_tool_trajectory_response(self, mem_str: str) -> dict:
        """
        Generete tool trajectory experience item by llm.
        Uses general_llm as this task is not fine-tuned for the main model.
        """
        try:
            lang = detect_lang(mem_str)
            template = TOOL_TRAJECTORY_PROMPT_ZH if lang == "zh" else TOOL_TRAJECTORY_PROMPT_EN
            prompt = template.replace("{messages}", mem_str)
            # Use general_llm for tool trajectory (not fine-tuned for this task)
            rsp = self.general_llm.generate([{"role": "user", "content": prompt}])
            rsp = rsp.replace("```json", "").replace("```", "")
            return json.loads(rsp)
        except Exception as e:
            logger.error(f"[MultiModalFine] Error calling LLM for tool trajectory: {e}")
            return []

    @timed
    def _process_tool_trajectory_fine(
        self, fast_memory_items: list[TextualMemoryItem], info: dict[str, Any], **kwargs
    ) -> list[TextualMemoryItem]:
        """
        Process tool trajectory memory items through LLM to generate fine mode memories.
        """
        if not fast_memory_items:
            return []

        fine_memory_items = []

        # Extract manager_user_id and project_id from user_context
        user_context: UserContext | None = kwargs.get("user_context")
        manager_user_id = user_context.manager_user_id if user_context else None
        project_id = user_context.project_id if user_context else None

        for fast_item in fast_memory_items:
            sources = fast_item.metadata.sources or []
            if not isinstance(sources, list):
                sources = [sources]

            # Extract memory text (string content)
            mem_str = fast_item.memory or ""
            if not mem_str.strip() or (
                "tool:" not in mem_str
                and "[tool_calls]:" not in mem_str
                and not re.search(r"<tool_schema>.*?</tool_schema>", mem_str, re.DOTALL)
            ):
                continue
            try:
                resp = self._get_llm_tool_trajectory_response(mem_str)
            except Exception as e:
                logger.error(f"[MultiModalFine] Error calling LLM for tool trajectory: {e}")
                continue
            for m in resp:
                try:
                    # Normalize memory_type (same as simple_struct)
                    memory_type = "ToolTrajectoryMemory"

                    node = self._make_memory_item(
                        value=m.get("trajectory", ""),
                        info=info,
                        memory_type=memory_type,
                        correctness=m.get("correctness", ""),
                        experience=m.get("experience", ""),
                        tool_used_status=m.get("tool_used_status", []),
                        manager_user_id=manager_user_id,
                        project_id=project_id,
                        sources=sources,
                    )
                    fine_memory_items.append(node)
                except Exception as e:
                    logger.error(f"[MultiModalFine] parse error for tool trajectory: {e}")

        return fine_memory_items

    @timed
    def _process_multi_modal_data(
        self, scene_data_info: MessagesType, info, mode: str = "fine", **kwargs
    ) -> list[TextualMemoryItem]:
        """
        Process multimodal data using MultiModalParser.

        Args:
            scene_data_info: MessagesType input
            info: Dictionary containing user_id and session_id
            mode: mem-reader mode, fast for quick process while fine for
            better understanding via calling llm
            **kwargs: Additional parameters (mode, etc.)
        """
        # Pop custom_tags from info (same as simple_struct.py)
        # must pop here, avoid add to info, only used in sync fine mode
        custom_tags = info.pop("custom_tags", None) if isinstance(info, dict) else None

        # Stage: parse — parallel message parsing + sliding-window aggregation
        with timed_stage("add", "parse") as ts_parse:
            if isinstance(scene_data_info, list):
                expanded_messages = self._expand_multimodal_messages(scene_data_info)
                ts_parse.set(msg_count=len(expanded_messages))

                all_memory_items = []
                with ContextThreadPoolExecutor(max_workers=30) as executor:
                    futures = [
                        executor.submit(
                            self.multi_modal_parser.parse,
                            msg,
                            info,
                            mode="fast",
                            need_emb=False,
                            **kwargs,
                        )
                        for msg in expanded_messages
                    ]
                    for future in futures:
                        try:
                            items = future.result()
                            all_memory_items.extend(items)
                        except Exception as e:
                            logger.error(f"[MultiModalFine] Error in parallel parsing: {e}")
            else:
                ts_parse.set(msg_count=1)
                all_memory_items = self.multi_modal_parser.parse(
                    scene_data_info, info, mode="fast", need_emb=False, **kwargs
                )

            fast_memory_items = self._concat_multi_modal_memories(all_memory_items)
            ts_parse.set(window_count=len(fast_memory_items))

        if mode == "fast":
            return fast_memory_items

        # Stage: llm_extract — fine mode 4-way parallel LLM + per-source serial
        is_upload_skill = kwargs.pop("is_upload_skill", False)
        non_file_url_fast_items = [
            item for item in fast_memory_items if not self._is_file_url_only_item(item)
        ]

        with timed_stage("add", "llm_extract") as ts_llm:
            fine_memory_items = []

            with ContextThreadPoolExecutor(max_workers=4) as executor:
                future_string = executor.submit(
                    self._process_string_fine, non_file_url_fast_items, info, custom_tags, **kwargs
                )
                future_tool = executor.submit(
                    self._process_tool_trajectory_fine, non_file_url_fast_items, info, **kwargs
                )
                future_skill = executor.submit(
                    process_skill_memory_fine,
                    fast_memory_items=fast_memory_items
                    if is_upload_skill
                    else non_file_url_fast_items,
                    info=info,
                    searcher=self.searcher,
                    graph_db=self.graph_db,
                    llm=self.general_llm,
                    embedder=self.embedder,
                    oss_config=self.oss_config,
                    skills_dir_config=self.skills_dir_config,
                    is_upload_skill=is_upload_skill,
                    **kwargs,
                )
                future_pref = executor.submit(
                    process_preference_fine,
                    non_file_url_fast_items,
                    info,
                    getattr(self, "preference_extractor_llm", self.general_llm),
                    self.embedder,
                    **kwargs,
                )

                fine_memory_items_string_parser = future_string.result()
                fine_memory_items_tool_trajectory_parser = future_tool.result()
                fine_memory_items_skill_memory_parser = future_skill.result()
                fine_memory_items_pref_parser = future_pref.result()

            fine_memory_items.extend(fine_memory_items_string_parser)
            fine_memory_items.extend(fine_memory_items_tool_trajectory_parser)
            fine_memory_items.extend(fine_memory_items_skill_memory_parser)
            fine_memory_items.extend(fine_memory_items_pref_parser)

            # Part B: per-source serial processing
            if is_upload_skill:
                # (skip for upload skill to avoid zip being parsed)
                return fine_memory_items

            with timed_stage("add", "per_source") as ts_ps:
                for fast_item in fast_memory_items:
                    sources = fast_item.metadata.sources
                    for source in sources:
                        lang = getattr(source, "lang", "en")
                        items = self.multi_modal_parser.process_transfer(
                            source,
                            context_items=[fast_item],
                            custom_tags=custom_tags,
                            info=info,
                            lang=lang,
                            user_context=kwargs.get("user_context"),
                        )
                        fine_memory_items.extend(items)

            ts_llm.set(fine_memory_count=len(fine_memory_items), per_source_ms=ts_ps.duration_ms)

        return fine_memory_items

    @timed
    def _process_transfer_multi_modal_data(
        self, raw_nodes: list[TextualMemoryItem], custom_tags: list[str] | None = None, **kwargs
    ) -> list[TextualMemoryItem]:
        """
        Process transfer for multimodal data.

        Each source is processed independently by its corresponding parser,
        which knows how to rebuild the original message and parse it in fine mode.
        """
        if not raw_nodes:
            logger.warning("[MultiModalStruct] No raw nodes found.")
            return []

        is_upload_skill = kwargs.pop("is_upload_skill", False)

        # Extract info from raw_nodes (same as simple_struct.py)
        info = {
            "user_id": raw_nodes[0].metadata.user_id,
            "session_id": raw_nodes[0].metadata.session_id,
            **(raw_nodes[0].metadata.info or {}),
        }

        # Filter out file-URL-only items for Part A fine processing (same as _process_multi_modal_data)
        non_file_url_nodes = [node for node in raw_nodes if not self._is_file_url_only_item(node)]

        fine_memory_items = []
        # Part A: call llm in parallel using thread pool
        with ContextThreadPoolExecutor(max_workers=4) as executor:
            future_string = executor.submit(
                self._process_string_fine, non_file_url_nodes, info, custom_tags, **kwargs
            )
            future_tool = executor.submit(
                self._process_tool_trajectory_fine, non_file_url_nodes, info, **kwargs
            )
            future_skill = executor.submit(
                process_skill_memory_fine,
                raw_nodes if is_upload_skill else non_file_url_nodes,
                info,
                searcher=self.searcher,
                llm=self.general_llm,
                embedder=self.embedder,
                graph_db=self.graph_db,
                oss_config=self.oss_config,
                skills_dir_config=self.skills_dir_config,
                is_upload_skill=is_upload_skill,
                **kwargs,
            )
            # Add preference memory extraction
            future_pref = executor.submit(
                process_preference_fine,
                non_file_url_nodes,
                info,
                getattr(self, "preference_extractor_llm", self.general_llm),
                self.embedder,
                **kwargs,
            )

            # Collect results
            fine_memory_items_string_parser = future_string.result()
            fine_memory_items_tool_trajectory_parser = future_tool.result()
            fine_memory_items_skill_memory_parser = future_skill.result()
            fine_memory_items_pref_parser = future_pref.result()

        fine_memory_items.extend(fine_memory_items_string_parser)
        fine_memory_items.extend(fine_memory_items_tool_trajectory_parser)
        fine_memory_items.extend(fine_memory_items_skill_memory_parser)
        fine_memory_items.extend(fine_memory_items_pref_parser)

        # Part B: get fine multimodal items
        if is_upload_skill:
            # (skip for upload skill to avoid zip being parsed)
            return fine_memory_items
        for raw_node in raw_nodes:
            sources = raw_node.metadata.sources
            for source in sources:
                lang = getattr(source, "lang", "en")
                items = self.multi_modal_parser.process_transfer(
                    source,
                    context_items=[raw_node],
                    info=info,
                    custom_tags=custom_tags,
                    lang=lang,
                    user_context=kwargs.get("user_context"),
                )
                fine_memory_items.extend(items)
        return fine_memory_items

    @staticmethod
    def _expand_multimodal_messages(messages: list) -> list:
        """
        Expand messages whose ``content`` is a list into individual
        sub-messages so that each modality is routed to its specialised
        parser during fast-mode parsing.

        For a message like::

            {
                "content": [
                    {"type": "text", "text": "Analyze this file"},
                    {"type": "file", "file": {"file_data": "https://...", ...}},
                    {"type": "image_url", "image_url": {"url": "https://..."}},
                ],
                "role": "user",
                "chat_time": "03:14 PM on 13 March, 2026",
            }

        The result will be::

            [
                {"content": "Analyze this file", "role": "user", "chat_time": "..."},
                {"type": "file", "file": {"file_data": "https://...", ...}},
                {"type": "image_url", "image_url": {"url": "https://..."}},
            ]

        Messages whose ``content`` is already a plain string (or that are
        not dicts) are passed through unchanged.
        """
        expanded: list = []
        for msg in messages:
            if not isinstance(msg, dict):
                expanded.append(msg)
                continue

            content = msg.get("content")
            if not isinstance(content, list):
                expanded.append(msg)
                continue

            # ---- content is a list: split by modality ----
            text_parts: list[str] = []
            for part in content:
                if not isinstance(part, dict):
                    text_parts.append(str(part))
                    continue

                part_type = part.get("type", "")
                if part_type == "text":
                    text_parts.append(part.get("text", ""))
                elif part_type in ("file", "image", "image_url"):
                    # Extract as a standalone message for its specialised parser
                    expanded.append(part)
                else:
                    text_parts.append(f"[{part_type}]")

            # Reconstruct a text-only version of the original message
            # (preserving role, chat_time, message_id, etc.)
            text_content = "\n".join(t for t in text_parts if t.strip())
            if text_content.strip():
                text_msg = {k: v for k, v in msg.items() if k != "content"}
                text_msg["content"] = text_content
                expanded.append(text_msg)

        return expanded

    @staticmethod
    def _is_file_url_only_item(item: TextualMemoryItem) -> bool:
        """
        Check if a fast memory item contains only file-URL sources.
        Args:
            item: TextualMemoryItem to check

        Returns:
            True if all sources are file-type with URL info (metadata only)
        """
        sources = item.metadata.sources or []
        if not sources:
            return False
        return all(
            getattr(s, "type", None) == "file" and getattr(s, "file_info", None) for s in sources
        )

    def get_scene_data_info(self, scene_data: list, type: str) -> list[list[Any]]:
        """
        Convert normalized MessagesType scenes into scene data info.
        For MultiModalStructMemReader, this is a simplified version that returns the scenes as-is.

        Args:
            scene_data: List of MessagesType scenes
            type: Type of scene_data: ['doc', 'chat']

        Returns:
            List of scene data info
        """
        # TODO: split messages
        return scene_data

    def _read_memory(
        self,
        messages: list[MessagesType],
        type: str,
        info: dict[str, Any],
        mode: str = "fine",
        **kwargs,
    ) -> list[list[TextualMemoryItem]]:
        list_scene_data_info = self.get_scene_data_info(messages, type)

        memory_list = []
        # Process Q&A pairs concurrently with context propagation
        with ContextThreadPoolExecutor() as executor:
            futures = [
                executor.submit(
                    self._process_multi_modal_data, scene_data_info, info, mode=mode, **kwargs
                )
                for scene_data_info in list_scene_data_info
            ]
            for future in concurrent.futures.as_completed(futures):
                try:
                    res_memory = future.result()
                    if res_memory is not None:
                        memory_list.append(res_memory)
                except Exception as e:
                    logger.error(f"Task failed with exception: {e}")
                    logger.error(traceback.format_exc())
        return memory_list

    def fine_transfer_simple_mem(
        self,
        input_memories: list[TextualMemoryItem],
        type: str,
        custom_tags: list[str] | None = None,
        **kwargs,
    ) -> list[list[TextualMemoryItem]]:
        if not input_memories:
            return []

        # Process Q&A pairs concurrently with context propagation
        memory_list = self._process_transfer_multi_modal_data(input_memories, custom_tags, **kwargs)

        return [memory_list]
