"""Context-aware checkpoint system for compaction-resilient recovery.

Provides three-tier checkpointing that fires when context usage crosses 50%
of the estimated context window. Tiers map to the SessionState injection
budget (essential / standard / full):

    Tier 1 — Conversational: current task, blockers, next steps, active files
    Tier 2 — Active work: work item, plan, decisions, progress
    Tier 3 — Session summary: tasks, milestones, custom context

Also maintains a sparse checkpoint fallback (via sessions.checkpoint) so
recovery works even if structured state is unavailable.

Constants:
    CONTEXT_WINDOW_BYTES: Estimated context window (~20 MB).
    CHECKPOINT_THRESHOLD_BYTES: Trigger at 50% usage.
    MIN_CHECKPOINT_INTERVAL_BYTES: Suppress rapid checkpoints (2 MB gap).
"""
from __future__ import annotations

import json
import logging
from datetime import datetime, timezone
from typing import Any

from .db import Database
from .session_state import (
    TIER_STANDARD,
    PlanStep,
    SessionState,
    SessionStateManager,
)
from .sessions import SessionManager

logger = logging.getLogger(__name__)

# ─── Constants ──────────────────────────────────────────────

# Estimated context window size in bytes (~20 MB).
CONTEXT_WINDOW_BYTES: int = 20_000_000

# Checkpoint triggers at 50% context usage.
CHECKPOINT_THRESHOLD_BYTES: int = CONTEXT_WINDOW_BYTES // 2

# Minimum bytes of new context between checkpoints to prevent thrashing.
MIN_CHECKPOINT_INTERVAL_BYTES: int = 2_000_000


# ─── Manager ────────────────────────────────────────────────


