"""Persistent checkpoint store for paused crew runs."""
from __future__ import annotations

from dataclasses import dataclass
from datetime import UTC, datetime
import json
from pathlib import Path
import sqlite3
from typing import Any

from logos_adk.checkpoint_json import JSONObject, JSONValue, clone_json_object, ensure_json_object
from logos_adk.errors import StateBackendError
from logos_adk.persistence import (
    get_postgresql_schema_version,
    get_sqlite_schema_version,
    set_postgresql_schema_version,
    set_sqlite_schema_version,
)
from logos_adk.postgres_utils import import_psycopg, resolve_database_url


CREW_CHECKPOINT_SCHEMA_KIND = "crew_checkpoint_store"
CREW_CHECKPOINT_SCHEMA_VERSION = 1


def _utcnow() -> str:
    return datetime.now(UTC).isoformat()


@dataclass
class CrewCheckpointRecord:
    """Serialized paused crew run snapshot."""

    run_id: str
    team_name: str
    process: str
    status: str
    created_at: str
    updated_at: str
    payload: JSONObject


class CrewCheckpointStore:
    """Abstract persistent store for paused crew runs."""

    def save_checkpoint(self, record: CrewCheckpointRecord) -> None:
        raise NotImplementedError

    def load_checkpoint(self, run_id: str) -> CrewCheckpointRecord | None:
        raise NotImplementedError

    def delete_checkpoint(self, run_id: str) -> None:
        raise NotImplementedError


class InMemoryCrewCheckpointStore(CrewCheckpointStore):
    """In-memory checkpoint store for tests and local development."""

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

    def save_checkpoint(self, record: CrewCheckpointRecord) -> None:
        payload = clone_json_object(record.payload, path="checkpoint.payload")
        self._items[record.run_id] = CrewCheckpointRecord(
            run_id=record.run_id,
            team_name=record.team_name,
            process=record.process,
            status=record.status,
            created_at=record.created_at,
            updated_at=record.updated_at,
            payload=payload,
        )

    def load_checkpoint(self, run_id: str) -> CrewCheckpointRecord | None:
        record = self._items.get(run_id)
        if record is None:
            return None
        payload = clone_json_object(record.payload, path="checkpoint.payload")
        return CrewCheckpointRecord(
            run_id=record.run_id,
            team_name=record.team_name,
            process=record.process,
            status=record.status,
            created_at=record.created_at,
            updated_at=record.updated_at,
            payload=payload,
        )

    def delete_checkpoint(self, run_id: str) -> None:
        self._items.pop(run_id, None)


