"""
Session Replay & Agent Replay - "Save States" for AI Agents.

NO ONE in the agent optimization space has this yet.

When an agent fails:
1. Record the exact state at every step (checkpoint)
2. Let users "replay" from any checkpoint with a fix applied
3. The fix becomes a pattern for future runs

Think of it as save states / time-travel debugging for AI agents.

Use cases:
- "The agent failed at step 4. Let me replay from step 3 with a different tool"
- "I want to see what would happen if the agent had more context at step 2"
- "Replay this failed run but with the prompt patch applied"
"""

import copy
import hashlib
import json
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional


@dataclass
class Checkpoint:
    """A snapshot of agent state at a specific step."""

    checkpoint_id: str
    trajectory_id: str
    step_index: int
    step_action: str
    step_tool: Optional[str]

    # State snapshot
    input_state: Dict[str, Any]  # What was fed to this step
    output_state: Optional[Any]  # What this step produced
    context: Dict[str, Any]  # Agent context at this point
    accumulated_tokens: int
    accumulated_latency_ms: int

    # For replay
    was_successful: bool
    error_detail: Optional[str] = None
    created_at: datetime = field(default_factory=datetime.utcnow)

    def to_dict(self) -> Dict:
        return {
            "checkpoint_id": self.checkpoint_id,
            "trajectory_id": self.trajectory_id,
            "step_index": self.step_index,
            "step_action": self.step_action,
            "step_tool": self.step_tool,
            "input_state": self.input_state,
            "output_state": self.output_state,
            "context": self.context,
            "accumulated_tokens": self.accumulated_tokens,
            "accumulated_latency_ms": self.accumulated_latency_ms,
            "was_successful": self.was_successful,
            "error_detail": self.error_detail,
            "created_at": self.created_at.isoformat(),
        }


@dataclass
class ReplayRequest:
    """A request to replay from a checkpoint with modifications."""

    replay_id: str
    source_trajectory_id: str
    checkpoint_id: str
    from_step_index: int

    # What's different in the replay
    modifications: Dict[str, Any] = field(default_factory=dict)
    # e.g., {"prompt_patch": "Always verify before calling",
    #        "replace_tool": {"old": "search", "new": "web_search"},
    #        "add_context": "The API requires auth headers"}

    created_at: datetime = field(default_factory=datetime.utcnow)

    def to_dict(self) -> Dict:
        return {
            "replay_id": self.replay_id,
            "source_trajectory_id": self.source_trajectory_id,
            "checkpoint_id": self.checkpoint_id,
            "from_step_index": self.from_step_index,
            "modifications": self.modifications,
            "created_at": self.created_at.isoformat(),
        }


@dataclass
class ReplayResult:
    """Result of a replay attempt."""

    replay_id: str
    source_trajectory_id: str
    new_trajectory_id: Optional[str]
    from_step_index: int
    modifications_applied: Dict[str, Any]
    outcome: str  # "success", "failure", "partial"
    improvement: Optional[str] = None  # "The replay succeeded where original failed"
    steps_replayed: int = 0
    pattern_created: bool = False  # Did this replay create a new pattern?

    def to_dict(self) -> Dict:
        return {
            "replay_id": self.replay_id,
            "source_trajectory_id": self.source_trajectory_id,
            "new_trajectory_id": self.new_trajectory_id,
            "from_step_index": self.from_step_index,
            "modifications_applied": self.modifications_applied,
            "outcome": self.outcome,
            "improvement": self.improvement,
            "steps_replayed": self.steps_replayed,
            "pattern_created": self.pattern_created,
        }


# =============================================================================
# Session Recorder
# =============================================================================


