app/db.py

"""SQLite persistence for the legacy pilot and modality-revision study.

Migrations are additive: the original ``exchanges`` and ``annotations`` rows
remain exportable. New study items and their six predeclared condition tasks
live in separate tables, so deploying a new manifest never resets pilot data.
"""

from __future__ import annotations

import json
import hashlib
import os
from datetime import datetime, timedelta, timezone
from pathlib import Path

from sqlalchemy import (
    Boolean,
    DateTime,
    ForeignKey,
    Integer,
    String,
    Text,
    UniqueConstraint,
    create_engine,
    func,
    select,
    text,
)
from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column, sessionmaker

DB_PATH = os.environ.get("ANNOTATOR_DB_PATH", "/data/annotator.db")
Path(DB_PATH).parent.mkdir(parents=True, exist_ok=True)
engine = create_engine(f"sqlite:///{DB_PATH}", connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(bind=engine, expire_on_commit=False)


class Base(DeclarativeBase):
    pass


def _utcnow() -> datetime:
    return datetime.now(timezone.utc).replace(tzinfo=None)


class Exchange(Base):
    """Legacy pilot item. Kept intact for backward-compatible export."""

    __tablename__ = "exchanges"
    id: Mapped[str] = mapped_column(String, primary_key=True)
    call_id: Mapped[str] = mapped_column(String, index=True)
    exchange_index: Mapped[int] = mapped_column(Integer)
    year: Mapped[int | None] = mapped_column(Integer, nullable=True)
    question_text: Mapped[str] = mapped_column(Text)
    answer_text: Mapped[str] = mapped_column(Text)
    audio_clip_filename: Mapped[str] = mapped_column(String)
    duration_s: Mapped[float | None] = mapped_column(Integer, nullable=True)
    bootstrap_rasiah: Mapped[str | None] = mapped_column(String, nullable=True)
    flagged_broken: Mapped[bool] = mapped_column(Boolean, default=False)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=_utcnow)


class Annotator(Base):
    __tablename__ = "annotators"
    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    email: Mapped[str] = mapped_column(String, unique=True, index=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=_utcnow)
    current_exchange_id: Mapped[str | None] = mapped_column(ForeignKey("exchanges.id"), nullable=True)
    finance_familiarity: Mapped[str | None] = mapped_column(String, nullable=True)
    english_proficiency: Mapped[str | None] = mapped_column(String, nullable=True)
    consented_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
    training_passed_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
    training_attempts: Mapped[int] = mapped_column(Integer, default=0)
    current_task_id: Mapped[int | None] = mapped_column(Integer, nullable=True)


class StudyItem(Base):
    __tablename__ = "study_items"
    id: Mapped[str] = mapped_column(String, primary_key=True)
    study_id: Mapped[str] = mapped_column(String, index=True)
    partition: Mapped[str] = mapped_column(String, index=True)
    call_id: Mapped[str] = mapped_column(String, index=True)
    exchange_index: Mapped[int] = mapped_column(Integer)
    question_text: Mapped[str] = mapped_column(Text)
    answer_text: Mapped[str] = mapped_column(Text)
    audio_filename: Mapped[str | None] = mapped_column(String, nullable=True)
    duration_s: Mapped[int | None] = mapped_column(Integer, nullable=True)
    episode_reconstructed: Mapped[bool] = mapped_column(Boolean, default=False)
    eligibility_warning: Mapped[str | None] = mapped_column(Text, nullable=True)
    active: Mapped[bool] = mapped_column(Boolean, default=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=_utcnow)


class StudyTask(Base):
    __tablename__ = "study_tasks"
    __table_args__ = (UniqueConstraint("item_id", "slot", name="uq_study_item_slot"),)
    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    item_id: Mapped[str] = mapped_column(ForeignKey("study_items.id"), index=True)
    condition: Mapped[str] = mapped_column(String, index=True)
    slot: Mapped[int] = mapped_column(Integer)
    display_order_seed: Mapped[int] = mapped_column(Integer)
    assigned_annotator_id: Mapped[int | None] = mapped_column(ForeignKey("annotators.id"), nullable=True, index=True)
    assigned_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True)
    completed: Mapped[bool] = mapped_column(Boolean, default=False, index=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=_utcnow)


