"""Checkpoint management — save and restore execution state.

Manages checkpoint creation during plan execution. Applies configurable
rules to decide which steps get checkpoints (every N steps, on tool
changes, before risky operations). Supports save, restore, rollback,
and state diffing.
"""

import hashlib
import json
from dataclasses import dataclass, field
from typing import Any

from xyglue.plan.models import Step, Checkpoint


@dataclass
class CheckpointRules:
    """Rules for when to create checkpoints."""

    checkpoint_every_n_steps: int = 3
    checkpoint_on_tool_change: bool = True
    checkpoint_before_risky: bool = True
    checkpoint_first_step: bool = True
    checkpoint_last_step: bool = True
    risky_tools: list[str] = field(
        default_factory=lambda: ["database", "email", "files"]
    )
    risky_actions: list[str] = field(
        default_factory=lambda: ["delete", "drop", "truncate", "send", "publish"]
    )

    def should_checkpoint(
        self, step: Step, step_index: int, total_steps: int,
        prev_tool: str | None,
    ) -> bool:
        """Determine whether a step should get a checkpoint based on rules."""
        if step.checkpoint:
            return True  # Already marked
        if self.checkpoint_first_step and step_index == 0:
            return True
        if self.checkpoint_last_step and step_index == total_steps - 1:
            return True
        if (step_index + 1) % self.checkpoint_every_n_steps == 0:
            return True
        if self.checkpoint_on_tool_change and prev_tool and step.tool != prev_tool:
            return True
        if self.checkpoint_before_risky and step.tool in self.risky_tools:
            return True
        if self.checkpoint_before_risky:
            action_lower = step.action.lower()
            if any(r in action_lower for r in self.risky_actions):
                return True
        return False


class CheckpointManager:
    """Manages checkpoints during plan execution.

    Provides:
    - enrich_plan_checkpoints: add checkpoint flags to steps based on rules
    - save / get_latest / get_all / get_by_step: CRUD for checkpoints
    - rollback_to: find the right checkpoint to resume from
    - state_diff: compare two checkpoint states
    - integrity_check: verify checkpoint chain consistency
    """

    def __init__(self, rules: CheckpointRules | None = None) -> None:
        self.rules = rules or CheckpointRules()
        self._checkpoints: list[Checkpoint] = []

    def enrich_plan_checkpoints(self, steps: list[Step]) -> list[Step]:
        """Add checkpoint flags to steps based on rules."""
        prev_tool: str | None = None
        total = len(steps)
        for i, step in enumerate(steps):
            if self.rules.should_checkpoint(step, i, total, prev_tool):
                step.checkpoint = True
            prev_tool = step.tool
        return steps

    def save(
        self,
        step_number: int,
        state: dict[str, Any],
        outputs: list[Any] | None = None,
    ) -> Checkpoint:
        """Save a checkpoint at the given step."""
        cp = Checkpoint(
            step_number=step_number, state=state, outputs=outputs or []
        )
        self._checkpoints.append(cp)
        return cp

    def get_latest(self) -> Checkpoint | None:
        """Get the most recent checkpoint."""
        return self._checkpoints[-1] if self._checkpoints else None

    def get_by_step(self, step_number: int) -> Checkpoint | None:
        """Get checkpoint for a specific step number."""
        for cp in reversed(self._checkpoints):
            if cp.step_number == step_number:
                return cp
        return None

    def get_all(self) -> list[Checkpoint]:
        """Get all checkpoints in order."""
        return list(self._checkpoints)

    def rollback_to(self, step_number: int) -> Checkpoint | None:
        """Find the most recent checkpoint at or before the given step.

        Also trims all checkpoints after it so we have a clean resume point.
        """
        target: Checkpoint | None = None
        target_idx = -1
        for i, cp in enumerate(self._checkpoints):
            if cp.step_number <= step_number:
                target = cp
                target_idx = i
        if target is not None and target_idx >= 0:
            self._checkpoints = self._checkpoints[: target_idx + 1]
        return target

    def state_diff(self, cp_a: Checkpoint, cp_b: Checkpoint) -> dict[str, Any]:
        """Compare two checkpoint states and return differences."""
        diff: dict[str, Any] = {"added": {}, "removed": {}, "changed": {}}
        keys_a = set(cp_a.state.keys())
        keys_b = set(cp_b.state.keys())

        for key in keys_b - keys_a:
            diff["added"][key] = cp_b.state[key]
        for key in keys_a - keys_b:
            diff["removed"][key] = cp_a.state[key]
        for key in keys_a & keys_b:
            if cp_a.state[key] != cp_b.state[key]:
                diff["changed"][key] = {
                    "from": cp_a.state[key],
                    "to": cp_b.state[key],
                }

        return diff

    def integrity_check(self) -> list[str]:
        """Verify checkpoint chain is consistent. Returns issues."""
        issues: list[str] = []
        if not self._checkpoints:
            return issues

        prev_step = -1
        for i, cp in enumerate(self._checkpoints):
            if cp.step_number < prev_step:
                issues.append(
                    f"Checkpoint {i} step_number {cp.step_number} < "
                    f"previous {prev_step}"
                )
            prev_step = cp.step_number

        return issues

    def compute_state_hash(self, state: dict[str, Any]) -> str:
        """Compute a deterministic hash of a state dict for comparison."""
        serialized = json.dumps(state, sort_keys=True, default=str)
        return hashlib.sha256(serialized.encode()).hexdigest()[:16]

    def clear(self) -> None:
        """Remove all checkpoints."""
        self._checkpoints.clear()

    @property
    def count(self) -> int:
        return len(self._checkpoints)