class SessionRecorder:
    """
    Records checkpoints during agent execution.

    Automatically creates a checkpoint at each step so the user
    can replay from any point if the agent fails.
    """

    def __init__(self, max_checkpoints_per_trajectory: int = 50):
        self.max_checkpoints = max_checkpoints_per_trajectory
        # trajectory_id → [checkpoints]
        self._checkpoints: Dict[str, List[Checkpoint]] = {}
        # Store recent trajectories for replay (LRU-style)
        self._trajectory_steps: Dict[str, List[Dict]] = {}
        self._max_stored_trajectories = 100

    def record_checkpoint(
        self,
        trajectory_id: str,
        step_index: int,
        step_action: str,
        step_tool: Optional[str],
        input_state: Dict[str, Any],
        output_state: Optional[Any],
        context: Dict[str, Any],
        accumulated_tokens: int,
        accumulated_latency_ms: int,
        was_successful: bool,
        error_detail: Optional[str] = None,
    ) -> Checkpoint:
        """Record a checkpoint at this step."""
        checkpoint_id = f"cp_{trajectory_id}_{step_index}"

        checkpoint = Checkpoint(
            checkpoint_id=checkpoint_id,
            trajectory_id=trajectory_id,
            step_index=step_index,
            step_action=step_action,
            step_tool=step_tool,
            input_state=_safe_copy(input_state),
            output_state=_safe_copy(output_state),
            context=_safe_copy(context),
            accumulated_tokens=accumulated_tokens,
            accumulated_latency_ms=accumulated_latency_ms,
            was_successful=was_successful,
            error_detail=error_detail,
        )

        if trajectory_id not in self._checkpoints:
            self._checkpoints[trajectory_id] = []

        # Limit checkpoints per trajectory
        if len(self._checkpoints[trajectory_id]) < self.max_checkpoints:
            self._checkpoints[trajectory_id].append(checkpoint)

        return checkpoint

    def get_checkpoints(self, trajectory_id: str) -> List[Checkpoint]:
        """Get all checkpoints for a trajectory."""
        return self._checkpoints.get(trajectory_id, [])

    def get_checkpoint(self, trajectory_id: str, step_index: int) -> Optional[Checkpoint]:
        """Get checkpoint at a specific step."""
        checkpoints = self._checkpoints.get(trajectory_id, [])
        for cp in checkpoints:
            if cp.step_index == step_index:
                return cp
        return None

    def get_failure_checkpoint(self, trajectory_id: str) -> Optional[Checkpoint]:
        """Get the checkpoint where failure occurred."""
        checkpoints = self._checkpoints.get(trajectory_id, [])
        for cp in reversed(checkpoints):
            if not cp.was_successful:
                return cp
        return None

    def store_trajectory_steps(self, trajectory_id: str, steps: List[Dict]) -> None:
        """Store step data for replay capability."""
        self._trajectory_steps[trajectory_id] = _safe_copy(steps)
        # Evict oldest if over limit
        if len(self._trajectory_steps) > self._max_stored_trajectories:
            oldest_key = next(iter(self._trajectory_steps))
            del self._trajectory_steps[oldest_key]

    def get_trajectory_steps(self, trajectory_id: str) -> Optional[List[Dict]]:
        """Get stored steps for replay."""
        return self._trajectory_steps.get(trajectory_id)

    def clear_trajectory(self, trajectory_id: str) -> None:
        """Clear checkpoints for a trajectory (after successful completion)."""
        self._checkpoints.pop(trajectory_id, None)

    def get_stats(self) -> Dict:
        """Get recorder stats."""
        all_checkpoints = sum(len(cps) for cps in self._checkpoints.values())
        return {
            "trajectories_tracked": len(self._checkpoints),
            "total_checkpoints": all_checkpoints,
            "trajectories_replayable": len(self._trajectory_steps),
        }

    def export_checkpoints(self, trajectory_id: str) -> List[Dict]:
        """Export checkpoints for server storage."""
        return [cp.to_dict() for cp in self._checkpoints.get(trajectory_id, [])]


# =============================================================================
# Replay Engine
# =============================================================================