class ContextCheckpointManager:
    """Three-tier context-aware checkpoint manager.

    Coordinates structured session state persistence with sparse checkpoint
    fallback for compaction-resilient recovery.
    """

    def __init__(self, db: Database) -> None:
        self.db = db
        self._sessions = SessionManager(db)
        self._states = SessionStateManager(db)

    # ─── Trigger logic ──────────────────────────────────────

    def should_checkpoint(
        self,
        current_bytes: int,
        last_checkpoint_bytes: int = 0,
    ) -> bool:
        """Determine whether a checkpoint should fire.

        Returns True when *current_bytes* is at or above the 50% threshold
        and enough new context has accumulated since the last checkpoint
        to avoid thrashing.

        Args:
            current_bytes: Current estimated context usage in bytes.
            last_checkpoint_bytes: Byte count at the last checkpoint (0 = first).

        Returns:
            True if a checkpoint should be taken.
        """
        if current_bytes < CHECKPOINT_THRESHOLD_BYTES:
            return False

        # Thrash prevention: require minimum interval since last checkpoint
        if last_checkpoint_bytes > 0:
            delta = current_bytes - last_checkpoint_bytes
            if delta < MIN_CHECKPOINT_INTERVAL_BYTES:
                return False

        return True

    # ─── Save ───────────────────────────────────────────────

    def save_checkpoint(
        self,
        session_id: str = "",
        context_metrics: dict[str, Any] | None = None,
        conversational_state: dict[str, Any] | None = None,
        active_work: dict[str, Any] | None = None,
        session_summary: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Save a three-tier checkpoint.

        Args:
            session_id: Session to checkpoint. Auto-detects if empty.
            context_metrics: Dict with ``estimated_context_bytes`` etc.
            conversational_state: Tier 1 — current_task, blockers, next_steps,
                active_files.
            active_work: Tier 2 — work_item, plan_title, plan_steps, decisions,
                progress_summary.
            session_summary: Tier 3 — tasks, milestones, custom.

        Returns:
            Dict with session_id, checkpoint_id, state_id, context_pct,
            tier_summary, and checkpoint_at.
        """
        context_metrics = context_metrics or {}
        conversational_state = conversational_state or {}
        active_work = active_work or {}
        session_summary = session_summary or {}

        # Resolve session — auto-detect or create if needed
        if not session_id:
            current = self._sessions.get_current()
            if current:
                session_id = current["id"]
            else:
                new_session = self._sessions.start()
                session_id = new_session["id"]

        # Compaction number from session
        session = self._sessions.get(session_id)
        compaction_number = (
            session.get("compaction_count", 0) if session else 0
        )

        # Build structured state from the three tiers
        state = self._build_state(
            session_id,
            compaction_number,
            conversational_state,
            active_work,
            session_summary,
        )

        # Persist structured state
        save_result = self._states.save(state)

        # Context percentage
        est_bytes = context_metrics.get("estimated_context_bytes", 0)
        context_pct = (
            int(est_bytes / CONTEXT_WINDOW_BYTES * 100) if est_bytes else 0
        )

        # Build and persist sparse fallback checkpoint
        context_summary = self._build_context_summary(
            conversational_state, active_work, context_metrics,
        )
        open_items = self._build_open_items_json(active_work, session_summary)
        cp_result = self._sessions.checkpoint(
            session_id,
            context_summary=context_summary,
            open_items=open_items,
        )

        return {
            "session_id": session_id,
            "checkpoint_id": cp_result.get("checkpoint_id", ""),
            "state_id": save_result["state_id"],
            "context_pct": context_pct,
            "tier_summary": {
                "conversational": bool(conversational_state),
                "active_work": bool(active_work),
                "session_summary": bool(session_summary),
            },
            "checkpoint_at": save_result["captured_at"],
        }

    # ─── Recovery ───────────────────────────────────────────

    def load_for_recovery(
        self,
        session_id: str,
        max_chars: int = TIER_STANDARD,
    ) -> str | None:
        """Load recovery context for injection after compaction.

        Prefers structured state (via ``SessionState.format_for_injection``);
        falls back to the sparse checkpoint stored by ``sessions.checkpoint``.

        Args:
            session_id: The session to recover.
            max_chars: Character budget for the injection text.

        Returns:
            Formatted markdown string, or ``None`` if no state exists.
        """
        # Structured state takes priority
        state = self._states.load_latest(session_id)
        if state:
            return state.format_for_injection(max_chars)

        # Sparse checkpoint fallback
        cp = self._sessions.get_latest_checkpoint(session_id)
        if cp:
            return self._format_sparse_checkpoint(cp)

        return None

    def load_state(self, session_id: str) -> SessionState | None:
        """Load the raw SessionState for direct access.

        Args:
            session_id: The session to load state for.

        Returns:
            SessionState if found, ``None`` otherwise.
        """
        return self._states.load_latest(session_id)

    def get_checkpoint_history(
        self, session_id: str,
    ) -> list[dict[str, Any]]:
        """Get metadata for all structured-state snapshots in a session.

        Args:
            session_id: The session to query.

        Returns:
            List of dicts with compaction_number, captured_at, etc.
        """
        return self._states.get_history(session_id)

    # ─── Internal helpers ───────────────────────────────────

    def _build_state(
        self,
        session_id: str,
        compaction_number: int,
        conversational_state: dict[str, Any],
        active_work: dict[str, Any],
        session_summary: dict[str, Any],
    ) -> SessionState:
        """Assemble a SessionState from the three input tiers.

        Normalizes loose types:
        - string file paths -> ``{"path": ..., "role": "referenced"}``
        - string blockers -> single-element list
        - string plan steps -> ``PlanStep(description=..., status="pending")``
        - milestones merged into ``custom``
        """
        now = datetime.now(timezone.utc).isoformat()

        # Normalize active_files: bare strings -> dicts
        raw_files = conversational_state.get("active_files", [])
        active_files = [
            {"path": f, "role": "referenced"} if isinstance(f, str) else f
            for f in raw_files
        ]

        # Normalize blockers: single string -> list
        blockers = conversational_state.get("blockers", [])
        if isinstance(blockers, str):
            blockers = [blockers]

        # Normalize plan_steps: strings -> PlanStep objects
        raw_steps = active_work.get("plan_steps", [])
        plan_steps: list[PlanStep] = []
        for step in raw_steps:
            if isinstance(step, str):
                plan_steps.append(PlanStep(description=step, status="pending"))
            elif isinstance(step, dict):
                plan_steps.append(PlanStep.from_dict(step))
            else:
                plan_steps.append(step)

        # Merge milestones into custom
        custom = dict(session_summary.get("custom", {}))
        milestones = session_summary.get("milestones")
        if milestones:
            custom["milestones"] = milestones

        return SessionState(
            session_id=session_id,
            compaction_number=compaction_number,
            captured_at=now,
            active_work_item=active_work.get("work_item"),
            current_task=conversational_state.get("current_task", ""),
            plan_title=active_work.get("plan_title", ""),
            plan_steps=plan_steps,
            current_step_index=active_work.get("current_step_index", -1),
            tasks=session_summary.get("tasks", []),
            decisions=active_work.get("decisions", []),
            active_files=active_files,
            progress_summary=active_work.get("progress_summary", ""),
            next_steps=conversational_state.get("next_steps", ""),
            blockers=blockers,
            custom=custom,
        )

    def _build_context_summary(
        self,
        conversational_state: dict[str, Any],
        active_work: dict[str, Any],
        context_metrics: dict[str, Any],
    ) -> str:
        """Build a one-line context summary for the sparse checkpoint."""
        parts = ["Context-aware checkpoint"]

        est_bytes = context_metrics.get("estimated_context_bytes", 0)
        if est_bytes:
            pct = int(est_bytes / CONTEXT_WINDOW_BYTES * 100)
            parts.append(f"Context: {pct}% used")

        task = conversational_state.get("current_task")
        if task:
            parts.append(f"Task: {task}")

        work_item = active_work.get("work_item")
        if work_item:
            wi_id = work_item.get("id", "?")
            wi_title = work_item.get("title", "?")
            parts.append(f"Work item: {wi_id} \u2014 {wi_title}")

        return " | ".join(parts)

    def _build_open_items_json(
        self,
        active_work: dict[str, Any],
        session_summary: dict[str, Any],
    ) -> str:
        """Build a compact JSON array of open items for sparse checkpoint.

        Caps output at approximately 500 characters.
        """
        items: list[dict[str, Any]] = []

        work_item = active_work.get("work_item")
        if work_item:
            items.append(work_item)

        for task in session_summary.get("tasks", []):
            if task.get("status") != "completed":
                items.append(task)

        result = json.dumps(items, separators=(",", ":"))

        # Trim items from the end until within budget
        while len(result) > 500 and items:
            items.pop()
            result = json.dumps(items, separators=(",", ":"))

        return result

    @staticmethod
    def _format_sparse_checkpoint(checkpoint_data: dict[str, Any]) -> str:
        """Format a sparse checkpoint dict as markdown for context injection.

        Args:
            checkpoint_data: Row from the session_checkpoints table.

        Returns:
            Formatted markdown string.
        """
        parts = ["## Session Checkpoint Recovery"]

        checkpoint_at = checkpoint_data.get("checkpoint_at", "")
        if checkpoint_at:
            parts.append(f"**Checkpoint at:** {checkpoint_at}")

        summary = checkpoint_data.get("context_summary", "")
        if summary:
            parts.append(f"**Context:** {summary}")

        open_items_raw = checkpoint_data.get("open_items", "")
        if open_items_raw:
            try:
                items = json.loads(open_items_raw)
                if items:
                    parts.append("**Open items:**")
                    for item in items:
                        title = item.get(
                            "title",
                            item.get("content", str(item)),
                        )
                        parts.append(f"- {title}")
            except (json.JSONDecodeError, TypeError):
                parts.append(f"**Open items:** {open_items_raw}")

        return "\n".join(parts)
