src/database.py

"""
SQLite database layer — users + research sessions.
Uses only the stdlib sqlite3 module, no ORM needed.
"""
import sqlite3
import os
import json
from datetime import datetime

# HF Spaces mounts persistent storage at /data; fall back to local data/ dir
_HF_DATA = "/data"
if os.path.isdir(_HF_DATA) and os.access(_HF_DATA, os.W_OK):
    DB_PATH = os.path.join(_HF_DATA, "app.db")
else:
    DB_PATH = os.path.join(os.path.dirname(__file__), "..", "data", "app.db")


def get_conn() -> sqlite3.Connection:
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn


def init_db():
    """Create tables if they don't exist."""
    with get_conn() as conn:
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS users (
                id        INTEGER PRIMARY KEY AUTOINCREMENT,
                username  TEXT    UNIQUE NOT NULL,
                email     TEXT    UNIQUE NOT NULL,
                password  TEXT    NOT NULL,
                created_at TEXT   NOT NULL
            );

            CREATE TABLE IF NOT EXISTS sessions (
                id            TEXT    PRIMARY KEY,
                user_id       INTEGER NOT NULL,
                topic         TEXT    NOT NULL,
                draft         TEXT    NOT NULL,
                critique_score REAL   NOT NULL DEFAULT 0,
                iterations    INTEGER NOT NULL DEFAULT 0,
                research_notes TEXT   NOT NULL DEFAULT '[]',
                models        TEXT    NOT NULL DEFAULT '{}',
                created_at    TEXT    NOT NULL,
                FOREIGN KEY (user_id) REFERENCES users(id)
            );
        """)
        # Migrate existing DBs that predate the models column
        cols = [r[1] for r in conn.execute("PRAGMA table_info(sessions)").fetchall()]
        if "models" not in cols:
            conn.execute("ALTER TABLE sessions ADD COLUMN models TEXT NOT NULL DEFAULT '{}'")


# ── User helpers ──────────────────────────────────────────────────────────────

def create_user(username: str, email: str, hashed_password: str) -> dict:
    with get_conn() as conn:
        cur = conn.execute(
            "INSERT INTO users (username, email, password, created_at) VALUES (?,?,?,?)",
            (username, email, hashed_password, datetime.utcnow().isoformat()),
        )
        return {"id": cur.lastrowid, "username": username, "email": email}


def get_user_by_email(email: str) -> sqlite3.Row | None:
    with get_conn() as conn:
        return conn.execute("SELECT * FROM users WHERE email=?", (email,)).fetchone()


def get_user_by_id(user_id: int) -> sqlite3.Row | None:
    with get_conn() as conn:
        return conn.execute("SELECT * FROM users WHERE id=?", (user_id,)).fetchone()


# ── Session helpers ───────────────────────────────────────────────────────────

def save_session(session_id: str, user_id: int, topic: str, draft: str,
                 score: float, iterations: int, notes: list, models: str = "{}") -> dict:
    with get_conn() as conn:
        conn.execute(
            """INSERT INTO sessions
               (id, user_id, topic, draft, critique_score, iterations, research_notes, models, created_at)
               VALUES (?,?,?,?,?,?,?,?,?)""",
            (session_id, user_id, topic, draft, score, iterations,
             json.dumps(notes), models, datetime.utcnow().isoformat()),
        )
    return {"id": session_id}


def get_sessions_for_user(user_id: int) -> list[dict]:
    with get_conn() as conn:
        rows = conn.execute(
            "SELECT * FROM sessions WHERE user_id=? ORDER BY created_at DESC",
            (user_id,),
        ).fetchall()
    result = []
    for r in rows:
        d = dict(r)
        d["research_notes"] = json.loads(d["research_notes"])
        result.append(d)
    return result


def get_session(session_id: str, user_id: int) -> dict | None:
    with get_conn() as conn:
        row = conn.execute(
            "SELECT * FROM sessions WHERE id=? AND user_id=?",
            (session_id, user_id),
        ).fetchone()
    if not row:
        return None
    d = dict(row)
    d["research_notes"] = json.loads(d["research_notes"])
    return d


def delete_session(session_id: str, user_id: int) -> bool:
    with get_conn() as conn:
        cur = conn.execute(
            "DELETE FROM sessions WHERE id=? AND user_id=?",
            (session_id, user_id),
        )
    return cur.rowcount > 0