chatkit/backend/app/database.py

"""Data-access layer for ClassLens, backed by SQLAlchemy async (Supabase Postgres).

Public function names/signatures are kept stable for existing callers:
  - "session" in the API == a row in the ``quizzes`` table (a quiz/exam session)
  - "user" == a row in the ``teachers`` table

New capabilities (students as first-class entities, per-student results, and
analytics) are added alongside the legacy CRUD helpers.
"""

from __future__ import annotations

import re
import unicodedata
from datetime import datetime, timezone
from typing import Optional

from sqlalchemy import delete, func, select, update

from .config import get_settings
from .db import get_engine, get_sessionmaker
from .models import (
    Base,
    ParsedData,
    Prompt,
    Quiz,
    RawFile,
    Report,
    Student,
    StudentResult,
    Teacher,
)

ANSWER_GRID_DATA_TYPE = "answer_grid"


def _now() -> datetime:
    return datetime.now(timezone.utc)


def _row(obj) -> dict:
    """Serialize an ORM object to a plain dict of its columns."""
    return {c.key: getattr(obj, c.key) for c in obj.__table__.columns}


def normalize_name(name: str) -> str:
    """Canonical key for matching the same student across quizzes."""
    s = unicodedata.normalize("NFKC", (name or "").strip()).casefold()
    return re.sub(r"\s+", " ", s)


# =============================================================================
# Schema bootstrap
# =============================================================================

async def init_database():
    """Ensure schema exists and bootstrap admins.

    On Supabase, tables are managed by Alembic migrations; here we only create
    tables for the SQLite dev/test fallback. Admin bootstrap runs for both.
    """
    settings = get_settings()

    if settings.database_provider != "supabase":
        async with get_engine().begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

    await _bootstrap_admins()


async def _bootstrap_admins():
    settings = get_settings()
    emails = [e.strip().lower() for e in (settings.admin_emails or "").split(",") if e.strip()]
    if not emails:
        return
    async with get_sessionmaker()() as db:
        await db.execute(
            update(Teacher).where(func.lower(Teacher.email).in_(emails)).values(is_admin=True)
        )
        await db.commit()


# =============================================================================
# Teachers (users)
# =============================================================================

async def create_user(
    email: str,
    password_hash: str,
    display_name: str = "",
    full_name: str = "",
    school: str = "",
) -> int:
    async with get_sessionmaker()() as db:
        teacher = Teacher(
            email=email,
            password_hash=password_hash,
            display_name=display_name or email.split("@")[0],
            full_name=full_name,
            school=school,
        )
        db.add(teacher)
        await db.commit()
        await db.refresh(teacher)
        return teacher.id


async def get_user_by_email(email: str) -> Optional[dict]:
    async with get_sessionmaker()() as db:
        res = await db.execute(select(Teacher).where(Teacher.email == email))
        teacher = res.scalar_one_or_none()
        return _row(teacher) if teacher else None


async def get_user_by_id(user_id: int) -> Optional[dict]:
    async with get_sessionmaker()() as db:
        teacher = await db.get(Teacher, user_id)
        return _row(teacher) if teacher else None


async def update_last_login(user_id: int):
    async with get_sessionmaker()() as db:
        await db.execute(
            update(Teacher).where(Teacher.id == user_id).values(last_login=_now())
        )
        await db.commit()


async def set_user_admin(user_id: int, is_admin: bool):
    async with get_sessionmaker()() as db:
        await db.execute(
            update(Teacher).where(Teacher.id == user_id).values(is_admin=is_admin)
        )
        await db.commit()


async def list_all_users() -> list[dict]:
    """Admin: all teachers (without password hashes), with per-teacher counts."""
    async with get_sessionmaker()() as db:
        res = await db.execute(select(Teacher).order_by(Teacher.created_at.desc()))
        teachers = res.scalars().all()
        out = []
        for t in teachers:
            d = _row(t)
            d.pop("password_hash", None)
            d["quiz_count"] = (
                await db.execute(select(func.count(Quiz.id)).where(Quiz.teacher_id == t.id))
            ).scalar_one()
            d["student_count"] = (
                await db.execute(select(func.count(Student.id)).where(Student.teacher_id == t.id))
            ).scalar_one()
            out.append(d)
        return out


