frontend/src/lib/kg-selection.ts

import { KnowledgeGraph } from "@/types";

export type KGDisplayMode = "single" | "multiple";

/**
 * Select the best knowledge graph from a list of KGs for single display mode
 */
export function selectBestKnowledgeGraph(
  knowledgeGraphs: KnowledgeGraph[]
): KnowledgeGraph | null {
  if (!knowledgeGraphs || knowledgeGraphs.length === 0) {
    return null;
  }

  if (knowledgeGraphs.length === 1) {
    return knowledgeGraphs[0] || null;
  }

  // Sort by priority criteria
  const sortedKGs = [...knowledgeGraphs].sort((a, b) => {
    // 1. Prefer by status (analyzed > perturbed > enriched > created)
    const statusPriority = {
      analyzed: 4,
      perturbed: 3,
      enriched: 2,
      created: 1,
    };
    const statusDiff =
      (statusPriority[b.status as keyof typeof statusPriority] || 0) -
      (statusPriority[a.status as keyof typeof statusPriority] || 0);
    if (statusDiff !== 0) return statusDiff;

    // 2. Prefer production method over baseline methods
    const aMethod = a.processing_metadata?.method_name || "";
    const bMethod = b.processing_metadata?.method_name || "";
    if (aMethod === "production" && bMethod !== "production") return -1;
    if (bMethod === "production" && aMethod !== "production") return 1;

    // 3. Prefer more recent creation timestamp
    const aTime = a.created_at ? new Date(a.created_at).getTime() : 0;
    const bTime = b.created_at ? new Date(b.created_at).getTime() : 0;
    return bTime - aTime;
  });

  return sortedKGs[0] || null;
}

/**
 * Filter knowledge graphs based on display mode
 */
export function filterKnowledgeGraphsByMode(
  knowledgeGraphs: KnowledgeGraph[],
  mode: KGDisplayMode
): KnowledgeGraph[] {
  if (mode === "multiple") {
    return knowledgeGraphs;
  }

  // Group KGs by trace_id, then select the best one from each group
  const kgsByTrace = new Map<string, KnowledgeGraph[]>();

  knowledgeGraphs.forEach((kg) => {
    // Use a fallback key if trace_id is not available
    const traceKey =
      kg.processing_metadata?.method_name || kg.kg_id || "unknown";
    if (!kgsByTrace.has(traceKey)) {
      kgsByTrace.set(traceKey, []);
    }
    kgsByTrace.get(traceKey)!.push(kg);
  });

  // Select the best KG from each trace group
  const selectedKGs: KnowledgeGraph[] = [];
  kgsByTrace.forEach((kgsForTrace) => {
    const bestKG = selectBestKnowledgeGraph(kgsForTrace);
    if (bestKG !== null) {
      selectedKGs.push(bestKG);
    }
  });

  return selectedKGs;
}

/**
 * Get display mode description for UI
 */
export function getDisplayModeDescription(mode: KGDisplayMode): string {
  switch (mode) {
    case "single":
      return "Showing the most recent/best knowledge graph per trace";
    case "multiple":
      return "Showing all knowledge graphs for research and comparison";
    default:
      return "";
  }
}