"""Sandbox checkpoint and rollback manager for CleverAgents.

Provides snapshot-based checkpoint/rollback hooks that preserve sandbox state
at key points during plan execute/apply flows.  When a failure occurs, the
``CheckpointManager`` can restore a sandbox to a previously captured
checkpoint.

Stage M4 of the implementation plan.
"""

from __future__ import annotations

import os
import shutil
import tempfile
import threading
from datetime import datetime
from typing import Any, Protocol, runtime_checkable

import structlog
from pydantic import BaseModel, ConfigDict, Field
from ulid import ULID

from cleveragents.infrastructure.sandbox.protocol import (
    SandboxError,
)

logger = structlog.get_logger(__name__)


# ---------------------------------------------------------------------------
# Checkpointable protocol
# ---------------------------------------------------------------------------


@runtime_checkable
class Checkpointable(Protocol):
    """Minimal protocol for objects that can be checkpointed.

    Any object with a ``sandbox_id`` property and an optional ``context``
    providing ``sandbox_path`` can be used with :class:`CheckpointManager`.
    This is satisfied by the full ``Sandbox`` protocol as well as
    lightweight proxy objects.
    """

    @property
    def sandbox_id(self) -> str:
        """Unique identifier for the sandbox."""
        ...

    @property
    def context(self) -> Any:
        """Context with ``sandbox_path`` attribute, or ``None``."""
        ...


# ---------------------------------------------------------------------------
# Checkpoint data model
# ---------------------------------------------------------------------------


class SandboxCheckpoint(BaseModel):
    """Immutable record of a sandbox state snapshot.

    Captures the sandbox working directory state at a specific point
    in time, annotated with plan lifecycle metadata.
    """

    checkpoint_id: str = Field(
        ..., description="Unique ULID identifier for this checkpoint"
    )
    sandbox_id: str = Field(..., description="ID of the sandbox that was checkpointed")
    plan_id: str = Field(..., description="ID of the plan that owns the sandbox")
    phase: str = Field(
        ..., description="Lifecycle phase (pre_execute/post_execute/pre_apply)"
    )
    created_at: datetime = Field(..., description="When this checkpoint was captured")
    metadata: dict[str, str] = Field(
        default_factory=dict,
        description="Arbitrary key-value metadata (status, reason, etc.)",
    )
    snapshot_path: str = Field(
        ..., description="Filesystem path to the snapshot directory"
    )

    model_config = ConfigDict(frozen=True)


# ---------------------------------------------------------------------------
# Checkpoint manager
# ---------------------------------------------------------------------------