async def delete_user(user_id: int):
    async with get_sessionmaker()() as db:
        await db.execute(delete(Teacher).where(Teacher.id == user_id))
        await db.commit()


# =============================================================================
# Quizzes (sessions)
# =============================================================================

async def create_session(user_id: int, title: str = "Untitled Session") -> int:
    async with get_sessionmaker()() as db:
        quiz = Quiz(teacher_id=user_id, title=title, status="draft")
        db.add(quiz)
        await db.commit()
        await db.refresh(quiz)
        return quiz.id


async def get_sessions(user_id: int) -> list[dict]:
    async with get_sessionmaker()() as db:
        res = await db.execute(
            select(Quiz).where(Quiz.teacher_id == user_id).order_by(Quiz.created_at.desc())
        )
        return [_row(q) for q in res.scalars().all()]


async def get_session(session_id: int) -> Optional[dict]:
    async with get_sessionmaker()() as db:
        quiz = await db.get(Quiz, session_id)
        if not quiz:
            return None
        d = _row(quiz)
        d["user_id"] = quiz.teacher_id  # back-compat key for existing callers
        return d


async def update_session_status(session_id: int, status: str):
    async with get_sessionmaker()() as db:
        await db.execute(
            update(Quiz).where(Quiz.id == session_id).values(status=status, updated_at=_now())
        )
        await db.commit()


async def update_session_meta(session_id: int, title: Optional[str] = None, subject: Optional[str] = None):
    values: dict = {"updated_at": _now()}
    if title is not None:
        values["title"] = title
    if subject is not None:
        values["subject"] = subject
    async with get_sessionmaker()() as db:
        await db.execute(update(Quiz).where(Quiz.id == session_id).values(**values))
        await db.commit()


async def delete_session(session_id: int):
    async with get_sessionmaker()() as db:
        await db.execute(delete(Quiz).where(Quiz.id == session_id))
        await db.commit()


# =============================================================================
# Parsed data
# =============================================================================

async def save_parsed_data(
    session_id: int, data_type: str, file_name: str, raw_text: str, structured_data: dict
) -> int:
    async with get_sessionmaker()() as db:
        row = ParsedData(
            quiz_id=session_id,
            data_type=data_type,
            file_name=file_name or "",
            raw_text=raw_text or "",
            structured_data=structured_data,
        )
        db.add(row)
        await db.commit()
        await db.refresh(row)
        return row.id


async def get_parsed_data(session_id: int) -> list[dict]:
    async with get_sessionmaker()() as db:
        res = await db.execute(
            select(ParsedData)
            .where(ParsedData.quiz_id == session_id)
            .order_by(ParsedData.data_type, ParsedData.created_at)
        )
        return [_row(p) for p in res.scalars().all()]


async def delete_parsed_data(session_id: int, data_type: Optional[str] = None):
    async with get_sessionmaker()() as db:
        stmt = delete(ParsedData).where(ParsedData.quiz_id == session_id)
        if data_type:
            stmt = stmt.where(ParsedData.data_type == data_type)
        await db.execute(stmt)
        await db.commit()


# ---- Answer grid (stored as a parsed_data row) ----

async def save_answer_grid(session_id: int, grid_dict: dict) -> int:
    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]:
    async with get_sessionmaker()() as db:
        res = await db.execute(
            select(ParsedData.structured_data)
            .where(ParsedData.quiz_id == session_id)
            .where(ParsedData.data_type == ANSWER_GRID_DATA_TYPE)
            .order_by(ParsedData.created_at.desc())
            .limit(1)
        )
        row = res.scalar_one_or_none()
        return row if row else None


# =============================================================================
# Raw files (metadata; bytes live in Supabase Storage)
# =============================================================================

async def save_raw_file(
    session_id: int,
    data_type: str,
    file_name: str,
    content_type: str,
    size_bytes: int,
    storage_path: str,
) -> int:
    async with get_sessionmaker()() as db:
        row = RawFile(
            quiz_id=session_id,
            data_type=data_type,
            file_name=file_name,
            content_type=content_type,
            size_bytes=size_bytes,
            storage_path=storage_path,
        )
        db.add(row)
        await db.commit()
        await db.refresh(row)
        return row.id