class Annotation(Base):
    __tablename__ = "annotations"
    __table_args__ = (
        UniqueConstraint("exchange_id", "annotator_id", "condition", name="uq_one_label_per_annotator"),
    )
    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    exchange_id: Mapped[str | None] = mapped_column(ForeignKey("exchanges.id"), nullable=True, index=True)
    annotator_id: Mapped[int] = mapped_column(ForeignKey("annotators.id"), index=True)
    condition: Mapped[str] = mapped_column(String, default="text_audio")
    idk: Mapped[bool] = mapped_column(Boolean, default=False)
    flagged_broken: Mapped[bool] = mapped_column(Boolean, default=False)
    flag_reason: Mapped[str | None] = mapped_column(String, nullable=True)
    rasiah: Mapped[str | None] = mapped_column(String, nullable=True)
    bavelas: Mapped[str | None] = mapped_column(String, nullable=True)
    bull: Mapped[str | None] = mapped_column(String, nullable=True)
    created_at: Mapped[datetime] = mapped_column(DateTime, default=_utcnow)

    # Additive modality-revision fields.
    task_id: Mapped[int | None] = mapped_column(ForeignKey("study_tasks.id"), nullable=True, unique=True, index=True)
    study_item_id: Mapped[str | None] = mapped_column(ForeignKey("study_items.id"), nullable=True, index=True)
    gate: Mapped[str | None] = mapped_column(String, nullable=True)
    gate_confidence: Mapped[int | None] = mapped_column(Integer, nullable=True)
    responsiveness: Mapped[int | None] = mapped_column(Integer, nullable=True)
    supplied_information: Mapped[str | None] = mapped_column(String, nullable=True)
    confidence: Mapped[int | None] = mapped_column(Integer, nullable=True)
    delivery_descriptors_json: Mapped[str | None] = mapped_column(Text, nullable=True)
    audio_revision: Mapped[str | None] = mapped_column(String, nullable=True)
    audible_event: Mapped[str | None] = mapped_column(String, nullable=True)
    audible_event_other: Mapped[str | None] = mapped_column(Text, nullable=True)
    rationale: Mapped[str | None] = mapped_column(Text, nullable=True)
    response_time_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)
    audio_played_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)
    audio_completed: Mapped[bool | None] = mapped_column(Boolean, nullable=True)


def _existing_columns(table_name: str) -> set[str]:
    with engine.connect() as connection:
        rows = connection.execute(text(f"PRAGMA table_info({table_name})")).mappings()
        return {str(row["name"]) for row in rows}


def _add_column(table: str, name: str, ddl: str) -> None:
    if name in _existing_columns(table):
        return
    with engine.begin() as connection:
        connection.execute(text(f"ALTER TABLE {table} ADD COLUMN {name} {ddl}"))


def init_db() -> None:
    """Create new tables and add columns to an existing pilot database."""
    Base.metadata.create_all(engine)
    for name, ddl in {
        "finance_familiarity": "VARCHAR",
        "english_proficiency": "VARCHAR",
        "consented_at": "DATETIME",
        "training_passed_at": "DATETIME",
        "training_attempts": "INTEGER DEFAULT 0",
        "current_task_id": "INTEGER",
    }.items():
        _add_column("annotators", name, ddl)
    for name, ddl in {
        "task_id": "INTEGER",
        "study_item_id": "VARCHAR",
        "gate": "VARCHAR",
        "gate_confidence": "INTEGER",
        "responsiveness": "INTEGER",
        "supplied_information": "VARCHAR",
        "confidence": "INTEGER",
        "delivery_descriptors_json": "TEXT",
        "audio_revision": "VARCHAR",
        "audible_event": "VARCHAR",
        "audible_event_other": "TEXT",
        "rationale": "TEXT",
        "response_time_ms": "INTEGER",
        "audio_played_ms": "INTEGER",
        "audio_completed": "BOOLEAN",
    }.items():
        _add_column("annotations", name, ddl)
    with engine.begin() as connection:
        connection.execute(
            text(
                "CREATE UNIQUE INDEX IF NOT EXISTS ux_annotations_task_id "
                "ON annotations(task_id) WHERE task_id IS NOT NULL"
            )
        )


def get_session() -> Session:
    return SessionLocal()


def get_annotator(session: Session, identity: str) -> Annotator | None:
    return session.execute(
        select(Annotator).where(Annotator.email == identity.strip().lower())
    ).scalar_one_or_none()


def enroll_annotator(
    session: Session,
    identity: str,
    finance_familiarity: str,
    english_proficiency: str,
) -> Annotator:
    identity = identity.strip().lower()
    row = get_annotator(session, identity)
    if row is None:
        row = Annotator(email=identity)
        session.add(row)
        session.flush()
    row.finance_familiarity = finance_familiarity
    row.english_proficiency = english_proficiency
    row.consented_at = row.consented_at or _utcnow()
    session.commit()
    session.refresh(row)
    return row


def record_training_attempt(session: Session, annotator: Annotator, passed: bool) -> None:
    annotator.training_attempts = (annotator.training_attempts or 0) + 1
    if passed:
        annotator.training_passed_at = _utcnow()
    session.commit()


