"""
Memory Orchestrator — Central Control Plane for all memory layers.

PRD §1 / §5 — Core orchestration flow:
    1. Analyze Task     (TaskClassifier)
    2. Select Memory    (MemorySelector)
    3. Rank Memory      (MemoryRanker)
    4. Allocate Tokens  (BudgetAllocator)
    5. Fetch Memory     (all layers)
    6. Build Prompt     (PromptBuilder)

PRD §15 — Performance:
    Latency < 80ms
    Token utilization > 85%

PRD §19 — Design Principle:
    "Not more memory — the RIGHT memory."
"""

import time
import logging
from typing import Optional, Dict, Any

from openmemo.memory_orchestrator.task_classifier import TaskClassifier
from openmemo.memory_orchestrator.selector import MemorySelector, SelectionPlan
from openmemo.memory_orchestrator.ranker import MemoryRanker
from openmemo.memory_orchestrator.budget_allocator import BudgetAllocator
from openmemo.memory_orchestrator.prompt_builder import PromptBuilder

logger = logging.getLogger("openmemo.orchestrator")


class MemoryOrchestrator:
    """
    Unified Memory Decision & Injection Engine.

    Wires together all memory layers and produces an optimally-injected prompt.

    Usage (basic):
        orchestrator = MemoryOrchestrator(db_path="openmemo.db")
        result = orchestrator.build_prompt(
            task="Deploy my FastAPI backend",
            user_id="alice",
            session_id="session_123",
        )
        print(result["context"])          # injected prompt
        print(result["task_type"])        # "deployment"
        print(result["token_usage"])      # e.g. 1450

    Usage (adapter):
        # In BaseMemoryAdapter.inject_context():
        result = self._orchestrator.build_prompt(task, user_id, session_id)
        injected_prompt = result["context"] + "\\n\\n" + original_prompt
    """

    def __init__(
        self,
        db_path: str = "openmemo.db",
        max_tokens: int = 2000,
        memory=None,                        # openmemo.api.sdk.Memory instance
        user_memory=None,                   # UserMemory instance
        context_memory=None,                # ContextMemory instance
        experience_registry=None,           # ExperienceRegistry instance
    ):
        self.db_path = db_path
        self.max_tokens = max_tokens

        self._classifier = TaskClassifier()
        self._selector = MemorySelector()
        self._ranker = MemoryRanker()
        self._allocator = BudgetAllocator(max_tokens=max_tokens)
        self._builder = PromptBuilder()

        # Memory layer instances (lazy-init if not provided)
        self._memory = memory
        self._user_memory = user_memory
        self._context_memory = context_memory
        self._experience_registry = experience_registry

        self._init_layers()

    def _init_layers(self):
        """Lazy-initialize memory layers that weren't provided."""
        if not self._memory:
            try:
                from openmemo.api.sdk import Memory
                self._memory = Memory(db_path=self.db_path)
            except Exception as e:
                logger.warning("orchestrator: memory init failed: %s", e)

        if not self._user_memory:
            try:
                from openmemo.user_memory import UserMemory
                self._user_memory = UserMemory(db_path=self.db_path)
            except Exception as e:
                logger.warning("orchestrator: user_memory init failed: %s", e)

        if not self._context_memory:
            try:
                from openmemo.context_memory import ContextMemory
                self._context_memory = ContextMemory(
                    db_path=self.db_path, max_tokens=self.max_tokens // 4
                )
            except Exception as e:
                logger.warning("orchestrator: context_memory init failed: %s", e)

        if not self._experience_registry:
            try:
                from openmemo.experience.registry import ExperienceRegistry
                self._experience_registry = ExperienceRegistry(
                    db_path=self.db_path
                )
            except Exception as e:
                logger.warning("orchestrator: experience_registry init failed: %s", e)

    # ── Public API ─────────────────────────────────────────────────────────────

    def build_prompt(
        self,
        task: str,
        user_id: str = "",
        session_id: str = "",
        agent_id: str = "",
        team_id: str = "",
        agent_role: str = "",
        task_type: str = "",          # override classifier if provided
        mode: str = "",               # override ranker mode if provided
        system_block: str = "",
        user_input: str = "",
    ) -> dict:
        """
        PRD §5 — Full orchestration pipeline.

        Returns:
            {
                "context":      str,   # full injected prompt
                "user_memory":  str,   # user memory block
                "context_block": str,  # context summary block
                "task_memory":  str,   # task memory block
                "team_memory":  str,   # team memory block
                "experience":   str,   # experience block
                "user_input":   str,   # original user input
                "task_type":    str,   # detected task type
                "mode":         str,   # selected mode
                "token_usage":  int,   # total tokens used
                "layer_tokens": dict,  # per-layer token counts
                "latency_ms":   float, # orchestration latency
                "plan":         dict,  # selection plan summary
            }
        """
        t0 = time.time()
        user_input = user_input or task

        # ── Step 1: Classify Task ──────────────────────────────────────────────
        effective_type = task_type or self._classifier.classify(task, agent_role)
        effective_mode = mode or self._ranker.get_mode(effective_type)

        # ── Step 2: Select Memory ──────────────────────────────────────────────
        plan = self._selector.select(
            task_type=effective_type,
            user_id=user_id,
            session_id=session_id,
            agent_id=agent_id,
            team_id=team_id,
        )

        # ── Step 3: Rank Memory ────────────────────────────────────────────────
        ranked_layers = self._ranker.rank(effective_type, mode=effective_mode)

        # ── Step 4: Allocate Token Budget ──────────────────────────────────────
        budget = self._allocator.allocate(effective_type)

        # ── Step 5: Fetch Memory from Each Layer ───────────────────────────────
        user_memory_block  = self._fetch_user_memory(plan, budget)
        context_block      = self._fetch_context(plan, budget, user_memory_block)
        task_memory_block  = self._fetch_task_memory(plan, budget, agent_id)
        team_memory_block  = self._fetch_team_memory(plan, budget, team_id)
        experience_block   = self._fetch_experience(plan, budget, task)

        # ── Step 6: Build Prompt ───────────────────────────────────────────────
        result = self._builder.build(
            user_input=user_input,
            user_memory_block=user_memory_block,
            context_block=context_block,
            task_memory_block=task_memory_block,
            team_memory_block=team_memory_block,
            experience_block=experience_block,
            budget=budget,
            system_block=system_block,
        )

        latency_ms = round((time.time() - t0) * 1000, 1)

        result.update({
            "task_type":  effective_type,
            "mode":       effective_mode,
            "latency_ms": latency_ms,
            "plan":       plan.summary(),
            "budget":     budget,
            "ranked_layers": ranked_layers,
        })

        logger.info(
            "orchestrator: task=%s mode=%s type=%s tokens=%d latency=%.0fms",
            task[:50], effective_mode, effective_type,
            result["token_usage"], latency_ms,
        )
        return result

    def classify(self, task: str, agent_role: str = "") -> str:
        """Classify a task string into a task type (shortcut)."""
        return self._classifier.classify(task, agent_role)

    def get_mode(self, task_type: str) -> str:
        """Get the mode for a task type."""
        return self._ranker.get_mode(task_type)

    # ── Memory Fetchers ────────────────────────────────────────────────────────

    def _fetch_user_memory(self, plan: SelectionPlan,
                            budget: dict) -> str:
        if not plan.include_user_memory or not self._user_memory:
            return ""
        try:
            block = self._user_memory.format_for_prompt(
                plan.user_id, top_n=plan.user_memory_top_n
            )
            return block or ""
        except Exception as e:
            logger.warning("orchestrator: user_memory fetch failed: %s", e)
            return ""

    def _fetch_context(self, plan: SelectionPlan, budget: dict,
                        user_memory_block: str = "") -> str:
        if not plan.include_context or not self._context_memory:
            return ""
        try:
            return self._context_memory.recall(
                plan.session_id,
                level=plan.context_level,
                user_memory_block=user_memory_block,
            )
        except Exception as e:
            logger.warning("orchestrator: context fetch failed: %s", e)
            return ""

    def _fetch_task_memory(self, plan: SelectionPlan, budget: dict,
                            agent_id: str = "") -> str:
        if not plan.include_task_memory or not self._memory:
            return ""
        try:
            result = self._memory.recall_context(
                query=f"task state for {plan.agent_id or agent_id}",
                agent_id=plan.agent_id or agent_id,
                scene="task",
                mode="kv",
                limit=3,
            )
            items = result.get("context", [])
            if not items:
                return ""
            lines = ["TASK MEMORY"]
            for item in items[:3]:
                lines.append(f"  - {item}")
            return "\n".join(lines)
        except Exception as e:
            logger.warning("orchestrator: task_memory fetch failed: %s", e)
            return ""

    def _fetch_team_memory(self, plan: SelectionPlan, budget: dict,
                            team_id: str = "") -> str:
        if not plan.include_team_memory or not self._memory:
            return ""
        try:
            result = self._memory.recall_context(
                query="team shared context",
                scene=f"team:{plan.team_id or team_id}",
                mode="kv",
                limit=3,
            )
            items = result.get("context", [])
            if not items:
                return ""
            lines = ["TEAM MEMORY"]
            for item in items[:3]:
                lines.append(f"  - {item}")
            return "\n".join(lines)
        except Exception as e:
            logger.warning("orchestrator: team_memory fetch failed: %s", e)
            return ""

    def _fetch_experience(self, plan: SelectionPlan, budget: dict,
                           task: str = "") -> str:
        if not plan.include_experience or not self._experience_registry:
            return ""
        try:
            experiences = self._experience_registry.recall(
                query=task, top_k=plan.experience_top_k
            )
            if not experiences:
                return ""
            lines = ["EXPERIENCE"]
            for exp in experiences:
                title = getattr(exp, "title", "") or getattr(exp, "name", "")
                desc = getattr(exp, "description", "") or getattr(exp, "content", "")
                if title:
                    lines.append(f"  Workflow: {title}")
                if desc:
                    lines.append(f"  {desc[:200]}")
            return "\n".join(lines) if len(lines) > 1 else ""
        except Exception as e:
            logger.warning("orchestrator: experience fetch failed: %s", e)
            return ""
