SEC-cyBERT/labelapp/app/api/label/route.ts

import { NextResponse } from "next/server";
import { db } from "@/db";
import {
  assignments,
  humanLabels,
  paragraphs,
  quizSessions,
} from "@/db/schema";
import { eq, and, desc, isNull, sql } from "drizzle-orm";
import { getSession } from "@/lib/auth";
import { WARMUP_PARAGRAPHS } from "@/lib/warmup-paragraphs";

const VALID_CATEGORIES = [
  "Board Governance",
  "Management Role",
  "Risk Management Process",
  "Third-Party Risk",
  "Incident Disclosure",
  "Strategy Integration",
  "None/Other",
] as const;

const VALID_SPECIFICITY = [1, 2, 3, 4] as const;

async function checkWarmupComplete(annotatorId: string, sessionCreatedAt: number): Promise<boolean> {
  const [quizSession] = await db
    .select()
    .from(quizSessions)
    .where(
      and(
        eq(quizSessions.annotatorId, annotatorId),
        eq(quizSessions.passed, true),
      ),
    )
    .orderBy(desc(quizSessions.startedAt))
    .limit(1);

  if (!quizSession) return false;

  try {
    const parsed = JSON.parse(quizSession.answers);
    if (typeof parsed === "object" && !Array.isArray(parsed) && parsed.warmupBySession) {
      const sessionKey = String(sessionCreatedAt);
      return (parsed.warmupBySession[sessionKey] ?? 0) >= WARMUP_PARAGRAPHS.length;
    }
    return false;
  } catch {
    return false;
  }
}

export async function GET() {
  const session = await getSession();
  if (!session) {
    return NextResponse.json({ error: "Unauthorized" }, { status: 401 });
  }

  const warmupDone = await checkWarmupComplete(session.annotatorId, session.createdAt);
  if (!warmupDone) {
    return NextResponse.json({ redirectToWarmup: true });
  }

  // Find first assigned paragraph that doesn't have a label from this annotator
  const result = await db
    .select({
      paragraphId: assignments.paragraphId,
      assignedAt: assignments.assignedAt,
      text: paragraphs.text,
      wordCount: paragraphs.wordCount,
      paragraphIndex: paragraphs.paragraphIndex,
      companyName: paragraphs.companyName,
      ticker: paragraphs.ticker,
      filingType: paragraphs.filingType,
      filingDate: paragraphs.filingDate,
      secItem: paragraphs.secItem,
      labelId: humanLabels.id,
    })
    .from(assignments)
    .innerJoin(paragraphs, eq(assignments.paragraphId, paragraphs.id))
    .leftJoin(
      humanLabels,
      and(
        eq(humanLabels.paragraphId, assignments.paragraphId),
        eq(humanLabels.annotatorId, session.annotatorId),
      ),
    )
    .where(eq(assignments.annotatorId, session.annotatorId))
    .orderBy(assignments.assignedAt)
    .limit(1000);

  const total = result.length;
  const completed = result.filter((r) => r.labelId !== null).length;
  const next = result.find((r) => r.labelId === null);

  if (!next) {
    return NextResponse.json({
      done: true,
      progress: { completed, total },
    });
  }

  return NextResponse.json({
    paragraph: {
      id: next.paragraphId,
      text: next.text,
      wordCount: next.wordCount,
      paragraphIndex: next.paragraphIndex,
      companyName: next.companyName,
      ticker: next.ticker,
      filingType: next.filingType,
      filingDate: next.filingDate,
      secItem: next.secItem,
    },
    progress: { completed, total },
  });
}

export async function POST(request: Request) {
  const session = await getSession();
  if (!session) {
    return NextResponse.json({ error: "Unauthorized" }, { status: 401 });
  }

  const body = await request.json();
  const { paragraphId, contentCategory, specificityLevel, notes, sessionId, durationMs, activeMs } =
    body as {
      paragraphId?: string;
      contentCategory?: string;
      specificityLevel?: number;
      notes?: string;
      sessionId?: string;
      durationMs?: number;
      activeMs?: number;
    };

  if (!paragraphId || !contentCategory || !specificityLevel || !sessionId) {
    return NextResponse.json(
      { error: "Missing required fields" },
      { status: 400 },
    );
  }

  if (!VALID_CATEGORIES.includes(contentCategory as (typeof VALID_CATEGORIES)[number])) {
    return NextResponse.json(
      { error: "Invalid content category" },
      { status: 400 },
    );
  }

  if (!VALID_SPECIFICITY.includes(specificityLevel as (typeof VALID_SPECIFICITY)[number])) {
    return NextResponse.json(
      { error: "Invalid specificity level" },
      { status: 400 },
    );
  }

  // Verify assignment exists
  const [assignment] = await db
    .select()
    .from(assignments)
    .where(
      and(
        eq(assignments.paragraphId, paragraphId),
        eq(assignments.annotatorId, session.annotatorId),
      ),
    )
    .limit(1);

  if (!assignment) {
    return NextResponse.json(
      { error: "No assignment found for this paragraph" },
      { status: 403 },
    );
  }

  // Check for duplicate
  const [existing] = await db
    .select({ id: humanLabels.id })
    .from(humanLabels)
    .where(
      and(
        eq(humanLabels.paragraphId, paragraphId),
        eq(humanLabels.annotatorId, session.annotatorId),
      ),
    )
    .limit(1);

  if (existing) {
    return NextResponse.json(
      { error: "Already labeled this paragraph" },
      { status: 409 },
    );
  }

  await db.insert(humanLabels).values({
    paragraphId,
    annotatorId: session.annotatorId,
    contentCategory,
    specificityLevel,
    notes: notes || null,
    sessionId,
    durationMs: durationMs ?? null,
    activeMs: activeMs ?? null,
  });

  // Get updated progress
  const countResult = await db
    .select({
      total: sql<number>`count(*)`,
      completed: sql<number>`count(${humanLabels.id})`,
    })
    .from(assignments)
    .leftJoin(
      humanLabels,
      and(
        eq(humanLabels.paragraphId, assignments.paragraphId),
        eq(humanLabels.annotatorId, session.annotatorId),
      ),
    )
    .where(eq(assignments.annotatorId, session.annotatorId));

  const progress = {
    completed: Number(countResult[0]?.completed ?? 0),
    total: Number(countResult[0]?.total ?? 0),
  };

  return NextResponse.json({ ok: true, progress });
}