src/storage.py

"""
src/storage.py — SQLite-backed chat history store and rotating file logger.

Every user message and assistant response is persisted with metadata so the
admin view can display full per-user audit trails.

Schema
──────
chat_messages
  id          TEXT PK   — UUID per row
  session_id  TEXT      — UUID per chat session (reset on "New Chat")
  user_id     TEXT      — stable user identifier from auth.py
  username    TEXT      — human-readable username
  user_role   TEXT      — "user" | "admin"
  timestamp   TEXT      — UTC ISO-8601
  msg_role    TEXT      — "user" | "assistant"
  content     TEXT      — message body
  model_name  TEXT      — Ollama model used for this turn
  status      TEXT      — "ok" | "blocked_scope" | "blocked_emergency" | "error"
  error_info  TEXT      — non-empty on error / blocked responses
"""
import logging
import logging.handlers
import sqlite3
import uuid
from datetime import datetime, timezone
from typing import Dict, List

from src.config import DB_PATH, LOG_PATH

# ─── Logging setup ────────────────────────────────────────────────────────────

def setup_logging() -> None:
    """Configure a rotating file logger plus a console handler."""
    LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
    root = logging.getLogger()
    if root.handlers:
        return  # already configured (Streamlit re-runs the script)

    root.setLevel(logging.INFO)

    fh = logging.handlers.RotatingFileHandler(
        LOG_PATH,
        maxBytes=5 * 1024 * 1024,   # 5 MB per file
        backupCount=5,
        encoding="utf-8",
    )
    fh.setFormatter(
        logging.Formatter(
            "%(asctime)s | %(levelname)-8s | %(name)-30s | %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S",
        )
    )

    sh = logging.StreamHandler()
    sh.setLevel(logging.WARNING)
    sh.setFormatter(logging.Formatter("%(levelname)s: %(message)s"))

    root.addHandler(fh)
    root.addHandler(sh)


# ─── Database helpers ─────────────────────────────────────────────────────────

def _connect() -> sqlite3.Connection:
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(DB_PATH), check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn


def init_db() -> None:
    """Create tables and indexes if they don't already exist."""
    with _connect() as conn:
        conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS chat_messages (
                id          TEXT PRIMARY KEY,
                session_id  TEXT NOT NULL,
                user_id     TEXT NOT NULL,
                username    TEXT NOT NULL,
                user_role   TEXT NOT NULL,
                timestamp   TEXT NOT NULL,
                msg_role    TEXT NOT NULL,
                content     TEXT NOT NULL,
                model_name  TEXT DEFAULT '',
                status      TEXT DEFAULT 'ok',
                error_info  TEXT DEFAULT ''
            );

            CREATE INDEX IF NOT EXISTS idx_cm_user_id    ON chat_messages (user_id);
            CREATE INDEX IF NOT EXISTS idx_cm_session    ON chat_messages (session_id);
            CREATE INDEX IF NOT EXISTS idx_cm_timestamp  ON chat_messages (timestamp);
            """
        )


def save_message(
    *,
    session_id: str,
    user_id: str,
    username: str,
    user_role: str,
    msg_role: str,
    content: str,
    model_name: str = "",
    status: str = "ok",
    error_info: str = "",
) -> None:
    """Persist a single chat turn (user or assistant) to the database."""
    with _connect() as conn:
        conn.execute(
            """
            INSERT INTO chat_messages
                (id, session_id, user_id, username, user_role,
                 timestamp, msg_role, content, model_name, status, error_info)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                str(uuid.uuid4()),
                session_id,
                user_id,
                username,
                user_role,
                datetime.now(timezone.utc).isoformat(),
                msg_role,
                content,
                model_name,
                status,
                error_info,
            ),
        )


# ─── Query helpers ────────────────────────────────────────────────────────────

def get_user_history(user_id: str, limit: int = 300) -> List[Dict]:
    """Return the most recent `limit` messages for a user, newest first."""
    with _connect() as conn:
        rows = conn.execute(
            """
            SELECT * FROM chat_messages
            WHERE user_id = ?
            ORDER BY timestamp DESC
            LIMIT ?
            """,
            (user_id, limit),
        ).fetchall()
    return [dict(r) for r in rows]


def get_session_messages(session_id: str) -> List[Dict]:
    """Return all messages for a session ordered oldest-first (for replay)."""
    with _connect() as conn:
        rows = conn.execute(
            """
            SELECT * FROM chat_messages
            WHERE session_id = ?
            ORDER BY timestamp ASC
            """,
            (session_id,),
        ).fetchall()
    return [dict(r) for r in rows]


def get_all_users_summary() -> List[Dict]:
    """
    Return one row per user with aggregate statistics.
    Used by the admin panel to populate the user selector.
    """
    with _connect() as conn:
        rows = conn.execute(
            """
            SELECT
                user_id,
                username,
                user_role,
                COUNT(*)             AS message_count,
                MAX(timestamp)       AS last_active,
                COUNT(DISTINCT session_id) AS session_count
            FROM chat_messages
            GROUP BY user_id
            ORDER BY last_active DESC
            """
        ).fetchall()
    return [dict(r) for r in rows]


def get_platform_stats() -> Dict:
    """Return aggregate platform-wide KPIs for the admin dashboard."""
    with _connect() as conn:
        row = conn.execute(
            """
            SELECT
                COUNT(DISTINCT user_id)    AS total_users,
                COUNT(DISTINCT session_id) AS total_sessions,
                COUNT(*)                   AS total_messages,
                SUM(CASE WHEN msg_role = 'user' THEN 1 ELSE 0 END) AS user_messages,
                SUM(CASE WHEN status = 'error' THEN 1 ELSE 0 END)  AS errors,
                SUM(CASE WHEN status = 'blocked_emergency' THEN 1 ELSE 0 END) AS emergencies,
                MAX(timestamp) AS last_activity
            FROM chat_messages
            """
        ).fetchone()
        # Active today: messages since midnight UTC
        today_row = conn.execute(
            """
            SELECT COUNT(DISTINCT user_id) AS active_today
            FROM chat_messages
            WHERE timestamp >= strftime('%Y-%m-%dT00:00:00', 'now')
            """
        ).fetchone()
    result = dict(row) if row else {}
    result["active_today"] = today_row["active_today"] if today_row else 0
    return result


def get_user_sessions(user_id: str) -> List[Dict]:
    """Return one row per session for a user, newest first."""
    with _connect() as conn:
        rows = conn.execute(
            """
            SELECT
                session_id,
                MIN(timestamp)  AS started_at,
                MAX(timestamp)  AS last_msg_at,
                COUNT(*)        AS message_count,
                SUM(CASE WHEN status != 'ok' AND status != '' THEN 1 ELSE 0 END) AS flagged
            FROM chat_messages
            WHERE user_id = ?
            GROUP BY session_id
            ORDER BY last_msg_at DESC
            """,
            (user_id,),
        ).fetchall()
    return [dict(r) for r in rows]