async def get_raw_files(session_id: int) -> list[dict]:
    async with get_sessionmaker()() as db:
        res = await db.execute(
            select(RawFile).where(RawFile.quiz_id == session_id).order_by(RawFile.created_at)
        )
        return [_row(f) for f in res.scalars().all()]


# =============================================================================
# Students (first-class, matched across quizzes by normalized name)
# =============================================================================

async def match_or_create_student(
    teacher_id: int, name: str, external_id: Optional[str] = None
) -> int:
    """Return the student id for this teacher, creating the row if new."""
    norm = normalize_name(name)
    async with get_sessionmaker()() as db:
        res = await db.execute(
            select(Student)
            .where(Student.teacher_id == teacher_id)
            .where(Student.normalized_name == norm)
        )
        student = res.scalar_one_or_none()
        if student:
            if external_id and not student.external_id:
                student.external_id = external_id
                await db.commit()
            return student.id
        student = Student(
            teacher_id=teacher_id, name=name, normalized_name=norm, external_id=external_id
        )
        db.add(student)
        await db.commit()
        await db.refresh(student)
        return student.id


async def get_students(teacher_id: int) -> list[dict]:
    async with get_sessionmaker()() as db:
        res = await db.execute(
            select(Student).where(Student.teacher_id == teacher_id).order_by(Student.name)
        )
        return [_row(s) for s in res.scalars().all()]


async def get_student(student_id: int) -> Optional[dict]:
    async with get_sessionmaker()() as db:
        student = await db.get(Student, student_id)
        return _row(student) if student else None


async def save_student_result(
    quiz_id: int,
    student_id: int,
    answers: dict,
    total_questions: int,
    correct_count: int,
    score: float,
) -> int:
    """Upsert one student's result for a quiz (unique per quiz+student)."""
    async with get_sessionmaker()() as db:
        res = await db.execute(
            select(StudentResult)
            .where(StudentResult.quiz_id == quiz_id)
            .where(StudentResult.student_id == student_id)
        )
        existing = res.scalar_one_or_none()
        if existing:
            existing.answers = answers
            existing.total_questions = total_questions
            existing.correct_count = correct_count
            existing.score = score
            await db.commit()
            return existing.id
        row = StudentResult(
            quiz_id=quiz_id,
            student_id=student_id,
            answers=answers,
            total_questions=total_questions,
            correct_count=correct_count,
            score=score,
        )
        db.add(row)
        await db.commit()
        await db.refresh(row)
        return row.id


async def get_student_timeline(student_id: int) -> list[dict]:
    """One student's results across all quizzes, oldest first — the trajectory."""
    async with get_sessionmaker()() as db:
        res = await db.execute(
            select(StudentResult, Quiz.title, Quiz.subject, Quiz.created_at)
            .join(Quiz, StudentResult.quiz_id == Quiz.id)
            .where(StudentResult.student_id == student_id)
            .order_by(Quiz.created_at)
        )
        out = []
        for sr, title, subject, created in res.all():
            d = _row(sr)
            d["quiz_title"] = title
            d["quiz_subject"] = subject
            d["quiz_created_at"] = created
            out.append(d)
        return out


# =============================================================================
# Prompts
# =============================================================================

async def save_prompt(user_id: int, name: str, content: str, is_default: bool = False) -> int:
    async with get_sessionmaker()() as db:
        row = Prompt(teacher_id=user_id, name=name, content=content, is_default=is_default)
        db.add(row)
        await db.commit()
        await db.refresh(row)
        return row.id


async def get_prompts(user_id: int) -> list[dict]:
    async with get_sessionmaker()() as db:
        res = await db.execute(
            select(Prompt)
            .where((Prompt.teacher_id == user_id) | (Prompt.is_default.is_(True)))
            .order_by(Prompt.is_default.desc(), Prompt.updated_at.desc())
        )
        return [_row(p) for p in res.scalars().all()]


async def get_all_prompts() -> list[dict]:
    async with get_sessionmaker()() as db:
        res = await db.execute(
            select(Prompt, Teacher.email)
            .join(Teacher, Prompt.teacher_id == Teacher.id)
            .order_by(Prompt.updated_at.desc())
        )
        out = []
        for p, email in res.all():
            d = _row(p)
            d["author_email"] = email
            out.append(d)
        return out


