"""SQLite database for ClassLens - users, sessions, parsed data, prompts, reports.""" from __future__ import annotations import aiosqlite import json from pathlib import Path from datetime import datetime from typing import Optional from .config import get_settings DATABASE_PATH = Path(__file__).parent.parent / "classlens.db" def _get_db_path() -> Path: settings = get_settings() if settings.database_path: return Path(settings.database_path) return DATABASE_PATH async def init_database(): """Initialize the database with required tables.""" db_path = _get_db_path() async with aiosqlite.connect(db_path) as db: await db.execute(""" CREATE TABLE IF NOT EXISTS users ( id INTEGER PRIMARY KEY AUTOINCREMENT, email TEXT UNIQUE NOT NULL, password_hash TEXT NOT NULL, display_name TEXT, created_at TEXT NOT NULL, last_login TEXT ) """) await db.execute(""" CREATE TABLE IF NOT EXISTS exam_sessions ( id INTEGER PRIMARY KEY AUTOINCREMENT, user_id INTEGER NOT NULL, title TEXT DEFAULT 'Untitled Session', status TEXT DEFAULT 'draft', created_at TEXT NOT NULL, updated_at TEXT NOT NULL, FOREIGN KEY (user_id) REFERENCES users(id) ) """) await db.execute(""" CREATE TABLE IF NOT EXISTS parsed_data ( id INTEGER PRIMARY KEY AUTOINCREMENT, session_id INTEGER NOT NULL, data_type TEXT NOT NULL, file_name TEXT, raw_text TEXT, structured_data TEXT NOT NULL, created_at TEXT NOT NULL, FOREIGN KEY (session_id) REFERENCES exam_sessions(id) ON DELETE CASCADE ) """) await db.execute(""" CREATE TABLE IF NOT EXISTS prompts ( id INTEGER PRIMARY KEY AUTOINCREMENT, user_id INTEGER NOT NULL, name TEXT NOT NULL, content TEXT NOT NULL, is_default BOOLEAN DEFAULT 0, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, FOREIGN KEY (user_id) REFERENCES users(id) ) """) await db.execute(""" CREATE TABLE IF NOT EXISTS reports ( id INTEGER PRIMARY KEY AUTOINCREMENT, session_id INTEGER NOT NULL, prompt_id INTEGER, html_content TEXT NOT NULL, created_at TEXT NOT NULL, FOREIGN KEY (session_id) REFERENCES exam_sessions(id) ON DELETE CASCADE, FOREIGN KEY (prompt_id) REFERENCES prompts(id) ) """) await db.commit() # ---- User CRUD ---- async def create_user(email: str, password_hash: str, display_name: str = "") -> int: db_path = _get_db_path() now = datetime.utcnow().isoformat() async with aiosqlite.connect(db_path) as db: cursor = await db.execute( "INSERT INTO users (email, password_hash, display_name, created_at) VALUES (?, ?, ?, ?)", (email, password_hash, display_name or email.split("@")[0], now), ) await db.commit() return cursor.lastrowid async def get_user_by_email(email: str) -> Optional[dict]: db_path = _get_db_path() async with aiosqlite.connect(db_path) as db: db.row_factory = aiosqlite.Row async with db.execute("SELECT * FROM users WHERE email = ?", (email,)) as cursor: row = await cursor.fetchone() return dict(row) if row else None async def get_user_by_id(user_id: int) -> Optional[dict]: db_path = _get_db_path() async with aiosqlite.connect(db_path) as db: db.row_factory = aiosqlite.Row async with db.execute("SELECT * FROM users WHERE id = ?", (user_id,)) as cursor: row = await cursor.fetchone() return dict(row) if row else None async def update_last_login(user_id: int): db_path = _get_db_path() now = datetime.utcnow().isoformat() async with aiosqlite.connect(db_path) as db: await db.execute("UPDATE users SET last_login = ? WHERE id = ?", (now, user_id)) await db.commit() # ---- Session CRUD ---- async def create_session(user_id: int, title: str = "Untitled Session") -> int: db_path = _get_db_path() now = datetime.utcnow().isoformat() async with aiosqlite.connect(db_path) as db: cursor = await db.execute( "INSERT INTO exam_sessions (user_id, title, status, created_at, updated_at) VALUES (?, ?, 'draft', ?, ?)", (user_id, title, now, now), ) await db.commit() return cursor.lastrowid async def get_sessions(user_id: int) -> list[dict]: db_path = _get_db_path() async with aiosqlite.connect(db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT * FROM exam_sessions WHERE user_id = ? ORDER BY created_at DESC", (user_id,) ) as cursor: rows = await cursor.fetchall() return [dict(row) for row in rows] async def get_session(session_id: int) -> Optional[dict]: db_path = _get_db_path() async with aiosqlite.connect(db_path) as db: db.row_factory = aiosqlite.Row async with db.execute("SELECT * FROM exam_sessions WHERE id = ?", (session_id,)) as cursor: row = await cursor.fetchone() return dict(row) if row else None async def update_session_status(session_id: int, status: str): db_path = _get_db_path() now = datetime.utcnow().isoformat() async with aiosqlite.connect(db_path) as db: await db.execute( "UPDATE exam_sessions SET status = ?, updated_at = ? WHERE id = ?", (status, now, session_id), ) await db.commit() # ---- Parsed Data CRUD ---- async def save_parsed_data( session_id: int, data_type: str, file_name: str, raw_text: str, structured_data: dict ) -> int: db_path = _get_db_path() now = datetime.utcnow().isoformat() async with aiosqlite.connect(db_path) as db: cursor = await db.execute( "INSERT INTO parsed_data (session_id, data_type, file_name, raw_text, structured_data, created_at) VALUES (?, ?, ?, ?, ?, ?)", (session_id, data_type, file_name, raw_text, json.dumps(structured_data, ensure_ascii=False), now), ) await db.commit() return cursor.lastrowid async def get_parsed_data(session_id: int) -> list[dict]: db_path = _get_db_path() async with aiosqlite.connect(db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT * FROM parsed_data WHERE session_id = ? ORDER BY data_type, created_at", (session_id,), ) as cursor: rows = await cursor.fetchall() result = [] for row in rows: d = dict(row) d["structured_data"] = json.loads(d["structured_data"]) result.append(d) return result async def delete_parsed_data(session_id: int, data_type: Optional[str] = None): db_path = _get_db_path() async with aiosqlite.connect(db_path) as db: if data_type: await db.execute( "DELETE FROM parsed_data WHERE session_id = ? AND data_type = ?", (session_id, data_type), ) else: await db.execute("DELETE FROM parsed_data WHERE session_id = ?", (session_id,)) await db.commit() # ---- Prompt CRUD ---- async def save_prompt(user_id: int, name: str, content: str, is_default: bool = False) -> int: db_path = _get_db_path() now = datetime.utcnow().isoformat() async with aiosqlite.connect(db_path) as db: cursor = await db.execute( "INSERT INTO prompts (user_id, name, content, is_default, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?)", (user_id, name, content, is_default, now, now), ) await db.commit() return cursor.lastrowid async def get_prompts(user_id: int) -> list[dict]: """Get user's own prompts plus all default prompts.""" db_path = _get_db_path() async with aiosqlite.connect(db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT * FROM prompts WHERE user_id = ? OR is_default = 1 ORDER BY is_default DESC, updated_at DESC", (user_id,), ) as cursor: rows = await cursor.fetchall() return [dict(row) for row in rows] async def get_all_prompts() -> list[dict]: """Get all prompts from all users (for the 'read others' dropdown).""" db_path = _get_db_path() async with aiosqlite.connect(db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT p.*, u.email as author_email FROM prompts p JOIN users u ON p.user_id = u.id ORDER BY p.updated_at DESC" ) as cursor: rows = await cursor.fetchall() return [dict(row) for row in rows] async def get_prompt(prompt_id: int) -> Optional[dict]: db_path = _get_db_path() async with aiosqlite.connect(db_path) as db: db.row_factory = aiosqlite.Row async with db.execute("SELECT * FROM prompts WHERE id = ?", (prompt_id,)) as cursor: row = await cursor.fetchone() return dict(row) if row else None async def update_prompt(prompt_id: int, name: str, content: str): db_path = _get_db_path() now = datetime.utcnow().isoformat() async with aiosqlite.connect(db_path) as db: await db.execute( "UPDATE prompts SET name = ?, content = ?, updated_at = ? WHERE id = ?", (name, content, now, prompt_id), ) await db.commit() async def delete_prompt(prompt_id: int): db_path = _get_db_path() async with aiosqlite.connect(db_path) as db: await db.execute("DELETE FROM prompts WHERE id = ?", (prompt_id,)) await db.commit() # ---- Report CRUD ---- async def save_report(session_id: int, prompt_id: Optional[int], html_content: str) -> int: db_path = _get_db_path() now = datetime.utcnow().isoformat() async with aiosqlite.connect(db_path) as db: cursor = await db.execute( "INSERT INTO reports (session_id, prompt_id, html_content, created_at) VALUES (?, ?, ?, ?)", (session_id, prompt_id, html_content, now), ) await db.commit() return cursor.lastrowid async def get_reports(session_id: int) -> list[dict]: db_path = _get_db_path() async with aiosqlite.connect(db_path) as db: db.row_factory = aiosqlite.Row async with db.execute( "SELECT * FROM reports WHERE session_id = ? ORDER BY created_at DESC", (session_id,), ) as cursor: rows = await cursor.fetchall() return [dict(row) for row in rows] async def get_report(report_id: int) -> Optional[dict]: db_path = _get_db_path() async with aiosqlite.connect(db_path) as db: db.row_factory = aiosqlite.Row async with db.execute("SELECT * FROM reports WHERE id = ?", (report_id,)) as cursor: row = await cursor.fetchone() return dict(row) if row else None