chatkit/backend/app/database.py

"""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()


# ---- Answer Grid CRUD ----

ANSWER_GRID_DATA_TYPE = "answer_grid"


async def save_answer_grid(session_id: int, grid_dict: dict) -> int:
    """Save confirmed answer grid, replacing any previous one for the session."""
    await delete_parsed_data(session_id, ANSWER_GRID_DATA_TYPE)
    return await save_parsed_data(
        session_id=session_id,
        data_type=ANSWER_GRID_DATA_TYPE,
        file_name="",
        raw_text="",
        structured_data=grid_dict,
    )


async def get_answer_grid(session_id: int) -> Optional[dict]:
    """Return the saved answer grid dict or None if not confirmed."""
    db_path = _get_db_path()
    async with aiosqlite.connect(db_path) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute(
            "SELECT structured_data FROM parsed_data WHERE session_id = ? AND data_type = ? ORDER BY created_at DESC LIMIT 1",
            (session_id, ANSWER_GRID_DATA_TYPE),
        ) as cursor:
            row = await cursor.fetchone()
            if not row:
                return None
            return json.loads(row["structured_data"])


# ---- 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