"""
Base Memory Adapter — Universal interface for all agent framework integrations.

All adapters must inherit from BaseMemoryAdapter and implement:
- write_memory()
- recall_memory()
- inject_context()
"""

import logging
import time
from typing import List, Dict, Optional
from dataclasses import dataclass, field

logger = logging.getLogger("openmemo")


@dataclass
class AdapterMetrics:
    writes: int = 0
    recalls: int = 0
    injections: int = 0
    errors: int = 0
    total_write_ms: float = 0
    total_recall_ms: float = 0

    @property
    def avg_write_ms(self) -> float:
        return self.total_write_ms / self.writes if self.writes else 0

    @property
    def avg_recall_ms(self) -> float:
        return self.total_recall_ms / self.recalls if self.recalls else 0

    def summary(self) -> dict:
        return {
            "writes": self.writes,
            "recalls": self.recalls,
            "injections": self.injections,
            "errors": self.errors,
            "avg_write_ms": round(self.avg_write_ms, 1),
            "avg_recall_ms": round(self.avg_recall_ms, 1),
        }


class BaseMemoryAdapter:
    adapter_name: str = "base"

    def __init__(self, db_path: str = "openmemo.db", agent_id: str = "",
                 memory=None, base_url: str = None, api_key: str = None,
                 default_scene: str = "", recall_limit: int = 5,
                 default_scope: str = "private", conversation_id: str = "",
                 user_id: str = "", enable_user_memory: bool = True):
        self.agent_id = agent_id
        self.user_id = user_id or agent_id
        self.default_scene = default_scene
        self.recall_limit = recall_limit
        self.default_scope = default_scope
        self.conversation_id = conversation_id
        self.metrics = AdapterMetrics()

        # User Memory Layer (PRD v1.0)
        self._user_memory = None
        if enable_user_memory:
            try:
                from openmemo.user_memory import UserMemory
                self._user_memory = UserMemory(db_path=db_path)
            except Exception as e:
                logger.warning("[openmemo:%s] user_memory init failed: %s",
                               self.adapter_name, e)

        # Context Memory Layer (PRD v1.0)
        self._context_memory = None
        try:
            from openmemo.context_memory import ContextMemory
            self._context_memory = ContextMemory(db_path=db_path)
        except Exception as e:
            logger.warning("[openmemo:%s] context_memory init failed: %s",
                           self.adapter_name, e)

        # Memory Orchestrator (PRD v1.0 — Control Plane)
        self._orchestrator = None
        try:
            from openmemo.memory_orchestrator import MemoryOrchestrator
            self._orchestrator = MemoryOrchestrator(
                db_path=db_path,
                user_memory=self._user_memory,
                context_memory=self._context_memory,
            )
        except Exception as e:
            logger.warning("[openmemo:%s] orchestrator init failed: %s",
                           self.adapter_name, e)

        if memory:
            self._memory = memory
        elif base_url:
            from openmemo.api.remote import RemoteMemory
            self._memory = RemoteMemory(base_url=base_url, api_key=api_key)
        else:
            from openmemo.api.sdk import Memory
            self._memory = Memory(db_path=db_path)

    def write_memory(self, content: str, scene: str = None,
                     memory_type: str = "fact", confidence: float = 0.8,
                     metadata: dict = None, scope: str = None,
                     conversation_id: str = None) -> str:
        effective_scene = scene or self.default_scene
        effective_scope = scope or self.default_scope
        effective_conv = conversation_id or self.conversation_id
        start = time.time()
        try:
            result = self._memory.write_memory(
                content=content,
                scene=effective_scene,
                memory_type=memory_type,
                confidence=confidence,
                agent_id=self.agent_id,
                metadata=metadata,
                scope=effective_scope,
                conversation_id=effective_conv,
            )
            elapsed = (time.time() - start) * 1000
            self.metrics.writes += 1
            self.metrics.total_write_ms += elapsed
            logger.info("[openmemo:%s] write_memory scene=%s type=%s scope=%s latency=%.0fms",
                        self.adapter_name, effective_scene, memory_type, effective_scope, elapsed)
            return result
        except Exception as e:
            self.metrics.errors += 1
            logger.warning("[openmemo:%s] write_memory failed: %s", self.adapter_name, e)
            return ""

    def recall_memory(self, query: str, scene: str = None,
                      limit: int = None, conversation_id: str = None) -> List[Dict]:
        effective_scene = scene or self.default_scene
        effective_limit = limit or self.recall_limit
        effective_conv = conversation_id or self.conversation_id
        start = time.time()
        try:
            results = self._memory.search_memory(
                query=query,
                scene=effective_scene,
                agent_id=self.agent_id,
                limit=effective_limit,
                conversation_id=effective_conv,
            )
            elapsed = (time.time() - start) * 1000
            self.metrics.recalls += 1
            self.metrics.total_recall_ms += elapsed
            logger.info("[openmemo:%s] recall_memory query=%s scene=%s hits=%d latency=%.0fms",
                        self.adapter_name, query[:30], effective_scene, len(results), elapsed)
            return results
        except Exception as e:
            self.metrics.errors += 1
            logger.warning("[openmemo:%s] recall_memory failed: %s", self.adapter_name, e)
            return []

    def recall_context(self, query: str, scene: str = None,
                       limit: int = None, mode: str = "kv",
                       conversation_id: str = None) -> dict:
        effective_scene = scene or self.default_scene
        effective_limit = limit or self.recall_limit
        effective_conv = conversation_id or self.conversation_id
        start = time.time()
        try:
            result = self._memory.recall_context(
                query=query,
                scene=effective_scene,
                agent_id=self.agent_id,
                limit=effective_limit,
                mode=mode,
                conversation_id=effective_conv,
            )
            elapsed = (time.time() - start) * 1000
            self.metrics.recalls += 1
            self.metrics.total_recall_ms += elapsed
            return result
        except Exception as e:
            self.metrics.errors += 1
            logger.warning("[openmemo:%s] recall_context failed: %s", self.adapter_name, e)
            return {"context": []}

    # ── §11.1  User Memory Recall ──────────────────────────────────────────────
    def user_memory_recall(self, top_n: int = 5) -> str:
        """
        PRD §11.1 — Recall user profile as a prompt-ready block.

        Recall order (PRD §11.2):
            user memory → private → shared → team → experience graph

        Returns:
            str: USER PROFILE block, e.g.
                 "USER PROFILE\\nUser Preferences:\\n  - prefers FastAPI"
                 Empty string when no user memory exists.
        """
        if not self._user_memory or not self.user_id:
            return ""
        try:
            return self._user_memory.format_for_prompt(self.user_id, top_n=top_n)
        except Exception as e:
            logger.warning("[openmemo:%s] user_memory_recall failed: %s",
                           self.adapter_name, e)
            return ""

    def user_memory_store(self, text: str) -> list:
        """Auto-extract and persist user preferences/constraints from free text."""
        if not self._user_memory or not self.user_id:
            return []
        try:
            return self._user_memory.store(self.user_id, text)
        except Exception as e:
            logger.warning("[openmemo:%s] user_memory_store failed: %s",
                           self.adapter_name, e)
            return []

    def rank_workflows_for_user(self, workflows: list,
                                 score_key: str = "score") -> list:
        """
        PRD §10 — Re-rank workflow list using user alignment score.

        final_score = workflow_score * 0.7 + user_alignment_score * 0.3
        """
        if not self._user_memory or not self.user_id:
            return workflows
        try:
            return self._user_memory.rank_workflows(
                self.user_id, workflows, score_key=score_key
            )
        except Exception as e:
            logger.warning("[openmemo:%s] rank_workflows_for_user failed: %s",
                           self.adapter_name, e)
            return workflows

    # ── Context Memory Layer ────────────────────────────────────────────────────
    def context_on_message(self, session_id: str, role: str,
                           content: str) -> str:
        """PRD §12.1 — Record a message to the context session."""
        if not self._context_memory:
            return ""
        try:
            return self._context_memory.add_message(
                session_id, role, content, user_id=self.user_id
            )
        except Exception as e:
            logger.warning("[openmemo:%s] context_on_message failed: %s",
                           self.adapter_name, e)
            return ""

    def context_on_agent_step(self, session_id: str, step_text: str) -> None:
        """PRD §12.1 — Record an agent step event."""
        if not self._context_memory:
            return
        try:
            self._context_memory.on_agent_step(
                session_id, step_text, user_id=self.user_id
            )
        except Exception as e:
            logger.warning("[openmemo:%s] context_on_agent_step failed: %s",
                           self.adapter_name, e)

    def context_on_task_complete(self, session_id: str,
                                  task: str = "", result: str = "") -> None:
        """PRD §12.1 — Trigger context summary on task completion."""
        if not self._context_memory:
            return
        try:
            self._context_memory.on_task_complete(session_id, task=task,
                                                   result=result)
        except Exception as e:
            logger.warning("[openmemo:%s] context_on_task_complete failed: %s",
                           self.adapter_name, e)

    def context_recall(self, session_id: str, level: int = 1) -> str:
        """PRD §9 — Progressive context recall (level 1/2/3)."""
        if not self._context_memory:
            return ""
        try:
            return self._context_memory.recall(session_id, level=level)
        except Exception as e:
            logger.warning("[openmemo:%s] context_recall failed: %s",
                           self.adapter_name, e)
            return ""

    # ── §9 Prompt Injection (context → user → memories) ───────────────────────
    def inject_context(self, prompt: str, query: str = None,
                       scene: str = None, limit: int = None,
                       include_user_memory: bool = True,
                       session_id: str = "") -> str:
        """
        PRD §11.1 — Inject context into prompt.

        Injection order (Context Memory PRD §11.1):
            1. CONTEXT SUMMARY  (context memory — conversation continuity)
            2. USER PROFILE     (user memory — preferences/constraints)
            3. Relevant memories (private → shared → team)
            4. Original prompt
        """
        effective_query = query or prompt
        result = self.recall_context(effective_query, scene=scene, limit=limit)
        context = result.get("context", [])

        self.metrics.injections += 1

        parts = []

        # §11.1 Priority 1: Context Memory (CONTEXT SUMMARY block — topmost)
        if session_id and self._context_memory:
            try:
                ctx_block = self._context_memory.inject(session_id)
                if ctx_block:
                    parts.append(ctx_block)
            except Exception as e:
                logger.warning("[openmemo:%s] context inject failed: %s",
                               self.adapter_name, e)

        # §11.1 Priority 2: User Memory (USER PROFILE block)
        if include_user_memory:
            user_block = self.user_memory_recall()
            if user_block:
                parts.append(user_block)

        # §11.1 Priority 3: Relevant memories (private → shared → team)
        if context:
            memory_block = "Relevant memories:\n"
            for i, mem in enumerate(context, 1):
                memory_block += f"{i}. {mem}\n"
            parts.append(memory_block)

        if not parts:
            return prompt

        parts.append(prompt)
        injected = "\n\n".join(parts)
        logger.info("[openmemo:%s] inject_context session=%s user_memory=%s memories=%d",
                    self.adapter_name, bool(session_id),
                    bool(include_user_memory), len(context))
        return injected

    # ── Memory Orchestrator ────────────────────────────────────────────────────
    def orchestrate_prompt(
        self,
        task: str,
        session_id: str = "",
        team_id: str = "",
        agent_role: str = "",
        task_type: str = "",
        mode: str = "",
        system_block: str = "",
        user_input: str = "",
    ) -> dict:
        """
        PRD §14 — Full orchestration pipeline via on_agent_step.

        Returns orchestrator result dict:
            {
                "context":      str,   # full injected prompt
                "task_type":    str,
                "mode":         str,
                "token_usage":  int,
                "latency_ms":   float,
                ...
            }
        Falls back to basic inject_context when orchestrator unavailable.
        """
        if not self._orchestrator:
            # Graceful fallback
            prompt = self.inject_context(
                task, session_id=session_id or self.conversation_id
            )
            return {"context": prompt, "task_type": "coding", "token_usage": 0}

        try:
            return self._orchestrator.build_prompt(
                task=task,
                user_id=self.user_id,
                session_id=session_id or self.conversation_id,
                agent_id=self.agent_id,
                team_id=team_id,
                agent_role=agent_role or self.adapter_name,
                task_type=task_type,
                mode=mode,
                system_block=system_block,
                user_input=user_input or task,
            )
        except Exception as e:
            logger.warning("[openmemo:%s] orchestrate_prompt failed: %s",
                           self.adapter_name, e)
            return {"context": task, "task_type": "coding", "token_usage": 0}

    def get_context(self, query: str, scene: str = None,
                    limit: int = 3) -> List[str]:
        result = self.recall_context(query, scene=scene, limit=limit)
        return result.get("context", [])

    def list_scenes(self) -> List[str]:
        try:
            return self._memory.list_scenes(agent_id=self.agent_id)
        except Exception as e:
            logger.warning("[openmemo:%s] list_scenes failed: %s", self.adapter_name, e)
            return []

    def get_metrics(self) -> dict:
        return self.metrics.summary()

    def close(self):
        if hasattr(self._memory, 'close'):
            self._memory.close()
