chatkit/backend/app/database.py

"""SQLite database for token storage."""

from __future__ import annotations

import aiosqlite
import json
import base64
import os
from pathlib import Path
from datetime import datetime
from typing import Optional

from cryptography.fernet import Fernet


DATABASE_PATH = Path(__file__).parent.parent / "classlens.db"


def get_encryption_key() -> bytes:
    """Get or generate encryption key."""
    key = os.getenv("ENCRYPTION_KEY", "")
    if not key:
        # Generate a new key for development (should be set in production)
        key = Fernet.generate_key().decode()
        print(f"⚠️  No ENCRYPTION_KEY set. Generated temporary key: {key}")
        print("⚠️  Set this in your .env file for persistent token storage.")
    return key.encode() if isinstance(key, str) else key


def encrypt_token(token_data: dict) -> str:
    """Encrypt token data for storage."""
    key = get_encryption_key()
    f = Fernet(key)
    json_data = json.dumps(token_data)
    encrypted = f.encrypt(json_data.encode())
    return base64.b64encode(encrypted).decode()


def decrypt_token(encrypted_data: str) -> dict:
    """Decrypt stored token data."""
    key = get_encryption_key()
    f = Fernet(key)
    encrypted = base64.b64decode(encrypted_data.encode())
    decrypted = f.decrypt(encrypted)
    return json.loads(decrypted.decode())


async def init_database():
    """Initialize the database with required tables."""
    async with aiosqlite.connect(DATABASE_PATH) as db:
        await db.execute("""
            CREATE TABLE IF NOT EXISTS oauth_tokens (
                teacher_email TEXT PRIMARY KEY,
                encrypted_tokens TEXT NOT NULL,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL
            )
        """)
        
        await db.execute("""
            CREATE TABLE IF NOT EXISTS analysis_reports (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                teacher_email TEXT NOT NULL,
                exam_title TEXT,
                report_markdown TEXT,
                report_json TEXT,
                created_at TEXT NOT NULL
            )
        """)
        
        await db.commit()


async def save_oauth_tokens(teacher_email: str, tokens: dict):
    """Save encrypted OAuth tokens for a teacher."""
    encrypted = encrypt_token(tokens)
    now = datetime.utcnow().isoformat()
    
    async with aiosqlite.connect(DATABASE_PATH) as db:
        await db.execute("""
            INSERT INTO oauth_tokens (teacher_email, encrypted_tokens, created_at, updated_at)
            VALUES (?, ?, ?, ?)
            ON CONFLICT(teacher_email) DO UPDATE SET
                encrypted_tokens = excluded.encrypted_tokens,
                updated_at = excluded.updated_at
        """, (teacher_email, encrypted, now, now))
        await db.commit()


async def get_oauth_tokens(teacher_email: str) -> Optional[dict]:
    """Retrieve OAuth tokens for a teacher."""
    async with aiosqlite.connect(DATABASE_PATH) as db:
        async with db.execute(
            "SELECT encrypted_tokens FROM oauth_tokens WHERE teacher_email = ?",
            (teacher_email,)
        ) as cursor:
            row = await cursor.fetchone()
            if row:
                return decrypt_token(row[0])
    return None


async def delete_oauth_tokens(teacher_email: str):
    """Delete OAuth tokens for a teacher."""
    async with aiosqlite.connect(DATABASE_PATH) as db:
        await db.execute(
            "DELETE FROM oauth_tokens WHERE teacher_email = ?",
            (teacher_email,)
        )
        await db.commit()


async def save_report(teacher_email: str, exam_title: str, report_markdown: str, report_json: str):
    """Save an analysis report."""
    now = datetime.utcnow().isoformat()
    
    async with aiosqlite.connect(DATABASE_PATH) as db:
        await db.execute("""
            INSERT INTO analysis_reports (teacher_email, exam_title, report_markdown, report_json, created_at)
            VALUES (?, ?, ?, ?, ?)
        """, (teacher_email, exam_title, report_markdown, report_json, now))
        await db.commit()


async def get_teacher_reports(teacher_email: str, limit: int = 10) -> list[dict]:
    """Get recent reports for a teacher."""
    async with aiosqlite.connect(DATABASE_PATH) as db:
        db.row_factory = aiosqlite.Row
        async with db.execute("""
            SELECT id, exam_title, report_markdown, report_json, created_at
            FROM analysis_reports
            WHERE teacher_email = ?
            ORDER BY created_at DESC
            LIMIT ?
        """, (teacher_email, limit)) as cursor:
            rows = await cursor.fetchall()
            return [dict(row) for row in rows]