db.py

"""
Database layer — SQLite for user data and thread metadata.
LangGraph checkpointing also uses SQLite (configured in graph.py).
"""

import sqlite3
import time
from config import DB_PATH


conn = sqlite3.connect(database=DB_PATH, check_same_thread=False)

# --- Schema ---

conn.execute("""
    CREATE TABLE IF NOT EXISTS users (
        google_id       TEXT PRIMARY KEY,
        email           TEXT NOT NULL,
        name            TEXT NOT NULL,
        picture         TEXT,
        openrouter_key  TEXT,
        tavily_key      TEXT,
        student_profile TEXT,
        created_at      REAL NOT NULL,
        last_login      REAL NOT NULL
    )
""")

# Upgrade path: add tavily_key column to pre-existing databases.
try:
    conn.execute("ALTER TABLE users ADD COLUMN tavily_key TEXT")
except sqlite3.OperationalError:
    pass  # column already exists

conn.execute("""
    CREATE TABLE IF NOT EXISTS thread_titles (
        thread_id   TEXT PRIMARY KEY,
        title       TEXT,
        last_active REAL,
        user_id     TEXT
    )
""")

conn.commit()


# --- User operations ---

def upsert_user(google_id: str, email: str, name: str, picture: str = "") -> dict:
    """Create or update a user from Google OAuth data."""
    now = time.time()
    existing = conn.execute(
        "SELECT google_id FROM users WHERE google_id = ?", (google_id,)
    ).fetchone()

    if existing:
        conn.execute(
            "UPDATE users SET last_login = ?, name = ?, picture = ? WHERE google_id = ?",
            (now, name, picture, google_id),
        )
    else:
        conn.execute(
            "INSERT INTO users (google_id, email, name, picture, created_at, last_login) "
            "VALUES (?, ?, ?, ?, ?, ?)",
            (google_id, email, name, picture, now, now),
        )
    conn.commit()
    return {"google_id": google_id, "email": email, "name": name, "picture": picture}


def get_user(google_id: str) -> dict | None:
    row = conn.execute(
        "SELECT google_id, email, name, picture, openrouter_key, student_profile, tavily_key "
        "FROM users WHERE google_id = ?",
        (google_id,),
    ).fetchone()
    if not row:
        return None
    return {
        "google_id": row[0], "email": row[1], "name": row[2],
        "picture": row[3], "openrouter_key": row[4] or "",
        "student_profile": row[5] or "",
        "tavily_key": row[6] or "",
    }


def save_user_api_key(google_id: str, key: str):
    conn.execute(
        "UPDATE users SET openrouter_key = ? WHERE google_id = ?", (key, google_id)
    )
    conn.commit()


def save_tavily_key(google_id: str, key: str):
    conn.execute(
        "UPDATE users SET tavily_key = ? WHERE google_id = ?", (key, google_id)
    )
    conn.commit()


def get_tavily_key(google_id: str) -> str:
    row = conn.execute(
        "SELECT tavily_key FROM users WHERE google_id = ?", (google_id,)
    ).fetchone()
    return row[0] if row and row[0] else ""


def save_student_profile(google_id: str, profile: str):
    conn.execute(
        "UPDATE users SET student_profile = ? WHERE google_id = ?", (profile, google_id)
    )
    conn.commit()


def get_user_api_key(google_id: str) -> str:
    row = conn.execute(
        "SELECT openrouter_key FROM users WHERE google_id = ?", (google_id,)
    ).fetchone()
    return row[0] if row and row[0] else ""


# --- Thread operations ---

def get_threads(user_id: str = ""):
    if user_id:
        cursor = conn.execute(
            "SELECT thread_id, title FROM thread_titles "
            "WHERE user_id = ? ORDER BY last_active DESC",
            (user_id,),
        )
    else:
        cursor = conn.execute(
            "SELECT thread_id, title FROM thread_titles ORDER BY last_active DESC"
        )
    return [{"id": row[0], "title": row[1]} for row in cursor]


def upsert_thread(thread_id: str, message: str, now: float, user_id: str = ""):
    row = conn.execute(
        "SELECT thread_id FROM thread_titles WHERE thread_id = ?", (thread_id,)
    ).fetchone()

    if row:
        conn.execute(
            "UPDATE thread_titles SET last_active = ? WHERE thread_id = ?",
            (now, thread_id),
        )
    else:
        title = message[:30] + ("..." if len(message) > 30 else "")
        conn.execute(
            "INSERT INTO thread_titles (thread_id, title, last_active, user_id) "
            "VALUES (?, ?, ?, ?)",
            (thread_id, title, now, user_id),
        )
    conn.commit()


def rename_thread(thread_id: str, title: str):
    conn.execute(
        "UPDATE thread_titles SET title = ? WHERE thread_id = ?", (title, thread_id)
    )
    conn.commit()


def delete_thread(thread_id: str):
    conn.execute("DELETE FROM thread_titles WHERE thread_id = ?", (thread_id,))
    conn.commit()