"""Tests for PostgreSQL-backed checkpoint stores."""

from __future__ import annotations

from logos_adk.crew_checkpoint_store import CrewCheckpointRecord, PostgreSQLCrewCheckpointStore
from logos_adk.graph_checkpoint_store import GraphCheckpointRecord, PostgreSQLGraphCheckpointStore


class _FakeCursor:
    def __init__(self, storage):
        self.storage = storage
        self._fetchone = None

    def execute(self, query, params=None):
        normalized = " ".join(str(query).split())
        self._fetchone = None

        if normalized.startswith("INSERT INTO graph_checkpoints"):
            self.storage["graph"][(params[1], params[0])] = {
                "correlation_id": params[0],
                "graph_name": params[1],
                "status": params[2],
                "created_at": params[3],
                "updated_at": params[4],
                "payload_json": params[5],
            }
            return

        if normalized.startswith("SELECT * FROM graph_checkpoints WHERE graph_name = %s"):
            self._fetchone = self.storage["graph"].get((params[0], params[1]))
            return

        if normalized.startswith("DELETE FROM graph_checkpoints"):
            self.storage["graph"].pop((params[0], params[1]), None)
            return

        if normalized.startswith("INSERT INTO crew_checkpoints"):
            self.storage["crew"][params[0]] = {
                "run_id": params[0],
                "team_name": params[1],
                "process": params[2],
                "status": params[3],
                "created_at": params[4],
                "updated_at": params[5],
                "payload_json": params[6],
            }
            return

        if normalized.startswith("SELECT * FROM crew_checkpoints WHERE run_id = %s"):
            self._fetchone = self.storage["crew"].get(params[0])
            return

        if normalized.startswith("DELETE FROM crew_checkpoints"):
            self.storage["crew"].pop(params[0], None)

    def fetchone(self):
        return self._fetchone

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


class _FakeConnection:
    def __init__(self):
        self.storage = {"graph": {}, "crew": {}}
        self.commits = 0

    def cursor(self):
        return _FakeCursor(self.storage)

    def commit(self):
        self.commits += 1

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


def test_postgresql_graph_checkpoint_store_round_trip_with_fake_backend(monkeypatch):
    connection = _FakeConnection()
    store = PostgreSQLGraphCheckpointStore("postgresql://user:pass@localhost/graph")
    store._connect = lambda: connection
    monkeypatch.setattr(PostgreSQLGraphCheckpointStore, "_ensure_schema", lambda self, conn: None)

    record = GraphCheckpointRecord(
        correlation_id="corr-1",
        graph_name="release_graph",
        status="paused",
        created_at="2026-01-01T00:00:00+00:00",
        updated_at="2026-01-01T00:00:05+00:00",
        payload={"step": "approval"},
    )
    store.save_checkpoint(record)

    loaded = store.load_checkpoint("release_graph", "corr-1")

    assert loaded is not None
    assert loaded.payload == {"step": "approval"}
    store.delete_checkpoint("release_graph", "corr-1")
    assert store.load_checkpoint("release_graph", "corr-1") is None


def test_postgresql_crew_checkpoint_store_round_trip_with_fake_backend(monkeypatch):
    connection = _FakeConnection()
    store = PostgreSQLCrewCheckpointStore("postgresql://user:pass@localhost/crew")
    store._connect = lambda: connection
    monkeypatch.setattr(PostgreSQLCrewCheckpointStore, "_ensure_schema", lambda self, conn: None)

    record = CrewCheckpointRecord(
        run_id="run-1",
        team_name="research-team",
        process="hierarchical",
        status="paused",
        created_at="2026-01-01T00:00:00+00:00",
        updated_at="2026-01-01T00:00:10+00:00",
        payload={"task": "approval"},
    )
    store.save_checkpoint(record)

    loaded = store.load_checkpoint("run-1")

    assert loaded is not None
    assert loaded.payload == {"task": "approval"}
    store.delete_checkpoint("run-1")
    assert store.load_checkpoint("run-1") is None
