from __future__ import annotations

import json
import sqlite3
from pathlib import Path

from app.core.contracts import TaskCheckpoint


class SQLiteCheckpointStore:
    """Durable checkpoint store for resumable runtime state."""

    def __init__(self, db_path: str | Path) -> None:
        self._db_path = Path(db_path)
        self._db_path.parent.mkdir(parents=True, exist_ok=True)
        self._initialize()

    def save(self, checkpoint: TaskCheckpoint) -> TaskCheckpoint:
        with sqlite3.connect(self._db_path) as conn:
            conn.execute(
                """
                INSERT OR REPLACE INTO checkpoints (
                    task_id, status, active_step_id, plan_snapshot_json,
                    context_snapshot_json, updated_at
                ) VALUES (?, ?, ?, ?, ?, ?)
                """,
                (
                    checkpoint.task_id,
                    checkpoint.status,
                    checkpoint.active_step_id,
                    json.dumps(checkpoint.plan_snapshot, default=str),
                    json.dumps(checkpoint.context_snapshot, default=str),
                    checkpoint.updated_at.isoformat(),
                ),
            )
            conn.commit()
        return checkpoint

    def load(self, task_id: str) -> TaskCheckpoint | None:
        with sqlite3.connect(self._db_path) as conn:
            row = conn.execute(
                """
                SELECT task_id, status, active_step_id, plan_snapshot_json,
                       context_snapshot_json, updated_at
                FROM checkpoints
                WHERE task_id = ?
                """,
                (task_id,),
            ).fetchone()
        if not row:
            return None
        return TaskCheckpoint(
            task_id=row[0],
            status=row[1],
            active_step_id=row[2],
            plan_snapshot=json.loads(row[3]),
            context_snapshot=json.loads(row[4]),
            updated_at=row[5],
        )

    def _initialize(self) -> None:
        with sqlite3.connect(self._db_path) as conn:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS checkpoints (
                    task_id TEXT PRIMARY KEY,
                    status TEXT NOT NULL,
                    active_step_id TEXT,
                    plan_snapshot_json TEXT NOT NULL,
                    context_snapshot_json TEXT NOT NULL,
                    updated_at TEXT NOT NULL
                )
                """
            )
            conn.commit()