"""SQLite-based checkpoint manager for workflow state."""

from typing import Any
from datetime import datetime
from .persistence import DatabaseManager


class CheckpointManager:
    """Manages saving and loading workflow checkpoints to SQLite database."""

    def __init__(self, workflow_name: str, db_path: str = ".ombra_executions.db"):
        """Initialize checkpoint manager.

        Args:
            workflow_name: Name of the workflow
            db_path: Path to SQLite database file
        """
        self.workflow_name = workflow_name
        self.db = DatabaseManager(db_path)

    def save(self, step_name: str, inputs: Any, outputs: Any, metadata: dict | None = None) -> str:
        """Save a checkpoint for a step execution.

        Args:
            step_name: Name of the step
            inputs: Step input arguments
            outputs: Step output/return value
            metadata: Optional metadata dictionary

        Returns:
            Checkpoint ID
        """
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        checkpoint_id = f"{step_name}_{timestamp}"

        checkpoint_data = {
            "checkpoint_id": checkpoint_id,
            "workflow_name": self.workflow_name,
            "step_name": step_name,
            "timestamp": timestamp,
            "inputs": inputs,
            "outputs": outputs,
            "metadata": metadata or {},
            "execution_id": metadata.get("execution_id") if metadata else None
        }

        self.db.save_checkpoint(checkpoint_data)
        return checkpoint_id

    def load(self, checkpoint_id: str) -> dict:
        """Load a specific checkpoint by ID.

        Args:
            checkpoint_id: Checkpoint ID to load

        Returns:
            Dictionary with checkpoint data

        Raises:
            FileNotFoundError: If checkpoint not found
        """
        checkpoint = self.db.load_checkpoint(checkpoint_id)

        if not checkpoint:
            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_id}")

        return checkpoint

    def list_checkpoints(self, step_name: str | None = None) -> list[dict]:
        """List all checkpoints, optionally filtered by step name.

        Args:
            step_name: Optional step name filter

        Returns:
            List of checkpoint dictionaries
        """
        return self.db.list_checkpoints(self.workflow_name, step_name)

    def get_latest(self, step_name: str) -> dict | None:
        """Get the most recent checkpoint for a step.

        Args:
            step_name: Name of the step

        Returns:
            Dictionary with checkpoint data or None if not found
        """
        return self.db.get_latest_checkpoint(self.workflow_name, step_name)
