session_store.py

"""
session_store.py
----------------
Persistence for `DialogueState`, keyed by session_id.

Two backends:
- SQLite (default for local / Render / Docker with persistent volume)
- In-memory dict (for Hugging Face Spaces where the filesystem is ephemeral)

The backend is selected via the `PAYPARSE_SESSION_BACKEND` env var:
  "sqlite"  (default) → SQLite file
  "memory"            → in-memory dict (lost on restart, fine for a demo)
"""

from __future__ import annotations

import os
import sqlite3
import threading
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from schema import DialogueState

DEFAULT_DB_PATH = Path(__file__).parent / "sessions.db"

_SCHEMA = """
CREATE TABLE IF NOT EXISTS sessions (
    session_id  TEXT PRIMARY KEY,
    state_json  TEXT NOT NULL,
    updated_at  TEXT NOT NULL
);
"""


class SessionStore:
    """Thread-safe session store. Delegates to SQLite or in-memory dict."""

    def __init__(self, db_path: Path | str = DEFAULT_DB_PATH):
        backend = os.getenv("PAYPARSE_SESSION_BACKEND", "sqlite").lower()
        self._lock = threading.Lock()
        if backend == "memory":
            self._backend = _MemoryBackend()
        else:
            self._backend = _SQLiteBackend(db_path)

    def get(self, session_id: str) -> Optional[DialogueState]:
        return self._backend.get(session_id)

    def get_or_create(self, session_id: str) -> DialogueState:
        return self.get(session_id) or DialogueState(session_id=session_id)

    def save(self, state: DialogueState) -> None:
        self._backend.save(state)

    def delete(self, session_id: str) -> bool:
        return self._backend.delete(session_id)

    def close(self) -> None:
        self._backend.close()


class _SQLiteBackend:
    """SQLite-backed persistence (local / Render / Docker with volume)."""

    def __init__(self, db_path: Path | str):
        self._conn = sqlite3.connect(str(db_path), check_same_thread=False)
        self._conn.execute(_SCHEMA)
        self._conn.commit()

    def get(self, session_id: str) -> Optional[DialogueState]:
        row = self._conn.execute(
            "SELECT state_json FROM sessions WHERE session_id = ?",
            (session_id,),
        ).fetchone()
        return DialogueState.model_validate_json(row[0]) if row else None

    def save(self, state: DialogueState) -> None:
        now = datetime.now(timezone.utc).isoformat()
        self._conn.execute(
            "INSERT INTO sessions (session_id, state_json, updated_at) "
            "VALUES (?, ?, ?) "
            "ON CONFLICT(session_id) DO UPDATE SET "
            "state_json = excluded.state_json, updated_at = excluded.updated_at",
            (state.session_id, state.model_dump_json(), now),
        )
        self._conn.commit()

    def delete(self, session_id: str) -> bool:
        cursor = self._conn.execute(
            "DELETE FROM sessions WHERE session_id = ?", (session_id,)
        )
        self._conn.commit()
        return cursor.rowcount > 0

    def close(self) -> None:
        self._conn.close()


class _MemoryBackend:
    """In-memory dict store (Hugging Face Spaces — ephemeral filesystem)."""

    def __init__(self):
        self._data: dict[str, str] = {}

    def get(self, session_id: str) -> Optional[DialogueState]:
        json_str = self._data.get(session_id)
        return DialogueState.model_validate_json(json_str) if json_str else None

    def save(self, state: DialogueState) -> None:
        self._data[state.session_id] = state.model_dump_json()

    def delete(self, session_id: str) -> bool:
        return self._data.pop(session_id, None) is not None

    def close(self) -> None:
        pass