"""
OpenMemo Context Memory Layer

Lightweight Context Continuity for AI agents.
PRD: v1.0 — session summary / timeline / progressive recall / token budget

    from openmemo.context_memory import ContextMemory

    cm = ContextMemory("openmemo.db")

    # Record messages
    cm.add_message("session_1", "user", "Deploy my FastAPI backend")
    cm.add_message("session_1", "assistant", "I'll use docker-compose for you")

    # Record events
    cm.add_event("session_1", "action", "Docker image built successfully")
    cm.add_event("session_1", "decision", "Using docker-compose")

    # Get context for prompt injection (PRD §11)
    block = cm.inject("session_1")

    # Recall at different levels
    summary = cm.recall("session_1", level=1)   # summary only
    full    = cm.recall("session_1", level=2)   # summary + timeline
"""

from openmemo.context_memory.context_store import (
    ContextStore, SessionNode, TimelineEvent, RawMessage
)
from openmemo.context_memory.context_summary import ContextSummary
from openmemo.context_memory.context_timeline import ContextTimeline
from openmemo.context_memory.context_recall import ContextRecall
from openmemo.context_memory.context_budget import ContextBudget, estimate_tokens


class ContextMemory:
    """
    Unified facade for the Context Memory Layer.

    Combines Store + Summary + Timeline + Recall + Budget into one interface.
    """

    def __init__(self, db_path: str = "openmemo.db",
                 max_tokens: int = 800,
                 summarize_every_n: int = 5,
                 llm_client=None):
        self._store = ContextStore(db_path=db_path)
        self._budget = ContextBudget(max_tokens=max_tokens)
        self._summary = ContextSummary(
            self._store,
            summarize_every_n=summarize_every_n,
            llm_client=llm_client,
        )
        self._timeline = ContextTimeline(self._store)
        self._recall = ContextRecall(self._store, budget=self._budget)

    # ── Write ──────────────────────────────────────────────────────────────────

    def add_message(self, session_id: str, role: str,
                    content: str, user_id: str = "") -> str:
        """
        Record a raw message and auto-trigger summary if needed (PRD §7.1).

        role: "user" | "assistant" | "tool"
        Returns message_id.
        """
        self._store.get_or_create_session(session_id, user_id=user_id)
        msg = RawMessage(session_id=session_id, role=role, content=content)
        mid = self._store.add_raw_message(msg)
        count = self._store.increment_message_count(session_id)

        # Auto-trigger summary every N messages
        if self._summary.should_summarize(session_id):
            self._summary.summarize(session_id)

        return mid

    def add_event(self, session_id: str, event_type: str,
                  content: str) -> str:
        """
        Record a timeline event (PRD §8.1).

        event_type: "decision" | "action" | "result" | "error"
        """
        self._store.get_or_create_session(session_id)
        return self._timeline.add(session_id, event_type, content) or ""

    def on_agent_step(self, session_id: str, step_text: str,
                      user_id: str = "") -> None:
        """PRD §12.1 — Hook: agent step."""
        self._store.get_or_create_session(session_id, user_id=user_id)
        self._timeline.on_agent_step(session_id, step_text)

    def on_tool_call(self, session_id: str, tool_name: str,
                     result: str) -> None:
        """PRD §12.1 — Hook: tool call."""
        self._timeline.on_tool_call(session_id, tool_name, result)

    def on_task_complete(self, session_id: str,
                         task: str = "", result: str = "") -> SessionNode:
        """PRD §7.1 + §12.1 — Hook: task complete. Triggers summarization."""
        self._store.get_or_create_session(session_id)
        if result:
            self._timeline.on_result(session_id, result)
        return self._summary.summarize_on_complete(session_id, task, result)

    # ── Recall ─────────────────────────────────────────────────────────────────

    def recall(self, session_id: str, level: int = 1,
               user_memory_block: str = "") -> str:
        """
        Progressive recall — returns context block.

        level=1: summary only (default)
        level=2: summary + timeline
        level=3: summary + timeline + raw messages
        """
        return self._recall.recall(session_id, level=level,
                                   user_memory_block=user_memory_block)

    def inject(self, session_id: str,
               user_memory_block: str = "") -> str:
        """
        Auto-level recall for prompt injection (PRD §11).

        Returns the context block to prepend to the prompt.
        Returns empty string when session is empty.
        """
        return self._recall.auto_level(session_id,
                                       user_memory_block=user_memory_block)

    def format_summary(self, session_id: str) -> str:
        """Return CONTEXT SUMMARY block (PRD §11.2)."""
        return self._summary.format_for_prompt(session_id)

    def format_timeline(self, session_id: str, limit: int = 5) -> str:
        """Return timeline block."""
        return self._timeline.format_for_prompt(session_id, limit=limit)

    def to_dict(self, session_id: str) -> dict:
        """Return structured dict for API response."""
        return self._recall.recall_as_dict(session_id)

    # ── Admin ───────────────────────────────────────────────────────────────────

    def get_session(self, session_id: str) -> SessionNode:
        return self._store.get_or_create_session(session_id)

    def list_sessions(self, user_id: str = "") -> list:
        return self._store.list_sessions(user_id=user_id)

    def delete_session(self, session_id: str) -> bool:
        return self._store.delete_session(session_id)

    def close(self):
        self._store.close()


__all__ = [
    "ContextMemory",
    "ContextStore", "SessionNode", "TimelineEvent", "RawMessage",
    "ContextSummary", "ContextTimeline", "ContextRecall", "ContextBudget",
    "estimate_tokens",
]