class CheckpointManager:
    """Manages sandbox checkpoint lifecycle: create, rollback, list, delete.

    Thread-safe: all mutable state is protected by ``_lock``.

    Usage::

        mgr = CheckpointManager()
        cp = mgr.create_checkpoint(sandbox, plan_id, "pre_execute", {})
        # ... execute phase runs ...
        if failure:
            mgr.rollback_to(cp)

        mgr.list_checkpoints(sandbox.sandbox_id)
        mgr.delete_checkpoint(cp.checkpoint_id)
    """

    def __init__(self) -> None:
        """Initialise the checkpoint manager."""
        self._checkpoints: dict[str, list[SandboxCheckpoint]] = {}
        self._lock: threading.RLock = threading.RLock()
        self._logger = logger.bind(component="checkpoint_manager")

    def create_checkpoint(
        self,
        sandbox: Checkpointable,
        plan_id: str,
        phase: str,
        metadata: dict[str, Any] | None = None,
    ) -> SandboxCheckpoint:
        """Snapshot the current sandbox state as a checkpoint.

        For sandboxes with a filesystem path (``context.sandbox_path``),
        the directory tree is copied to a temporary snapshot location.
        For sandboxes without a physical path (e.g. ``NoSandbox``), a
        lightweight metadata-only checkpoint is created.

        Args:
            sandbox: The sandbox to checkpoint.
            plan_id: The plan that owns the sandbox.
            phase: Lifecycle phase label (e.g. ``"pre_execute"``).
            metadata: Optional key-value metadata to attach.

        Returns:
            A frozen :class:`SandboxCheckpoint` record.

        Raises:
            SandboxError: If the snapshot copy fails.
        """
        meta = {k: str(v) for k, v in (metadata or {}).items()}
        checkpoint_id = str(ULID())
        sandbox_id = sandbox.sandbox_id
        sandbox_path: str | None = None

        if sandbox.context is not None:
            sandbox_path = sandbox.context.sandbox_path

        snapshot_path = self._snapshot_directory(sandbox_path, checkpoint_id)

        checkpoint = SandboxCheckpoint(
            checkpoint_id=checkpoint_id,
            sandbox_id=sandbox_id,
            plan_id=plan_id,
            phase=phase,
            created_at=datetime.now(),
            metadata=meta,
            snapshot_path=snapshot_path,
        )

        with self._lock:
            self._checkpoints.setdefault(sandbox_id, []).append(checkpoint)

        self._logger.info(
            "Checkpoint created",
            checkpoint_id=checkpoint_id,
            sandbox_id=sandbox_id,
            plan_id=plan_id,
            phase=phase,
        )

        return checkpoint

    def rollback_to(self, checkpoint: SandboxCheckpoint) -> bool:
        """Restore a sandbox to the state captured in *checkpoint*.

        Copies the snapshot directory back over the sandbox working
        directory.  Only works when the sandbox has a physical path
        and the snapshot directory exists.

        Args:
            checkpoint: The checkpoint to restore.

        Returns:
            ``True`` if the rollback succeeded, ``False`` otherwise.
        """
        snapshot = checkpoint.snapshot_path
        if not snapshot or not os.path.isdir(snapshot):
            self._logger.warning(
                "Rollback skipped: snapshot path missing or not a directory",
                checkpoint_id=checkpoint.checkpoint_id,
                snapshot_path=snapshot,
            )
            return False

        # Find sandbox path from the checkpoint's snapshot parent structure.
        # The snapshot is at <tmpdir>/ca-checkpoint-<id>/snapshot.
        # We need the *actual* sandbox path, which we embed in metadata.
        sandbox_path = checkpoint.metadata.get("sandbox_path")
        if not sandbox_path or not os.path.isdir(sandbox_path):
            self._logger.warning(
                "Rollback skipped: sandbox path unknown or missing",
                checkpoint_id=checkpoint.checkpoint_id,
            )
            return False

        try:
            # Clear the current sandbox contents and restore from snapshot
            for entry in os.listdir(sandbox_path):
                entry_path = os.path.join(sandbox_path, entry)
                if os.path.isdir(entry_path):
                    shutil.rmtree(entry_path)
                else:
                    os.remove(entry_path)

            # Copy snapshot contents into sandbox
            for entry in os.listdir(snapshot):
                src = os.path.join(snapshot, entry)
                dst = os.path.join(sandbox_path, entry)
                if os.path.isdir(src):
                    shutil.copytree(src, dst, symlinks=True)
                else:
                    shutil.copy2(src, dst)

        except OSError as exc:
            self._logger.error(
                "Rollback failed",
                checkpoint_id=checkpoint.checkpoint_id,
                error=str(exc),
            )
            return False

        self._logger.info(
            "Rollback completed",
            checkpoint_id=checkpoint.checkpoint_id,
            sandbox_id=checkpoint.sandbox_id,
        )
        return True

    def list_checkpoints(self, sandbox_id: str) -> list[SandboxCheckpoint]:
        """Return all checkpoints for a sandbox in creation order.

        Args:
            sandbox_id: The sandbox identifier.

        Returns:
            List of checkpoints ordered by creation time.
        """
        with self._lock:
            return list(self._checkpoints.get(sandbox_id, []))

    def delete_checkpoint(self, checkpoint_id: str) -> bool:
        """Remove a checkpoint and its snapshot data.

        Args:
            checkpoint_id: The checkpoint ULID to delete.

        Returns:
            ``True`` if the checkpoint was found and deleted,
            ``False`` if it was not found.
        """
        with self._lock:
            for sandbox_id, checkpoints in self._checkpoints.items():
                for i, cp in enumerate(checkpoints):
                    if cp.checkpoint_id == checkpoint_id:
                        removed = checkpoints.pop(i)
                        self._cleanup_snapshot(removed.snapshot_path)
                        self._logger.info(
                            "Checkpoint deleted",
                            checkpoint_id=checkpoint_id,
                            sandbox_id=sandbox_id,
                        )
                        return True

        self._logger.debug(
            "Checkpoint not found for deletion",
            checkpoint_id=checkpoint_id,
        )
        return False

    # -- internal helpers ----------------------------------------------------

    def _snapshot_directory(
        self,
        sandbox_path: str | None,
        checkpoint_id: str,
    ) -> str:
        """Copy the sandbox working directory to a snapshot location.

        Args:
            sandbox_path: Path to the sandbox directory, or ``None``
                for metadata-only checkpoints.
            checkpoint_id: Used to name the snapshot directory.

        Returns:
            Absolute path to the snapshot directory.

        Raises:
            SandboxError: If the copy fails.
        """
        parent = tempfile.mkdtemp(prefix=f"ca-checkpoint-{checkpoint_id[:8]}-")
        snapshot = os.path.join(parent, "snapshot")

        if sandbox_path is not None and os.path.isdir(sandbox_path):
            try:
                shutil.copytree(
                    sandbox_path, snapshot, symlinks=True, dirs_exist_ok=False
                )
            except OSError as exc:
                shutil.rmtree(parent, ignore_errors=True)
                raise SandboxError(
                    f"Failed to snapshot sandbox directory: {exc}"
                ) from exc
        else:
            os.makedirs(snapshot, exist_ok=True)

        return snapshot

    @staticmethod
    def _cleanup_snapshot(snapshot_path: str) -> None:
        """Remove a snapshot directory and its parent temp dir."""
        if not snapshot_path:
            return
        parent = os.path.dirname(snapshot_path)
        if os.path.isdir(parent):
            shutil.rmtree(parent, ignore_errors=True)
