from __future__ import annotations

import json
import sqlite3
import threading
import time
from pathlib import Path
from typing import Dict, List, Optional, Any

from pulselog.exceptions import InvalidStatusError

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
VALID_STATUSES = {"DONE", "IN_PROGRESS", "FAILED", "SKIPPED"}

_SCHEMA = """
CREATE TABLE IF NOT EXISTS checkpoints (
    name       TEXT PRIMARY KEY,
    data       TEXT NOT NULL,
    status     TEXT NOT NULL,
    note       TEXT DEFAULT '',
    progress   REAL,
    saved_at   REAL NOT NULL
)
"""


class CheckpointStore:
    """Thread-safe SQLite-backed key/value store for training checkpoints.

    Multiple threads may call save/load/delete concurrently; a threading.Lock
    serialises all SQLite operations so we never hit "database is locked".

    WAL mode is enabled so concurrent reads do not block writes.
    """

    def __init__(self, db_path: str = ":memory:") -> None:
        """Open (or create) the checkpoint database.

        Args:
            db_path: Filesystem path to the SQLite file, or ":memory:".
        """
        if db_path != ":memory:":
            Path(db_path).parent.mkdir(parents=True, exist_ok=True)

        # check_same_thread=False is safe because we guard every access with _lock.
        self._conn = sqlite3.connect(db_path, check_same_thread=False)  # guarded by self._lock
        self._lock = threading.Lock()

        with self._lock:
            self._conn.execute("PRAGMA journal_mode=WAL")
            self._conn.execute(_SCHEMA)
            self._conn.commit()

        self._saved_count = 0  # guarded by self._lock

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def save(
        self,
        name: str,
        data: Dict[str, Any],
        status: str = "DONE",
        note: str = "",
        progress: Optional[float] = None,
    ) -> Dict[str, Any]:
        """Persist a checkpoint, overwriting any previous value for *name*.

        Args:
            name:     Unique checkpoint identifier.
            data:     JSON-serialisable payload dict.
            status:   One of DONE | IN_PROGRESS | FAILED | SKIPPED.
            note:     Human-readable description.
            progress: Optional 0–100 progress value.

        Returns:
            The full record dict including a ``_meta`` key.

        Raises:
            InvalidStatusError: If *status* is not one of the allowed values.
        """
        upper_status = status.upper()
        if upper_status not in VALID_STATUSES:
            raise InvalidStatusError(
                f"Invalid checkpoint status {status!r}. "
                f"Allowed values: {', '.join(sorted(VALID_STATUSES))}"
            )

        saved_at = time.time()
        data_json = json.dumps(data)

        with self._lock:
            self._conn.execute(
                """
                INSERT OR REPLACE INTO checkpoints
                    (name, data, status, note, progress, saved_at)
                VALUES (?, ?, ?, ?, ?, ?)
                """,
                (name, data_json, upper_status, note, progress, saved_at),
            )
            self._conn.commit()
            self._saved_count += 1

        return {
            **data,
            "_meta": {
                "name": name,
                "status": upper_status,
                "note": note,
                "progress": progress,
                "saved_at": saved_at,
            },
        }

    def load(self, name: str) -> Optional[Dict[str, Any]]:
        """Load a checkpoint by name.

        Args:
            name: Checkpoint identifier.

        Returns:
            Dict with original data keys plus ``_meta``, or None if not found.
        """
        with self._lock:
            cur = self._conn.execute(
                "SELECT data, status, note, progress, saved_at FROM checkpoints WHERE name = ?",
                (name,),
            )
            row = cur.fetchone()

        if row is None:
            return None

        data_json, status, note, progress, saved_at = row
        data = json.loads(data_json)
        return {
            **data,
            "_meta": {
                "name": name,
                "status": status,
                "note": note,
                "progress": progress,
                "saved_at": saved_at,
            },
        }

    def checkpoints(self) -> List[str]:
        """Return all checkpoint names ordered by most-recently saved first.

        Returns:
            List of name strings.
        """
        with self._lock:
            cur = self._conn.execute(
                "SELECT name FROM checkpoints ORDER BY saved_at DESC"
            )
            return [row[0] for row in cur.fetchall()]

    def delete_checkpoint(self, name: str) -> None:
        """Delete a checkpoint by name.  Silent if the name doesn't exist.

        Args:
            name: Checkpoint identifier to remove.
        """
        with self._lock:
            self._conn.execute("DELETE FROM checkpoints WHERE name = ?", (name,))
            self._conn.commit()

    def saved_count(self) -> int:
        """Total number of save() calls since this store was created."""
        with self._lock:
            return self._saved_count