async def get_prompt(prompt_id: int) -> Optional[dict]:
    async with get_sessionmaker()() as db:
        p = await db.get(Prompt, prompt_id)
        if not p:
            return None
        d = _row(p)
        d["user_id"] = p.teacher_id  # back-compat key
        return d


async def update_prompt(prompt_id: int, name: str, content: str):
    async with get_sessionmaker()() as db:
        await db.execute(
            update(Prompt)
            .where(Prompt.id == prompt_id)
            .values(name=name, content=content, updated_at=_now())
        )
        await db.commit()


async def delete_prompt(prompt_id: int):
    async with get_sessionmaker()() as db:
        await db.execute(delete(Prompt).where(Prompt.id == prompt_id))
        await db.commit()


# =============================================================================
# Reports
# =============================================================================

async def save_report(
    session_id: int,
    prompt_id: Optional[int],
    html_content: str,
    student_id: Optional[int] = None,
    model: str = "",
) -> int:
    async with get_sessionmaker()() as db:
        row = Report(
            quiz_id=session_id,
            prompt_id=prompt_id,
            student_id=student_id,
            html_content=html_content,
            model=model,
        )
        db.add(row)
        await db.commit()
        await db.refresh(row)
        return row.id


async def get_reports(session_id: int) -> list[dict]:
    async with get_sessionmaker()() as db:
        res = await db.execute(
            select(Report).where(Report.quiz_id == session_id).order_by(Report.created_at.desc())
        )
        return [_row(r) for r in res.scalars().all()]


async def get_report(report_id: int) -> Optional[dict]:
    async with get_sessionmaker()() as db:
        r = await db.get(Report, report_id)
        return _row(r) if r else None


# =============================================================================
# Analytics (admin dashboard)
# =============================================================================

async def get_stats() -> dict:
    async with get_sessionmaker()() as db:
        async def count(model):
            return (await db.execute(select(func.count(model.id)))).scalar_one()

        return {
            "teachers": await count(Teacher),
            "students": await count(Student),
            "quizzes": await count(Quiz),
            "reports": await count(Report),
            "student_results": await count(StudentResult),
            "prompts": await count(Prompt),
        }


async def admin_list_students(limit: int = 500) -> list[dict]:
    """All students across teachers, with quiz count + average score (admin)."""
    async with get_sessionmaker()() as db:
        res = await db.execute(
            select(Student, Teacher.email).join(Teacher, Student.teacher_id == Teacher.id).order_by(Student.name).limit(limit)
        )
        out = []
        for s, email in res.all():
            d = _row(s)
            d["teacher_email"] = email
            agg = (
                await db.execute(
                    select(func.count(StudentResult.id), func.avg(StudentResult.score)).where(
                        StudentResult.student_id == s.id
                    )
                )
            ).one()
            d["result_count"] = agg[0]
            d["avg_score"] = round(float(agg[1]), 1) if agg[1] is not None else None
            out.append(d)
        return out


async def get_schema_version() -> Optional[str]:
    """Current Alembic migration revision applied to the DB."""
    from sqlalchemy import text

    from .db import active_schema

    schema = active_schema()
    table = f'"{schema}".alembic_version' if schema else "alembic_version"
    async with get_sessionmaker()() as db:
        try:
            res = await db.execute(text(f"SELECT version_num FROM {table} LIMIT 1"))
            row = res.first()
            return row[0] if row else None
        except Exception:
            return None


async def admin_list_sessions(limit: int = 200) -> list[dict]:
    """All quizzes across teachers, with author email + counts (admin)."""
    async with get_sessionmaker()() as db:
        res = await db.execute(
            select(Quiz, Teacher.email)
            .join(Teacher, Quiz.teacher_id == Teacher.id)
            .order_by(Quiz.created_at.desc())
            .limit(limit)
        )
        out = []
        for q, email in res.all():
            d = _row(q)
            d["teacher_email"] = email
            d["report_count"] = (
                await db.execute(select(func.count(Report.id)).where(Report.quiz_id == q.id))
            ).scalar_one()
            out.append(d)
        return out