class ReplayEngine:
    """
    Enables replaying agent executions from checkpoints.

    This is not a full agent re-execution engine — it prepares the
    replay state and modifications so the SDK can re-run the agent
    with the right context and patches applied.
    """

    def __init__(self, recorder: SessionRecorder):
        self.recorder = recorder
        self._replay_history: List[ReplayResult] = []

    def prepare_replay(
        self,
        trajectory_id: str,
        from_step: int,
        modifications: Dict[str, Any],
    ) -> Optional[ReplayRequest]:
        """
        Prepare a replay from a specific step with modifications.

        Modifications can include:
        - prompt_patch: str — inject instruction into prompt
        - replace_tool: {old: str, new: str} — swap tool at failure point
        - add_context: str — add extra context
        - skip_step: bool — skip the failing step entirely
        - retry_with_backoff: bool — retry with exponential backoff
        """
        checkpoint = self.recorder.get_checkpoint(trajectory_id, from_step)
        if not checkpoint:
            return None

        replay_id = (
            "replay_"
            + hashlib.md5(
                f"{trajectory_id}:{from_step}:{json.dumps(modifications, sort_keys=True, default=str)}".encode()
            ).hexdigest()[:12]
        )

        return ReplayRequest(
            replay_id=replay_id,
            source_trajectory_id=trajectory_id,
            checkpoint_id=checkpoint.checkpoint_id,
            from_step_index=from_step,
            modifications=modifications,
        )

    def get_replay_context(self, replay_request: ReplayRequest) -> Optional[Dict]:
        """
        Build the full context for replay execution.

        Returns everything the SDK needs to re-run the agent from the checkpoint:
        - Previous steps (pre-failure)
        - Modified state
        - Prompt patches to apply
        """
        trajectory_id = replay_request.source_trajectory_id
        steps = self.recorder.get_trajectory_steps(trajectory_id)
        checkpoint = self.recorder.get_checkpoint(trajectory_id, replay_request.from_step_index)

        if not checkpoint:
            return None

        # Build pre-failure steps (everything before the replay point)
        pre_steps = []
        if steps:
            pre_steps = steps[: replay_request.from_step_index]

        return {
            "replay_id": replay_request.replay_id,
            "source_trajectory_id": trajectory_id,
            "from_step_index": replay_request.from_step_index,
            "checkpoint_state": checkpoint.input_state,
            "checkpoint_context": checkpoint.context,
            "pre_steps": pre_steps,
            "accumulated_tokens": checkpoint.accumulated_tokens,
            "accumulated_latency_ms": checkpoint.accumulated_latency_ms,
            "modifications": replay_request.modifications,
            "prompt_patch": replay_request.modifications.get("prompt_patch"),
            "replace_tool": replay_request.modifications.get("replace_tool"),
            "add_context": replay_request.modifications.get("add_context"),
        }

    def record_replay_result(
        self,
        replay_request: ReplayRequest,
        new_trajectory_id: Optional[str],
        outcome: str,
        steps_replayed: int,
        original_outcome: str = "failure",
    ) -> ReplayResult:
        """Record the result of a replay attempt."""
        improvement = None
        if outcome == "success" and original_outcome == "failure":
            improvement = "Replay succeeded where original failed"
        elif outcome == "failure" and original_outcome == "failure":
            improvement = "Replay also failed — different fix needed"

        result = ReplayResult(
            replay_id=replay_request.replay_id,
            source_trajectory_id=replay_request.source_trajectory_id,
            new_trajectory_id=new_trajectory_id,
            from_step_index=replay_request.from_step_index,
            modifications_applied=replay_request.modifications,
            outcome=outcome,
            improvement=improvement,
            steps_replayed=steps_replayed,
            pattern_created=(outcome == "success"),
        )

        self._replay_history.append(result)
        return result

    def get_replay_history(self, trajectory_id: str = None) -> List[Dict]:
        """Get replay history, optionally filtered by source trajectory."""
        results = self._replay_history
        if trajectory_id:
            results = [r for r in results if r.source_trajectory_id == trajectory_id]
        return [r.to_dict() for r in results]

    def get_replay_stats(self) -> Dict:
        """Get replay stats for dashboard."""
        total = len(self._replay_history)
        successes = sum(1 for r in self._replay_history if r.outcome == "success")
        return {
            "total_replays": total,
            "successful_replays": successes,
            "replay_success_rate": round(successes / total, 3) if total > 0 else 0.0,
            "patterns_created": sum(1 for r in self._replay_history if r.pattern_created),
        }


# =============================================================================
# Helper Functions
# =============================================================================


def _safe_copy(obj: Any) -> Any:
    """Deep copy that handles non-copyable objects gracefully."""
    try:
        return copy.deepcopy(obj)
    except (TypeError, copy.Error):
        # Fall back to JSON round-trip for complex objects
        try:
            return json.loads(json.dumps(obj, default=str))
        except (TypeError, ValueError):
            return str(obj)