def seed_study(session: Session, manifest_path: Path, assignments_path: Path) -> dict:
    """Idempotently upsert study items/tasks; never delete or reset labels."""
    if not manifest_path.exists() or not assignments_path.exists():
        return {"items": 0, "tasks": 0, "warning": "study package missing"}
    manifest = json.loads(manifest_path.read_text())
    items = manifest["items"] if isinstance(manifest, dict) else manifest
    assignments = json.loads(assignments_path.read_text())
    for payload in items:
        row = session.get(StudyItem, payload["study_item_id"])
        if row is None:
            row = StudyItem(id=payload["study_item_id"])
            session.add(row)
        row.study_id = manifest.get("study_id", "modality_revision_v1")
        row.partition = payload["partition"]
        row.call_id = str(payload["call_id"])
        row.exchange_index = int(payload["exchange_index"])
        row.question_text = payload["question"]
        row.answer_text = payload["answer"]
        row.audio_filename = payload.get("audio_filename")
        row.duration_s = round(float(payload.get("duration_s") or 0))
        row.episode_reconstructed = bool(payload.get("episode_reconstructed"))
        row.eligibility_warning = payload.get("eligibility_warning")
        row.active = bool(payload.get("active", True))
        # Mirror the key into the legacy table because the original deployed
        # SQLite schema made annotations.exchange_id NOT NULL. Keeping this
        # additive avoids a destructive table rebuild during migration.
        legacy = session.get(Exchange, payload["study_item_id"])
        if legacy is None:
            session.add(
                Exchange(
                    id=payload["study_item_id"],
                    call_id=str(payload["call_id"]),
                    exchange_index=int(payload["exchange_index"]),
                    year=None,
                    question_text=payload["question"],
                    answer_text=payload["answer"],
                    audio_clip_filename=payload.get("audio_filename") or "",
                    duration_s=round(float(payload.get("duration_s") or 0)),
                    bootstrap_rasiah=None,
                )
            )
    session.flush()
    for payload in assignments:
        existing = session.execute(
            select(StudyTask).where(
                StudyTask.item_id == payload["study_item_id"],
                StudyTask.slot == int(payload["slot"]),
            )
        ).scalar_one_or_none()
        if existing is None:
            session.add(
                StudyTask(
                    item_id=payload["study_item_id"],
                    condition=payload["condition"],
                    slot=int(payload["slot"]),
                    display_order_seed=int(payload["display_order_seed"]),
                )
            )
        elif not existing.completed:
            existing.condition = payload["condition"]
            existing.display_order_seed = int(payload["display_order_seed"])
    session.commit()
    return {
        "items": session.scalar(select(func.count(StudyItem.id))) or 0,
        "tasks": session.scalar(select(func.count(StudyTask.id))) or 0,
    }


def claim_task(session: Session, annotator: Annotator, phase: str) -> tuple[StudyTask, StudyItem] | None:
    """Claim one task while preventing cross-condition repeat exposure."""
    if annotator.current_task_id:
        current = session.get(StudyTask, annotator.current_task_id)
        if current and not current.completed and current.assigned_annotator_id == annotator.id:
            item = session.get(StudyItem, current.item_id)
            if item:
                return current, item
        annotator.current_task_id = None
        session.commit()

    # Recover abandoned browser assignments after eight hours.
    expiry = _utcnow() - timedelta(hours=8)
    stale = session.scalars(
        select(StudyTask).where(
            StudyTask.completed.is_(False),
            StudyTask.assigned_at.is_not(None),
            StudyTask.assigned_at < expiry,
        )
    ).all()
    for task in stale:
        owner = session.get(Annotator, task.assigned_annotator_id) if task.assigned_annotator_id else None
        if owner and owner.current_task_id == task.id:
            owner.current_task_id = None
        task.assigned_annotator_id = None
        task.assigned_at = None
    if stale:
        session.commit()

    seen_items = select(StudyTask.item_id).where(StudyTask.assigned_annotator_id == annotator.id)
    partitions = {
        "development": ["development"],
        "confirmatory": ["locked_confirmatory"],
        "all": ["development", "locked_confirmatory"],
    }.get(phase, ["development"])
    statement = (
        select(StudyTask, StudyItem)
        .join(StudyItem, StudyTask.item_id == StudyItem.id)
        .where(
            StudyTask.completed.is_(False),
            StudyTask.assigned_annotator_id.is_(None),
            StudyItem.active.is_(True),
            StudyItem.partition.in_(partitions),
            StudyTask.item_id.notin_(seen_items),
        )
        .order_by(((StudyTask.display_order_seed + annotator.id * 7919) % 100000).asc())
        .limit(1)
    )
    result = session.execute(statement).first()
    if result is None:
        return None
    task, item = result
    task.assigned_annotator_id = annotator.id
    task.assigned_at = _utcnow()
    annotator.current_task_id = task.id
    session.commit()
    return task, item


