import type { Annotation } from "@sec-cybert/schemas/annotation.ts";
import type { ConsensusResult } from "@sec-cybert/schemas/consensus.ts";
import type { LabelOutput } from "@sec-cybert/schemas/label.ts";

/**
 * Compute consensus from 3 Stage 1 annotations for a single paragraph.
 * Returns the consensus result and whether the judge is needed.
 */
export function computeConsensus(
  paragraphId: string,
  annotations: Annotation[],
): { result: ConsensusResult; needsJudge: boolean } {
  if (annotations.length !== 3) {
    throw new Error(
      `Expected 3 annotations for ${paragraphId}, got ${annotations.length}`,
    );
  }

  const labels = annotations.map((a) => a.label);
  const modelIds = annotations.map((a) => a.provenance.modelId);

  // Category votes
  const catVotes: Record<string, number> = {};
  for (const l of labels) {
    catVotes[l.content_category] = (catVotes[l.content_category] ?? 0) + 1;
  }

  // Specificity votes
  const specVotes: Record<string, number> = {};
  const specValues = labels.map((l) => l.specificity_level);
  for (const s of specValues) {
    const key = String(s);
    specVotes[key] = (specVotes[key] ?? 0) + 1;
  }
  const specSpread =
    Math.max(...specValues) - Math.min(...specValues);

  // Check agreement
  const catMajority = Object.entries(catVotes).find(([, count]) => count >= 2);
  const specMajority = Object.entries(specVotes).find(([, count]) => count >= 2);
  const catAgreed = catMajority !== undefined;
  const specAgreed = specMajority !== undefined;
  const bothAgree = catAgreed && specAgreed;

  // Check if all confidences are low (confidence-aware routing)
  const allLowCatConf = labels.every(
    (l) => l.category_confidence === "low",
  );
  const allLowSpecConf = labels.every(
    (l) => l.specificity_confidence === "low",
  );
  const allLowConfidence = allLowCatConf || allLowSpecConf;

  // Determine method and whether judge is needed
  let method: ConsensusResult["method"];
  let needsJudge: boolean;

  if (!bothAgree) {
    method = "unresolved";
    needsJudge = true;
  } else if (allLowConfidence) {
    // Even when models agree, if all report low confidence, get judge opinion
    method = "unresolved";
    needsJudge = true;
  } else {
    const isUnanimous =
      Object.values(catVotes).some((v) => v === 3) &&
      Object.values(specVotes).some((v) => v === 3);
    method = isUnanimous ? "unanimous" : "majority";
    needsJudge = false;
  }

  // Build final label from majority when available
  let finalLabel: LabelOutput;
  if (catMajority && specMajority) {
    // Find the annotation that matches the majority on both dims, prefer highest confidence
    const majorityCategory = catMajority[0]!;
    const majoritySpec = Number(specMajority[0]!);
    const matching = labels.filter(
      (l) =>
        l.content_category === majorityCategory &&
        l.specificity_level === majoritySpec,
    );
    finalLabel = matching[0] ?? labels.find((l) => l.content_category === majorityCategory) ?? labels[0]!;
  } else if (catMajority) {
    finalLabel = labels.find((l) => l.content_category === catMajority[0]!)!;
  } else {
    // No agreement — pick first label as placeholder (judge will override)
    finalLabel = labels[0]!;
  }

  // Confidence: proportion of agreement
  const catMaxVote = Math.max(...Object.values(catVotes));
  const specMaxVote = Math.max(...Object.values(specVotes));
  const confidence = ((catMaxVote / 3) + (specMaxVote / 3)) / 2;

  return {
    result: {
      paragraphId,
      finalLabel,
      method,
      categoryAgreement: { votes: catVotes, agreed: catAgreed },
      specificityAgreement: {
        votes: specVotes,
        agreed: specAgreed,
        spread: specSpread,
      },
      stage1ModelIds: modelIds,
      stage2JudgeModelId: null,
      confidence,
    },
    needsJudge,
  };
}

/**
 * Update a consensus result after the Stage 2 judge has ruled.
 */
export function applyJudgeResult(
  consensus: ConsensusResult,
  judgeAnnotation: Annotation,
): ConsensusResult {
  return {
    ...consensus,
    finalLabel: judgeAnnotation.label,
    method: "judge-resolved",
    stage2JudgeModelId: judgeAnnotation.provenance.modelId,
    confidence: judgeAnnotation.label.category_confidence === "high" ? 0.9 : 0.7,
  };
}
