utils/chat_store.py

"""
Chat history persistence module using SQLite.

On Hugging Face Spaces:
  - Local writes work while the Space is running.
  - To survive restarts, enable HF Persistent Storage (mount to /data)
    and set the DB_PATH env var to /data/smarttrip.db
"""

import os
import sqlite3
from datetime import datetime
from typing import List, Dict, Optional

# Allow override via env var so HF persistent storage can be used
DB_PATH: str = os.getenv("DB_PATH", "data/smarttrip.db")


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _get_connection() -> sqlite3.Connection:
    """Open (and if needed create) the SQLite database."""
    os.makedirs(os.path.dirname(DB_PATH) if os.path.dirname(DB_PATH) else "data", exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


# ---------------------------------------------------------------------------
# Schema initialisation
# ---------------------------------------------------------------------------

def init_db() -> None:
    """Create tables if they don't exist yet. Safe to call on every startup."""
    conn = _get_connection()
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS chat_history (
            id            INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id    TEXT    NOT NULL,
            username      TEXT    NOT NULL,
            role          TEXT    NOT NULL DEFAULT 'user',
            user_message  TEXT    NOT NULL,
            bot_response  TEXT    NOT NULL,
            timestamp     TEXT    NOT NULL,
            status        TEXT    NOT NULL DEFAULT 'success'
        )
    """)
    conn.commit()
    conn.close()


# ---------------------------------------------------------------------------
# Write
# ---------------------------------------------------------------------------

def save_chat(
    session_id: str,
    username: str,
    role: str,
    user_message: str,
    bot_response: str,
    status: str = "success",
) -> None:
    """Persist one user-bot exchange."""
    conn = _get_connection()
    cursor = conn.cursor()
    cursor.execute(
        """
        INSERT INTO chat_history
            (session_id, username, role, user_message, bot_response, timestamp, status)
        VALUES (?, ?, ?, ?, ?, ?, ?)
        """,
        (
            session_id,
            username,
            role,
            user_message,
            bot_response,
            datetime.now().isoformat(sep=" ", timespec="seconds"),
            status,
        ),
    )
    conn.commit()
    conn.close()


# ---------------------------------------------------------------------------
# Read
# ---------------------------------------------------------------------------

def get_user_chats(username: str) -> List[Dict]:
    """Return all chat records for a specific user (newest first)."""
    conn = _get_connection()
    cursor = conn.cursor()
    cursor.execute(
        "SELECT * FROM chat_history WHERE username = ? ORDER BY timestamp DESC",
        (username,),
    )
    rows = [dict(row) for row in cursor.fetchall()]
    conn.close()
    return rows


def get_all_chats() -> List[Dict]:
    """Return all chat records across all users (newest first)."""
    conn = _get_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM chat_history ORDER BY timestamp DESC")
    rows = [dict(row) for row in cursor.fetchall()]
    conn.close()
    return rows


def get_users_with_chats() -> List[str]:
    """Return distinct non-admin usernames that have at least one chat record."""
    conn = _get_connection()
    cursor = conn.cursor()
    cursor.execute(
        "SELECT DISTINCT username FROM chat_history WHERE role != 'admin' ORDER BY username"
    )
    rows = [row["username"] for row in cursor.fetchall()]
    conn.close()
    return rows


def get_chat_stats() -> Dict:
    """Return aggregate statistics for the admin dashboard."""
    conn = _get_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT COUNT(*) AS total FROM chat_history")
    total = cursor.fetchone()["total"]
    cursor.execute("SELECT COUNT(DISTINCT username) AS users FROM chat_history WHERE role != 'admin'")
    users = cursor.fetchone()["users"]
    cursor.execute("SELECT COUNT(DISTINCT session_id) AS sessions FROM chat_history")
    sessions = cursor.fetchone()["sessions"]
    conn.close()
    return {"total_messages": total, "unique_users": users, "unique_sessions": sessions}