"""
Context Recall — Progressive Recall with token budget enforcement.

PRD §9 — Recall Mechanism:
    Level 1 (default): Summary only
    Level 2 (medium):  Summary + top timeline events
    Level 3 (full):    Summary + timeline + raw messages

PRD §9.3 — Progressive Recall:
    default → summary
    if needed → timeline
    if critical → raw
"""

import logging
from typing import Optional
from openmemo.context_memory.context_store import ContextStore, SessionNode
from openmemo.context_memory.context_summary import ContextSummary
from openmemo.context_memory.context_timeline import ContextTimeline
from openmemo.context_memory.context_budget import ContextBudget, estimate_tokens

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


class ContextRecall:
    """
    Progressive context recall with token budget control.

    Usage:
        recall = ContextRecall(store, budget=ContextBudget(800))

        # Default: summary only (Level 1)
        block = recall.recall("session_1")

        # Medium: summary + timeline (Level 2)
        block = recall.recall("session_1", level=2)

        # Full: summary + timeline + raw (Level 3)
        block = recall.recall("session_1", level=3)
    """

    def __init__(self, store: ContextStore,
                 budget: Optional[ContextBudget] = None,
                 timeline_limit: int = 5):
        self.store = store
        self.budget = budget or ContextBudget(max_tokens=800)
        self.summary_gen = ContextSummary(store)
        self.timeline = ContextTimeline(store)
        self.timeline_limit = timeline_limit

    # ── Progressive Recall ─────────────────────────────────────────────────────

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

        level=1: summary only (default)
        level=2: summary + timeline
        level=3: summary + timeline + raw
        """
        summary_block = self.summary_gen.format_for_prompt(session_id)

        timeline_block = ""
        if level >= 2:
            timeline_block = self.timeline.format_for_prompt(
                session_id, limit=self.timeline_limit
            )

        raw_block = ""
        if level >= 3:
            raw_msgs = self.store.get_raw_messages(session_id, limit=20)
            if raw_msgs:
                raw_lines = [f"[{m.role}] {m.content[:200]}" for m in raw_msgs[-10:]]
                raw_block = "Recent messages:\n" + "\n".join(raw_lines)

        return self.budget.combine(
            user_memory_block=user_memory_block,
            summary_block=summary_block,
            timeline_block=timeline_block,
            raw_block=raw_block,
        )

    def recall_as_dict(self, session_id: str, level: int = 2) -> dict:
        """
        Return structured recall result for API response (PRD §9.2).
        """
        session = self.store.get_session(session_id)
        events = self.store.get_recent_events(session_id, limit=self.timeline_limit)

        return {
            "session_id": session_id,
            "summary": session.summary if session else "",
            "goal": session.goal if session else "",
            "progress": session.progress if session else "",
            "decisions": session.decisions if session else "",
            "issues": session.issues if session else "",
            "timeline": [e.to_dict() for e in events],
            "message_count": session.message_count if session else 0,
        }

    def auto_level(self, session_id: str,
                   user_memory_block: str = "") -> str:
        """
        Auto-select recall level based on budget:
        - Start at level 1 (summary)
        - If budget allows, expand to level 2 (+ timeline)
        - Rarely go to level 3 (raw)
        """
        summary_block = self.summary_gen.format_for_prompt(session_id)
        summary_tokens = estimate_tokens(summary_block)
        um_tokens = estimate_tokens(user_memory_block)
        used = summary_tokens + um_tokens

        remaining = self.budget.max_tokens - used
        if remaining > 200:
            # Enough room for timeline
            return self.recall(session_id, level=2,
                               user_memory_block=user_memory_block)
        return self.recall(session_id, level=1,
                           user_memory_block=user_memory_block)
