ThothII/frontend/src/viewers/SchemaLinkingViewer.tsx

import React, { useEffect, useState } from "react";
import { renderMermaid } from "./mermaid";

const CAP = 45;

export interface Candidate {
  kind: "table" | "column";
  name: string;
  decision: "promoted" | "excluded" | "pending";
  signals?: Record<string, unknown>;
  evidence?: string[];
}

export interface Join {
  from: string;
  to: string;
  source?: string;
}

export interface SchemaLinking {
  question?: string;
  candidates: Candidate[];
  joins: Join[];
  excluded: { kind: string; name: string }[];
  open_questions: string[];
}

/** The inline "perché" comment for a candidate: evidence first, then signal keys. */
function reasonFor(c: Candidate): string {
  if (c.evidence?.length) return c.evidence[0];
  if (c.signals && Object.keys(c.signals).length)
    return Object.keys(c.signals).join(", ");
  return "";
}

// The real DB type for a promoted column, when the schema-linking data carries it
// (`data_type`/`type`, at top level or under `signals`). Today's schema_linking.json
// has no column types, so this falls back to the generic `col` token that Mermaid
// still needs as the attribute type; a real type shows as soon as the pipeline
// provides one. The token is sanitized to keep the erDiagram valid.
function attrType(c: Candidate): string {
  const sig = (c.signals ?? {}) as Record<string, unknown>;
  const rec = c as unknown as Record<string, unknown>;
  const raw = sig["data_type"] ?? sig["type"] ?? rec["data_type"] ?? rec["type"];
  const t = typeof raw === "string" && raw.trim() ? raw.trim() : "col";
  return t.replace(/[^a-zA-Z0-9]/g, "_");
}

export function buildErDiagram(promoted: Candidate[], joins: Join[]): string {
  const sanitize = (name: string) => name.replace(/[^a-zA-Z0-9]/g, "_");
  const tables = promoted.filter((c) => c.kind === "table");
  const tableNames = new Set(tables.map((t) => t.name));
  const columns = promoted.filter((c) => c.kind === "column");

  const lines: string[] = ["erDiagram"];

  for (const t of tables) {
    const id = sanitize(t.name);
    const cols = columns.filter((col) => col.name.startsWith(t.name + "."));
    lines.push(`  ${id} {`);
    for (const col of cols) {
      lines.push(`    ${attrType(col)} ${sanitize(col.name.slice(t.name.length + 1))}`);
    }
    lines.push(`  }`);
  }

  // A join endpoint may be "table" or "table.column"; resolve to its owning table.
  // Draw a relationship only between two DISTINCT promoted tables, once per pair.
  const owningTable = (ref: string) =>
    ref.includes(".") ? ref.slice(0, ref.indexOf(".")) : ref;
  const seen = new Set<string>();
  for (const j of joins) {
    const a = owningTable(j.from);
    const b = owningTable(j.to);
    if (a === b || !tableNames.has(a) || !tableNames.has(b)) continue;
    const key = [a, b].sort().join("::");
    if (seen.has(key)) continue;
    seen.add(key);
    lines.push(`  ${sanitize(a)} }o--o{ ${sanitize(b)} : join`);
  }

  return lines.join("\n");
}

function buildHierarchicalRows(promoted: Candidate[]) {
  const tables = promoted.filter((c) => c.kind === "table");
  const columns = promoted.filter((c) => c.kind === "column");

  // Group columns under their table (column name format: "table.column")
  const grouped: { table: Candidate; columns: Candidate[] }[] = tables.map(
    (t) => ({
      table: t,
      columns: columns.filter((col) => col.name.startsWith(t.name + ".")),
    })
  );

  // Orphan columns not attached to any promoted table
  const attachedCols = new Set(
    grouped.flatMap((g) => g.columns.map((c) => c.name))
  );
  const orphans = columns.filter((c) => !attachedCols.has(c.name));

  return { grouped, orphans };
}

export function SchemaLinkingViewer({ linking }: { linking: SchemaLinking }) {
  const promoted = linking.candidates.filter((c) => c.decision === "promoted");
  // The graph only contains promoted candidates, so cap on that count.
  const oversized = promoted.length > CAP;
  const [view, setView] = useState<"chart" | "table">(
    oversized ? "table" : "chart"
  );
  const [svg, setSvg] = useState<string>("");

  useEffect(() => {
    if (view !== "chart" || oversized) return;
    let cancelled = false;
    const def = buildErDiagram(promoted, linking.joins);
    renderMermaid(def).then((s) => {
      if (!cancelled) setSvg(s);
    });
    return () => {
      cancelled = true;
    };
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [view, linking, oversized]);

  const { grouped, orphans } = buildHierarchicalRows(promoted);

  return (
    <div className="schema-linking-viewer">
      <div className="flex gap-2 mb-2">
        <button
          onClick={() => setView("chart")}
          disabled={oversized}
          className={`rounded-md border border-border/70 px-2 py-1 text-xs shadow-xs transition-colors hover:bg-accent disabled:cursor-not-allowed disabled:opacity-50 ${view === "chart" ? "font-semibold" : ""}`}
        >
          Graph
        </button>
        <button
          onClick={() => setView("table")}
          className={`rounded-md border border-border/70 px-2 py-1 text-xs shadow-xs transition-colors hover:bg-accent ${view === "table" ? "font-semibold" : ""}`}
        >
          Table
        </button>
      </div>

      {oversized && (
        <p className="text-sm text-amber-600">
          Too many elements for the graph, use the table
        </p>
      )}

      {view === "chart" && !oversized && (
        <div
          className="mermaid-container overflow-auto"
          // eslint-disable-next-line react/no-danger
          dangerouslySetInnerHTML={{ __html: svg }}
        />
      )}

      {view === "table" && (
        <table className="w-full text-sm border-collapse">
          <thead>
            <tr>
              <th className="text-left border-b py-1">Name</th>
              <th className="text-left border-b py-1">Type</th>
              <th className="text-left border-b py-1">Why</th>
            </tr>
          </thead>
          <tbody>
            {grouped.map(({ table, columns: cols }) => (
              <React.Fragment key={table.name}>
                <tr className="font-semibold bg-muted/30">
                  <td className="py-1 pl-1">{table.name}</td>
                  <td>table</td>
                  <td className="text-xs text-muted-foreground">
                    {reasonFor(table)}
                  </td>
                </tr>
                {cols.map((col) => (
                  <tr key={col.name} className="pl-4">
                    <td className="py-0.5 pl-4 text-muted-foreground">
                      {col.name}
                    </td>
                    <td>column</td>
                    <td className="text-xs text-muted-foreground">
                      {reasonFor(col)}
                    </td>
                  </tr>
                ))}
              </React.Fragment>
            ))}
            {orphans.map((col) => (
              <tr key={col.name}>
                <td className="py-0.5 pl-1 text-muted-foreground">
                  {col.name}
                </td>
                <td>column</td>
                <td className="text-xs text-muted-foreground">
                  {reasonFor(col)}
                </td>
              </tr>
            ))}
          </tbody>
        </table>
      )}
    </div>
  );
}