"""Persistence adapters for durable workflow checkpoints.

The workflow runtime depends on a small StateStore contract rather than a specific
database. InMemoryStateStore keeps tests deterministic, while SQLiteStateStore
provides a zero-infrastructure durable store for the portfolio demo.
"""

from __future__ import annotations

import sqlite3
from pathlib import Path
from typing import Protocol

from src.schemas import WorkflowRun


class StateStoreError(RuntimeError):
    """Base error for workflow state-store operations."""


class RunNotFoundError(StateStoreError):
    """Raised when a requested workflow run does not exist."""


class StateStore(Protocol):
    """Minimal persistence contract used by the workflow runtime."""

    def save(self, run: WorkflowRun) -> None:
        """Persist the latest complete checkpoint for a workflow run."""

    def load(self, run_id: str) -> WorkflowRun:
        """Load a detached workflow run snapshot by identifier."""


class InMemoryStateStore:
    """Deterministic state store for tests and local execution."""

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

    def save(self, run: WorkflowRun) -> None:
        self._runs[run.run_id] = run.model_dump_json()

    def load(self, run_id: str) -> WorkflowRun:
        try:
            payload = self._runs[run_id]
        except KeyError as exc:
            raise RunNotFoundError(f"workflow run '{run_id}' was not found") from exc
        return WorkflowRun.model_validate_json(payload)


class SQLiteStateStore:
    """SQLite-backed state store using one JSON snapshot per workflow run."""

    def __init__(self, database_path: str | Path) -> None:
        self.database_path = str(database_path)
        self._initialize()

    def _connect(self) -> sqlite3.Connection:
        return sqlite3.connect(self.database_path)

    def _initialize(self) -> None:
        with self._connect() as connection:
            connection.execute(
                """
                CREATE TABLE IF NOT EXISTS workflow_runs (
                    run_id TEXT PRIMARY KEY,
                    payload TEXT NOT NULL
                )
                """
            )

    def save(self, run: WorkflowRun) -> None:
        payload = run.model_dump_json()
        with self._connect() as connection:
            connection.execute(
                """
                INSERT INTO workflow_runs (run_id, payload)
                VALUES (?, ?)
                ON CONFLICT(run_id) DO UPDATE SET payload = excluded.payload
                """,
                (run.run_id, payload),
            )

    def load(self, run_id: str) -> WorkflowRun:
        with self._connect() as connection:
            row = connection.execute(
                "SELECT payload FROM workflow_runs WHERE run_id = ?",
                (run_id,),
            ).fetchone()

        if row is None:
            raise RunNotFoundError(f"workflow run '{run_id}' was not found")

        return WorkflowRun.model_validate_json(row[0])