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

import { NextResponse } from "next/server";
import { db } from "@/db";
import {
  adjudications,
  annotators,
  humanLabels,
  paragraphs,
} from "@/db/schema";
import { eq } from "drizzle-orm";
import { getSession } from "@/lib/auth";

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;

interface LabelRecord {
  annotatorId: string;
  contentCategory: string;
  specificityLevel: number;
  notes: string | null;
}

interface QueueItem {
  paragraphId: string;
  paragraphText: string;
  labels: {
    annotatorId: string;
    displayName: string;
    category: string;
    specificity: number;
    notes: string | null;
  }[];
  stage1Category: string | null;
  stage1Specificity: number | null;
  splitSeverity: number; // 3 = three-way split, 2 = two-way split
}

function findMajority(values: string[]): string | null {
  const counts = new Map<string, number>();
  for (const v of values) {
    counts.set(v, (counts.get(v) ?? 0) + 1);
  }
  for (const [val, count] of counts) {
    if (count >= 2) return val;
  }
  return null;
}

function findMajorityNum(values: number[]): number | null {
  const counts = new Map<number, number>();
  for (const v of values) {
    counts.set(v, (counts.get(v) ?? 0) + 1);
  }
  for (const [val, count] of counts) {
    if (count >= 2) return val;
  }
  return null;
}

function countDistinct(values: string[]): number {
  return new Set(values).size;
}

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

  if (session.annotatorId !== "joey") {
    return NextResponse.json({ error: "Admin access required" }, { status: 403 });
  }

  // Load all human labels
  const allLabels = await db.select().from(humanLabels);

  // Load all existing adjudications
  const allAdj = await db
    .select({ paragraphId: adjudications.paragraphId })
    .from(adjudications);
  const adjudicatedSet = new Set(allAdj.map((a) => a.paragraphId));

  // Load all paragraphs (for text and stage1 data)
  const allParagraphs = await db.select().from(paragraphs);
  const paragraphMap = new Map(allParagraphs.map((p) => [p.id, p]));

  // Load all annotators (for display names)
  const allAnnotators = await db.select().from(annotators);
  const annotatorMap = new Map(allAnnotators.map((a) => [a.id, a.displayName]));

  // Group labels by paragraph
  const byParagraph = new Map<string, LabelRecord[]>();
  for (const label of allLabels) {
    if (!byParagraph.has(label.paragraphId)) {
      byParagraph.set(label.paragraphId, []);
    }
    byParagraph.get(label.paragraphId)!.push({
      annotatorId: label.annotatorId,
      contentCategory: label.contentCategory,
      specificityLevel: label.specificityLevel,
      notes: label.notes,
    });
  }

  const queue: QueueItem[] = [];

  for (const [paragraphId, labels] of byParagraph) {
    // Only process paragraphs with 3+ labels that aren't already adjudicated
    if (labels.length < 3 || adjudicatedSet.has(paragraphId)) continue;

    const categories = labels.map((l) => l.contentCategory);
    const specificities = labels.map((l) => l.specificityLevel);

    const allCategoriesAgree = categories.every((c) => c === categories[0]);
    const allSpecificitiesAgree = specificities.every(
      (s) => s === specificities[0],
    );

    if (allCategoriesAgree && allSpecificitiesAgree) {
      // Full consensus - auto-resolve
      await db.insert(adjudications).values({
        paragraphId,
        finalCategory: categories[0],
        finalSpecificity: specificities[0],
        method: "consensus",
        adjudicatorId: null,
        notes: null,
      });
      continue;
    }

    const majorityCategory = findMajority(categories);
    const majoritySpecificity = findMajorityNum(specificities);

    if (majorityCategory !== null && majoritySpecificity !== null) {
      // 2/3 majority on both dims - auto-resolve
      await db.insert(adjudications).values({
        paragraphId,
        finalCategory: majorityCategory,
        finalSpecificity: majoritySpecificity,
        method: "majority",
        adjudicatorId: null,
        notes: null,
      });
      continue;
    }

    // Needs manual adjudication
    const para = paragraphMap.get(paragraphId);
    if (!para) continue;

    // Severity: 3-way category split is worse than 2-way
    const categorySplitCount = countDistinct(categories);

    queue.push({
      paragraphId,
      paragraphText: para.text,
      labels: labels.map((l) => ({
        annotatorId: l.annotatorId,
        displayName: annotatorMap.get(l.annotatorId) ?? l.annotatorId,
        category: l.contentCategory,
        specificity: l.specificityLevel,
        notes: l.notes,
      })),
      stage1Category: para.stage1Category,
      stage1Specificity: para.stage1Specificity,
      splitSeverity: categorySplitCount,
    });
  }

  // Sort: 3-way splits first, then 2-way
  queue.sort((a, b) => b.splitSeverity - a.splitSeverity);

  return NextResponse.json({ queue });
}

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

  if (session.annotatorId !== "joey") {
    return NextResponse.json({ error: "Admin access required" }, { status: 403 });
  }

  const body = await request.json();
  const { paragraphId, finalCategory, finalSpecificity, notes } = body as {
    paragraphId?: string;
    finalCategory?: string;
    finalSpecificity?: number;
    notes?: string;
  };

  if (!paragraphId || !finalCategory || finalSpecificity === undefined) {
    return NextResponse.json(
      { error: "Missing required fields" },
      { status: 400 },
    );
  }

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

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

  // Check paragraph exists
  const [para] = await db
    .select({ id: paragraphs.id })
    .from(paragraphs)
    .where(eq(paragraphs.id, paragraphId))
    .limit(1);

  if (!para) {
    return NextResponse.json(
      { error: "Paragraph not found" },
      { status: 404 },
    );
  }

  await db
    .insert(adjudications)
    .values({
      paragraphId,
      finalCategory,
      finalSpecificity,
      method: "discussion",
      adjudicatorId: session.annotatorId,
      notes: notes ?? null,
    })
    .onConflictDoUpdate({
      target: adjudications.paragraphId,
      set: {
        finalCategory,
        finalSpecificity,
        method: "discussion",
        adjudicatorId: session.annotatorId,
        notes: notes ?? null,
        resolvedAt: new Date(),
      },
    });

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