"""Structured session state for compaction-resilient recovery.

Captures everything needed to resume work after context compaction:
plan state, tasks, decisions, file context, and progress. Stored as
versioned JSON in the session_states table.

The key insight: recovery artifacts tell Claude *about* what it knew,
but structured state tells Claude *exactly where it was* — which step,
which files, which decisions were made and why.
"""
import json
import uuid
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from typing import Any

from .db import Database

# Current schema version — bump when adding fields
STATE_VERSION = 1

# Retention: keep last N states per session
MAX_STATES_PER_SESSION = 5

# Budget tiers for recovery injection (chars, not tokens)
TIER_ESSENTIAL = 2000   # Always included: task, plan step, blockers
TIER_STANDARD = 6000    # + decisions, file context, progress
TIER_FULL = 12000       # + all tasks, full plan, custom context


@dataclass
class PlanStep:
    """A single step in an implementation plan."""

    description: str
    status: str = "pending"  # pending | in_progress | completed | skipped
    notes: str = ""

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dict."""
        return {"description": self.description, "status": self.status, "notes": self.notes}

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "PlanStep":
        """Deserialize from dict."""
        return cls(
            description=data.get("description", ""),
            status=data.get("status", "pending"),
            notes=data.get("notes", ""),
        )


@dataclass
class SessionState:
    """Complete structured state for session recovery after compaction.

    This is the central data structure that survives compaction. Every field
    is something Claude needs to resume work without re-reading the full
    conversation history.
    """

    # Identity
    session_id: str = ""
    compaction_number: int = 0
    captured_at: str = ""

    # What we're doing
    active_work_item: dict[str, Any] | None = None  # {id, title, description, status}
    current_task: str = ""  # Natural language description

    # Plan state
    plan_title: str = ""
    plan_steps: list[PlanStep] = field(default_factory=list)
    current_step_index: int = -1  # -1 = no plan active

    # Todo items
    tasks: list[dict[str, Any]] = field(default_factory=list)  # [{content, status}]

    # Decision trail
    decisions: list[dict[str, Any]] = field(default_factory=list)
    # [{decision, rationale, timestamp}]

    # File context — which files are actively being worked on
    active_files: list[dict[str, Any]] = field(default_factory=list)
    # [{path, role}] role: editing | reading | testing | created

    # Progress
    progress_summary: str = ""  # What's been accomplished so far
    next_steps: str = ""  # What to do next

    # Blockers
    blockers: list[str] = field(default_factory=list)

    # Extensible context
    custom: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Serialize to a JSON-safe dict."""
        return {
            "session_id": self.session_id,
            "compaction_number": self.compaction_number,
            "captured_at": self.captured_at,
            "active_work_item": self.active_work_item,
            "current_task": self.current_task,
            "plan_title": self.plan_title,
            "plan_steps": [s.to_dict() for s in self.plan_steps],
            "current_step_index": self.current_step_index,
            "tasks": self.tasks,
            "decisions": self.decisions,
            "active_files": self.active_files,
            "progress_summary": self.progress_summary,
            "next_steps": self.next_steps,
            "blockers": self.blockers,
            "custom": self.custom,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "SessionState":
        """Deserialize from dict, handling missing fields gracefully."""
        steps_raw = data.get("plan_steps", [])
        steps = [PlanStep.from_dict(s) if isinstance(s, dict) else s for s in steps_raw]

        return cls(
            session_id=data.get("session_id", ""),
            compaction_number=data.get("compaction_number", 0),
            captured_at=data.get("captured_at", ""),
            active_work_item=data.get("active_work_item"),
            current_task=data.get("current_task", ""),
            plan_title=data.get("plan_title", ""),
            plan_steps=steps,
            current_step_index=data.get("current_step_index", -1),
            tasks=data.get("tasks", []),
            decisions=data.get("decisions", []),
            active_files=data.get("active_files", []),
            progress_summary=data.get("progress_summary", ""),
            next_steps=data.get("next_steps", ""),
            blockers=data.get("blockers", []),
            custom=data.get("custom", {}),
        )

    def format_for_injection(self, max_chars: int = TIER_STANDARD) -> str:
        """Format state for injection into Claude's context after compaction.

        Uses a tiered approach: essential fields always included, optional
        fields added if budget allows. Output is structured markdown that
        Claude can parse and act on immediately.

        Args:
            max_chars: Maximum characters for the output.

        Returns:
            Formatted markdown string for context injection.
        """
        parts: list[str] = []
        chars_used = 0

        def add(section: str) -> bool:
            nonlocal chars_used
            cost = len(section)
            if chars_used + cost > max_chars:
                return False
            parts.append(section)
            chars_used += cost
            return True

        # --- Tier 1: Essential (always included) ---
        header = f"## Structured Session State (compaction #{self.compaction_number})"
        add(header)

        if self.current_task:
            add(f"**Current task:** {self.current_task}")

        if self.active_work_item:
            wi = self.active_work_item
            add(f"**Work item:** {wi.get('id', '?')} — {wi.get('title', '?')}")

        if self.plan_title and self.plan_steps:
            step_idx = self.current_step_index
            total = len(self.plan_steps)
            completed = sum(1 for s in self.plan_steps if s.status == "completed")
            current_desc = ""
            if 0 <= step_idx < total:
                current_desc = f" → {self.plan_steps[step_idx].description}"
            add(f"**Plan:** {self.plan_title} [{completed}/{total} done, step {step_idx + 1}{current_desc}]")

        if self.blockers:
            blocker_text = "; ".join(self.blockers[:3])
            add(f"**Blockers:** {blocker_text}")

        if self.next_steps:
            add(f"**Next:** {self.next_steps}")

        if chars_used >= max_chars:
            return "\n".join(parts)

        # --- Tier 2: Standard (decisions, files, progress) ---
        if self.progress_summary:
            add(f"**Progress:** {self.progress_summary}")

        if self.decisions:
            decision_lines = ["**Decisions made:**"]
            for d in self.decisions[:5]:
                dec = d.get("decision", "?")
                rationale = d.get("rationale", "")
                line = f"- {dec}"
                if rationale:
                    line += f" (reason: {rationale})"
                decision_lines.append(line)
            add("\n".join(decision_lines))

        if self.active_files:
            file_lines = ["**Active files:**"]
            for f in self.active_files[:10]:
                path = f.get("path", "?")
                role = f.get("role", "")
                line = f"- {path}"
                if role:
                    line += f" [{role}]"
                file_lines.append(line)
            add("\n".join(file_lines))

        if chars_used >= max_chars:
            return "\n".join(parts)

        # --- Tier 3: Full (all tasks, full plan steps, custom) ---
        if self.tasks:
            task_lines = ["**Tasks:**"]
            for t in self.tasks:
                status_icon = {"completed": "x", "in_progress": "~"}.get(
                    t.get("status", ""), " "
                )
                task_lines.append(f"- [{status_icon}] {t.get('content', '?')}")
            add("\n".join(task_lines))

        if self.plan_steps and len(self.plan_steps) > 1:
            plan_lines = [f"**Full plan — {self.plan_title}:**"]
            for i, step in enumerate(self.plan_steps):
                marker = "→" if i == self.current_step_index else " "
                status_icon = {"completed": "✓", "in_progress": "→", "skipped": "−"}.get(
                    step.status, " "
                )
                plan_lines.append(f"{marker} {i + 1}. [{status_icon}] {step.description}")
                if step.notes:
                    plan_lines.append(f"     Notes: {step.notes}")
            add("\n".join(plan_lines))

        if self.custom:
            custom_lines = ["**Additional context:**"]
            for k, v in self.custom.items():
                val_str = str(v)[:200]
                custom_lines.append(f"- {k}: {val_str}")
            add("\n".join(custom_lines))

        return "\n".join(parts)


class SessionStateManager:
    """Manages structured session state persistence.

    Provides save/load/merge operations for SessionState objects,
    backed by the session_states table in SQLite.
    """

    def __init__(self, db: Database):
        self.db = db
        self._ensure_table()

    def _ensure_table(self) -> None:
        """Create session_states table if it doesn't exist."""
        self.db.execute("""
            CREATE TABLE IF NOT EXISTS session_states (
                id TEXT PRIMARY KEY,
                session_id TEXT NOT NULL,
                compaction_number INTEGER NOT NULL,
                state_version INTEGER NOT NULL DEFAULT 1,
                state_data TEXT NOT NULL,
                captured_at TEXT DEFAULT (datetime('now')),
                FOREIGN KEY (session_id) REFERENCES sessions(id)
            )
        """)
        # Index for fast lookup
        self.db.execute("""
            CREATE INDEX IF NOT EXISTS idx_session_states_session
                ON session_states(session_id, compaction_number DESC)
        """)

    def save(self, state: SessionState) -> dict[str, Any]:
        """Save a session state snapshot.

        Args:
            state: The SessionState to persist.

        Returns:
            Dict with state_id, session_id, and captured_at.
        """
        state_id = f"ss-{uuid.uuid4().hex[:8]}"
        now = datetime.now(timezone.utc).isoformat()

        # Ensure captured_at is set
        if not state.captured_at:
            state.captured_at = now

        state_json = json.dumps(state.to_dict(), default=str)

        self.db.execute(
            """INSERT INTO session_states
               (id, session_id, compaction_number, state_version, state_data, captured_at)
               VALUES (?, ?, ?, ?, ?, ?)""",
            (state_id, state.session_id, state.compaction_number,
             STATE_VERSION, state_json, now),
        )

        # Enforce retention limit
        self._cleanup(state.session_id)

        return {
            "state_id": state_id,
            "session_id": state.session_id,
            "captured_at": now,
        }

    def load_latest(self, session_id: str) -> SessionState | None:
        """Load the most recent state for a session.

        Args:
            session_id: The session to load state for.

        Returns:
            SessionState if found, None otherwise.
        """
        row = self.db.execute_one(
            """SELECT state_data, state_version FROM session_states
               WHERE session_id = ?
               ORDER BY compaction_number DESC, captured_at DESC
               LIMIT 1""",
            (session_id,),
        )
        if not row:
            return None

        return self._deserialize(row["state_data"], row.get("state_version", 1))

    def load_by_compaction(
        self, session_id: str, compaction_number: int
    ) -> SessionState | None:
        """Load state for a specific compaction number.

        Args:
            session_id: The session ID.
            compaction_number: Which compaction's state to load.

        Returns:
            SessionState if found, None otherwise.
        """
        row = self.db.execute_one(
            """SELECT state_data, state_version FROM session_states
               WHERE session_id = ? AND compaction_number = ?
               ORDER BY captured_at DESC LIMIT 1""",
            (session_id, compaction_number),
        )
        if not row:
            return None

        return self._deserialize(row["state_data"], row.get("state_version", 1))

    def get_history(self, session_id: str) -> list[dict[str, Any]]:
        """Get metadata for all states in a session (without full state_data).

        Args:
            session_id: The session to query.

        Returns:
            List of dicts with id, compaction_number, state_version, captured_at.
        """
        rows = self.db.execute(
            """SELECT id, compaction_number, state_version, captured_at
               FROM session_states
               WHERE session_id = ?
               ORDER BY compaction_number DESC""",
            (session_id,),
        )
        return rows if isinstance(rows, list) else list(rows)

    def merge_conversation_state(
        self, state: SessionState, conversation_state: dict[str, Any]
    ) -> SessionState:
        """Merge regex-extracted conversation state into a SessionState.

        This bridges the existing capture_conversation_state() output
        into the new structured format, filling in fields that weren't
        explicitly set.

        Args:
            state: The SessionState to merge into (modified in place).
            conversation_state: Output from capture_conversation_state().

        Returns:
            The merged SessionState.
        """
        if not conversation_state:
            return state

        # Fill current_task if not explicitly set
        if not state.current_task and conversation_state.get("current_task"):
            state.current_task = conversation_state["current_task"]

        # Fill next_steps if not explicitly set
        if not state.next_steps and conversation_state.get("next_steps"):
            state.next_steps = conversation_state["next_steps"]

        # Merge file references (deduplicate)
        if conversation_state.get("recent_files"):
            existing_paths = {f.get("path") for f in state.active_files}
            for path in conversation_state["recent_files"]:
                if path not in existing_paths:
                    state.active_files.append({"path": path, "role": "referenced"})
                    existing_paths.add(path)

        return state

    def _deserialize(self, state_data: str, version: int) -> SessionState:
        """Deserialize state_data JSON, handling version migration.

        Args:
            state_data: JSON string of state data.
            version: Schema version of the stored data.

        Returns:
            SessionState instance.
        """
        data = json.loads(state_data)

        # Version migration: add new fields with defaults for older versions
        # Currently only version 1 exists, but this is the extension point
        if version < STATE_VERSION:
            pass  # Future: migrate older schema versions here

        return SessionState.from_dict(data)

    def _cleanup(self, session_id: str) -> None:
        """Enforce retention limit: keep only last MAX_STATES_PER_SESSION states."""
        rows = self.db.execute(
            """SELECT id FROM session_states
               WHERE session_id = ?
               ORDER BY compaction_number DESC, captured_at DESC""",
            (session_id,),
        )
        if isinstance(rows, list):
            ids_to_delete = [r["id"] for r in rows[MAX_STATES_PER_SESSION:]]
        else:
            all_rows = list(rows)
            ids_to_delete = [r["id"] for r in all_rows[MAX_STATES_PER_SESSION:]]

        for state_id in ids_to_delete:
            self.db.execute(
                "DELETE FROM session_states WHERE id = ?", (state_id,)
            )