def complete_task(session: Session, annotator: Annotator, task_id: int, clean: dict) -> bool:
    task = session.get(StudyTask, task_id)
    if (
        task is None
        or task.completed
        or task.assigned_annotator_id != annotator.id
        or annotator.current_task_id != task.id
    ):
        return False
    annotation = Annotation(
        exchange_id=task.item_id,
        annotator_id=annotator.id,
        condition=task.condition,
        task_id=task.id,
        study_item_id=task.item_id,
        **clean,
    )
    session.add(annotation)
    task.completed = True
    annotator.current_task_id = None
    session.commit()
    return True


def flag_task(session: Session, annotator: Annotator, task_id: int, reason: str) -> bool:
    task = session.get(StudyTask, task_id)
    if task is None or task.assigned_annotator_id != annotator.id or task.completed:
        return False
    session.add(
        Annotation(
            exchange_id=task.item_id,
            annotator_id=annotator.id,
            condition=task.condition,
            task_id=task.id,
            study_item_id=task.item_id,
            flagged_broken=True,
            flag_reason=reason[:500],
        )
    )
    task.completed = True
    annotator.current_task_id = None
    session.commit()
    return True


def study_progress(session: Session, annotator_id: int, phase: str) -> dict:
    partitions = {
        "development": ["development"],
        "confirmatory": ["locked_confirmatory"],
        "all": ["development", "locked_confirmatory"],
    }.get(phase, ["development"])
    total = session.scalar(
        select(func.count(StudyTask.id))
        .join(StudyItem, StudyTask.item_id == StudyItem.id)
        .where(StudyItem.partition.in_(partitions), StudyItem.active.is_(True))
    ) or 0
    completed = session.scalar(
        select(func.count(StudyTask.id))
        .join(StudyItem, StudyTask.item_id == StudyItem.id)
        .where(
            StudyItem.partition.in_(partitions),
            StudyItem.active.is_(True),
            StudyTask.completed.is_(True),
        )
    ) or 0
    own = session.scalar(
        select(func.count(StudyTask.id)).where(
            StudyTask.assigned_annotator_id == annotator_id,
            StudyTask.completed.is_(True),
        )
    ) or 0
    return {"own_done": own, "study_done": completed, "study_total": total}


def export_payload(session: Session) -> dict:
    annotators = {row.id: row for row in session.scalars(select(Annotator)).all()}
    rows = []
    for annotation in session.scalars(select(Annotation)).all():
        annotator = annotators.get(annotation.annotator_id)
        identity = annotator.email if annotator else "unknown"
        salt = os.environ.get("ANNOTATOR_EXPORT_SALT") or os.environ.get("ANNOTATOR_SECRET", "local")
        participant_hash = hashlib.sha256(f"{salt}|{identity}".encode()).hexdigest()[:20]
        rows.append(
            {
                "annotation_id": annotation.id,
                "legacy_exchange_id": annotation.exchange_id,
                "study_item_id": annotation.study_item_id,
                "task_id": annotation.task_id,
                "participant_hash": participant_hash,
                "finance_familiarity": annotator.finance_familiarity if annotator else None,
                "english_proficiency": annotator.english_proficiency if annotator else None,
                "training_attempts": annotator.training_attempts if annotator else None,
                "training_passed_at": (
                    annotator.training_passed_at.isoformat()
                    if annotator and annotator.training_passed_at
                    else None
                ),
                "condition": annotation.condition,
                "idk": annotation.idk,
                "flagged_broken": annotation.flagged_broken,
                "flag_reason": annotation.flag_reason,
                "gate": annotation.gate,
                "gate_confidence": annotation.gate_confidence,
                "responsiveness": annotation.responsiveness,
                "rasiah": annotation.rasiah,
                "bavelas": annotation.bavelas,
                "bull": annotation.bull,
                "supplied_information": annotation.supplied_information,
                "confidence": annotation.confidence,
                "delivery_descriptors": (
                    json.loads(annotation.delivery_descriptors_json)
                    if annotation.delivery_descriptors_json
                    else None
                ),
                "audio_revision": annotation.audio_revision,
                "audible_event": annotation.audible_event,
                "audible_event_other": annotation.audible_event_other,
                "rationale": annotation.rationale,
                "response_time_ms": annotation.response_time_ms,
                "audio_played_ms": annotation.audio_played_ms,
                "audio_completed": annotation.audio_completed,
                "created_at": annotation.created_at.isoformat(),
            }
        )
    return {"schema_version": "annotation-export-v2", "annotations": rows}