from __future__ import annotations

from concurrent.futures import as_completed
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

from memos.context.context import ContextThreadPoolExecutor
from memos.multi_mem_cube.views import MemCubeView
from memos.utils import timed_stage


if TYPE_CHECKING:
    from memos.api.product_models import APIADDRequest, APIFeedbackRequest, APISearchRequest
    from memos.multi_mem_cube.single_cube import SingleCubeView


@dataclass
class CompositeCubeView(MemCubeView):
    """
    A composite view over multiple logical cubes.

    For now (fast mode), it simply fan-out writes to all cubes;
    later we can add smarter routing / slow mode here.
    """

    cube_views: list[SingleCubeView]
    logger: Any

    def add_memories(self, add_req: APIADDRequest) -> list[dict[str, Any]]:
        all_results: list[dict[str, Any]] = []
        cube_count = len(self.cube_views)

        with timed_stage("add", "multi_cube", cube_count=cube_count):
            for idx, view in enumerate(self.cube_views):
                self.logger.info(
                    "[CompositeCubeView] fan-out add to cube=%s (%d/%d)",
                    view.cube_id,
                    idx + 1,
                    cube_count,
                )
                results = view.add_memories(add_req)
                all_results.extend(results)

        return all_results

    def search_memories(self, search_req: APISearchRequest) -> dict[str, Any]:
        # aggregated MOSSearchResult
        merged_results: dict[str, Any] = {
            "text_mem": [],
            "act_mem": [],
            "para_mem": [],
            "pref_mem": [],
            "pref_note": "",
            "tool_mem": [],
            "skill_mem": [],
        }

        def _search_single_cube(view: SingleCubeView) -> dict[str, Any]:
            self.logger.info(f"[CompositeCubeView] fan-out search to cube={view.cube_id}")
            return view.search_memories(search_req)

        # parallel search for each cube
        with ContextThreadPoolExecutor(max_workers=2) as executor:
            future_to_view = {
                executor.submit(_search_single_cube, view): view for view in self.cube_views
            }

            for future in as_completed(future_to_view):
                cube_result = future.result()
                merged_results["text_mem"].extend(cube_result.get("text_mem", []))
                merged_results["act_mem"].extend(cube_result.get("act_mem", []))
                merged_results["para_mem"].extend(cube_result.get("para_mem", []))
                merged_results["pref_mem"].extend(cube_result.get("pref_mem", []))
                merged_results["tool_mem"].extend(cube_result.get("tool_mem", []))
                merged_results["skill_mem"].extend(cube_result.get("skill_mem", []))
                note = cube_result.get("pref_note")
                if note:
                    if merged_results["pref_note"]:
                        merged_results["pref_note"] += " | " + note
                    else:
                        merged_results["pref_note"] = note

        return merged_results

    def feedback_memories(self, feedback_req: APIFeedbackRequest) -> list[dict[str, Any]]:
        all_results: list[dict[str, Any]] = []

        for view in self.cube_views:
            self.logger.info(f"[CompositeCubeView] fan-out add to cube={view.cube_id}")
            results = view.feedback_memories(feedback_req)
            all_results.extend(results)

        return all_results