class SQLiteCrewCheckpointStore(CrewCheckpointStore):
    """SQLite-backed checkpoint store for paused crew runs."""

    def __init__(self, path: str = ".logos_adk/crew_checkpoints.db", table_name: str = "crew_checkpoints") -> None:
        self.path = path
        self.table_name = table_name

    def _connect(self) -> sqlite3.Connection:
        try:
            Path(self.path).parent.mkdir(parents=True, exist_ok=True)
            connection = sqlite3.connect(self.path)
            connection.row_factory = sqlite3.Row
            return connection
        except Exception as exc:
            raise StateBackendError(f"Failed to open SQLite crew checkpoint store: {exc}") from exc

    def _ensure_schema(self, connection: sqlite3.Connection) -> None:
        connection.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {self.table_name} (
                run_id TEXT PRIMARY KEY,
                team_name TEXT NOT NULL,
                process TEXT NOT NULL,
                status TEXT NOT NULL,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL,
                payload_json TEXT NOT NULL
            )
            """
        )
        connection.execute(
            f"""
            CREATE INDEX IF NOT EXISTS idx_{self.table_name}_team_updated
            ON {self.table_name} (team_name, updated_at DESC)
            """
        )
        recorded_version = get_sqlite_schema_version(connection, CREW_CHECKPOINT_SCHEMA_KIND)
        if recorded_version != CREW_CHECKPOINT_SCHEMA_VERSION:
            set_sqlite_schema_version(connection, CREW_CHECKPOINT_SCHEMA_KIND, CREW_CHECKPOINT_SCHEMA_VERSION)
        connection.commit()

    def save_checkpoint(self, record: CrewCheckpointRecord) -> None:
        payload_json = json.dumps(
            ensure_json_object(record.payload, path="checkpoint.payload"),
            ensure_ascii=False,
            allow_nan=False,
        )
        with self._connect() as connection:
            self._ensure_schema(connection)
            connection.execute(
                f"""
                INSERT INTO {self.table_name}
                (run_id, team_name, process, status, created_at, updated_at, payload_json)
                VALUES (?, ?, ?, ?, ?, ?, ?)
                ON CONFLICT(run_id) DO UPDATE SET
                    team_name = excluded.team_name,
                    process = excluded.process,
                    status = excluded.status,
                    created_at = excluded.created_at,
                    updated_at = excluded.updated_at,
                    payload_json = excluded.payload_json
                """,
                (
                    record.run_id,
                    record.team_name,
                    record.process,
                    record.status,
                    record.created_at,
                    record.updated_at,
                    payload_json,
                ),
            )
            connection.commit()

    def load_checkpoint(self, run_id: str) -> CrewCheckpointRecord | None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            row = connection.execute(
                f"SELECT * FROM {self.table_name} WHERE run_id = ?",
                (run_id,),
            ).fetchone()
        if row is None:
            return None
        try:
            payload = ensure_json_object(
                json.loads(row["payload_json"]),
                path="checkpoint.payload",
            )
        except (TypeError, ValueError, json.JSONDecodeError) as exc:
            raise StateBackendError(
                f"Failed to load SQLite crew checkpoint '{run_id}': invalid payload: {exc}"
            ) from exc
        return CrewCheckpointRecord(
            run_id=row["run_id"],
            team_name=row["team_name"],
            process=row["process"],
            status=row["status"],
            created_at=row["created_at"],
            updated_at=row["updated_at"],
            payload=payload,
        )

    def delete_checkpoint(self, run_id: str) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            connection.execute(
                f"DELETE FROM {self.table_name} WHERE run_id = ?",
                (run_id,),
            )
            connection.commit()


class PostgreSQLCrewCheckpointStore(CrewCheckpointStore):
    """PostgreSQL-backed checkpoint store for paused crew runs."""

    def __init__(
        self,
        dsn: str | None = None,
        table_name: str = "crew_checkpoints",
    ) -> None:
        self.dsn = resolve_database_url(dsn)
        if not self.dsn:
            raise ValueError(
                "PostgreSQL crew checkpoint store requires a DSN or DATABASE_URL/LOGOS_ADK_DATABASE_URL"
            )
        self.table_name = table_name

    def _connect(self) -> Any:
        try:
            return import_psycopg().connect(self.dsn)
        except Exception as exc:
            raise StateBackendError(f"Failed to open PostgreSQL crew checkpoint store: {exc}") from exc

    def _ensure_schema(self, connection: Any) -> None:
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    run_id TEXT PRIMARY KEY,
                    team_name TEXT NOT NULL,
                    process TEXT NOT NULL,
                    status TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL,
                    payload_json JSONB NOT NULL
                )
                """
            )
            cursor.execute(
                f"""
                CREATE INDEX IF NOT EXISTS idx_{self.table_name}_team_updated
                ON {self.table_name} (team_name, updated_at DESC)
                """
            )
        recorded_version = get_postgresql_schema_version(connection, CREW_CHECKPOINT_SCHEMA_KIND)
        if recorded_version != CREW_CHECKPOINT_SCHEMA_VERSION:
            set_postgresql_schema_version(connection, CREW_CHECKPOINT_SCHEMA_KIND, CREW_CHECKPOINT_SCHEMA_VERSION)
        connection.commit()

    def save_checkpoint(self, record: CrewCheckpointRecord) -> None:
        payload_json = json.dumps(
            ensure_json_object(record.payload, path="checkpoint.payload"),
            ensure_ascii=False,
            allow_nan=False,
        )
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"""
                    INSERT INTO {self.table_name}
                    (run_id, team_name, process, status, created_at, updated_at, payload_json)
                    VALUES (%s, %s, %s, %s, %s, %s, %s::jsonb)
                    ON CONFLICT(run_id) DO UPDATE SET
                        team_name = excluded.team_name,
                        process = excluded.process,
                        status = excluded.status,
                        created_at = excluded.created_at,
                        updated_at = excluded.updated_at,
                        payload_json = excluded.payload_json
                    """,
                    (
                        record.run_id,
                        record.team_name,
                        record.process,
                        record.status,
                        record.created_at,
                        record.updated_at,
                        payload_json,
                    ),
                )
            connection.commit()

    def load_checkpoint(self, run_id: str) -> CrewCheckpointRecord | None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"SELECT * FROM {self.table_name} WHERE run_id = %s",
                    (run_id,),
                )
                row = cursor.fetchone()
        if row is None:
            return None
        try:
            raw_payload = _row_value(row, "payload_json", 6)
            payload = ensure_json_object(
                json.loads(raw_payload) if isinstance(raw_payload, str) else raw_payload,
                path="checkpoint.payload",
            )
        except (TypeError, ValueError, json.JSONDecodeError) as exc:
            raise StateBackendError(
                f"Failed to load PostgreSQL crew checkpoint '{run_id}': invalid payload: {exc}"
            ) from exc
        return CrewCheckpointRecord(
            run_id=str(_row_value(row, "run_id", 0)),
            team_name=str(_row_value(row, "team_name", 1)),
            process=str(_row_value(row, "process", 2)),
            status=str(_row_value(row, "status", 3)),
            created_at=str(_row_value(row, "created_at", 4)),
            updated_at=str(_row_value(row, "updated_at", 5)),
            payload=payload,
        )

    def delete_checkpoint(self, run_id: str) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"DELETE FROM {self.table_name} WHERE run_id = %s",
                    (run_id,),
                )
            connection.commit()


def _row_value(row: Any, name: str, index: int) -> Any:
    if row is None:
        return None
    try:
        return row[name]
    except (KeyError, TypeError, IndexError):
        return row[index]


__all__ = [
    "CREW_CHECKPOINT_SCHEMA_KIND",
    "CREW_CHECKPOINT_SCHEMA_VERSION",
    "CrewCheckpointRecord",
    "CrewCheckpointStore",
    "InMemoryCrewCheckpointStore",
    "JSONObject",
    "JSONValue",
    "PostgreSQLCrewCheckpointStore",
    "SQLiteCrewCheckpointStore",
    "clone_json_object",
    "ensure_json_object",
]
