"""
Checkpoint — durable execution via state snapshots.

Design notes:
  - Every task execution can be checkpointed (snapshot state + conversation)
  - Supports pause/resume, failure recovery, and time-travel debugging
  - Pluggable storage backends: in-memory (default), filesystem, Redis, etc.
  - Each checkpoint is immutable once created
"""

from __future__ import annotations

import json
import logging
from abc import ABC, abstractmethod
from datetime import datetime
from pathlib import Path
from typing import Any

from pydantic import BaseModel, Field

from nexus.types import new_id

logger = logging.getLogger("nexus.runtime.checkpoint")


class Checkpoint(BaseModel):
    """
    An immutable snapshot of execution state at a point in time.

    Contains everything needed to resume a task from this point.
    """

    id: str = Field(default_factory=new_id)
    task_id: str
    agent_id: str = ""
    version: int = 0

    # State snapshot
    state_data: dict[str, Any] = Field(default_factory=dict)
    conversation_messages: list[dict[str, Any]] = Field(default_factory=list)

    # Execution context
    iteration: int = 0
    token_usage: dict[str, Any] = Field(default_factory=dict)
    status: str = "working"

    # Metadata
    created_at: datetime = Field(default_factory=datetime.utcnow)
    description: str = ""
    metadata: dict[str, Any] = Field(default_factory=dict)


class CheckpointStore(ABC):
    """Abstract base class for checkpoint storage backends."""

    @abstractmethod
    async def save(self, checkpoint: Checkpoint) -> str:
        """Save a checkpoint and return its ID."""
        ...

    @abstractmethod
    async def load(self, checkpoint_id: str) -> Checkpoint | None:
        """Load a checkpoint by ID."""
        ...

    @abstractmethod
    async def list_checkpoints(
        self,
        task_id: str | None = None,
        limit: int = 50,
    ) -> list[Checkpoint]:
        """List checkpoints, optionally filtered by task_id."""
        ...

    @abstractmethod
    async def delete(self, checkpoint_id: str) -> bool:
        """Delete a checkpoint. Returns True if it existed."""
        ...

    async def get_latest(self, task_id: str) -> Checkpoint | None:
        """Get the most recent checkpoint for a task."""
        checkpoints = await self.list_checkpoints(task_id=task_id, limit=1)
        return checkpoints[0] if checkpoints else None


class InMemoryCheckpointStore(CheckpointStore):
    """
    In-memory checkpoint store (for development and testing).

    Checkpoints are lost when the process exits.
    """

    def __init__(self) -> None:
        self._store: dict[str, Checkpoint] = {}

    async def save(self, checkpoint: Checkpoint) -> str:
        self._store[checkpoint.id] = checkpoint
        logger.debug(f"Checkpoint saved: {checkpoint.id} (task={checkpoint.task_id})")
        return checkpoint.id

    async def load(self, checkpoint_id: str) -> Checkpoint | None:
        return self._store.get(checkpoint_id)

    async def list_checkpoints(
        self,
        task_id: str | None = None,
        limit: int = 50,
    ) -> list[Checkpoint]:
        checkpoints = list(self._store.values())
        if task_id:
            checkpoints = [c for c in checkpoints if c.task_id == task_id]
        # Sort by creation time descending
        checkpoints.sort(key=lambda c: c.created_at, reverse=True)
        return checkpoints[:limit]

    async def delete(self, checkpoint_id: str) -> bool:
        return self._store.pop(checkpoint_id, None) is not None

    @property
    def size(self) -> int:
        return len(self._store)


class FileCheckpointStore(CheckpointStore):
    """
    Filesystem-based checkpoint store.

    Each checkpoint is stored as a JSON file in the specified directory.
    Suitable for local development and small-scale usage.
    """

    def __init__(self, directory: str | Path) -> None:
        self._dir = Path(directory)
        self._dir.mkdir(parents=True, exist_ok=True)

    def _path(self, checkpoint_id: str) -> Path:
        return self._dir / f"{checkpoint_id}.json"

    async def save(self, checkpoint: Checkpoint) -> str:
        path = self._path(checkpoint.id)
        path.write_text(checkpoint.model_dump_json(indent=2))
        logger.debug(f"Checkpoint saved to {path}")
        return checkpoint.id

    async def load(self, checkpoint_id: str) -> Checkpoint | None:
        path = self._path(checkpoint_id)
        if not path.exists():
            return None
        try:
            data = json.loads(path.read_text())
            return Checkpoint.model_validate(data)
        except Exception as exc:
            logger.error(f"Failed to load checkpoint {checkpoint_id}: {exc}")
            return None

    async def list_checkpoints(
        self,
        task_id: str | None = None,
        limit: int = 50,
    ) -> list[Checkpoint]:
        checkpoints = []
        for path in self._dir.glob("*.json"):
            try:
                data = json.loads(path.read_text())
                cp = Checkpoint.model_validate(data)
                if task_id and cp.task_id != task_id:
                    continue
                checkpoints.append(cp)
            except Exception:
                continue
        checkpoints.sort(key=lambda c: c.created_at, reverse=True)
        return checkpoints[:limit]

    async def delete(self, checkpoint_id: str) -> bool:
        path = self._path(checkpoint_id)
        if path.exists():
            path.unlink()
            return True
        return False
