next-tjwater-frontend/src/features/workbench/utils/investigation-diagram-style.ts

import { mermaid, type MermaidConfig } from "@streamdown/mermaid";
import type { InvestigationDiagram } from "../api/investigation-report";

export const INVESTIGATION_DIAGRAM_FONT =
  '"Noto Sans SC Variable", "PingFang SC", "Microsoft YaHei", sans-serif';

export const investigationDiagramConfig: MermaidConfig = {
  securityLevel: "strict",
  theme: "base",
  htmlLabels: true,
  fontFamily: INVESTIGATION_DIAGRAM_FONT,
  themeVariables: {
    fontFamily: INVESTIGATION_DIAGRAM_FONT,
    fontSize: "15px",
    background: "#FFFFFF",
    primaryColor: "#FFFFFF",
    primaryTextColor: "#26334D",
    primaryBorderColor: "#A5B5CC",
    secondaryColor: "#FFF6E8",
    secondaryTextColor: "#78531B",
    secondaryBorderColor: "#C3A575",
    tertiaryColor: "#F6F8FC",
    tertiaryBorderColor: "#CCD6E5",
    tertiaryTextColor: "#475976",
    lineColor: "#7689A6",
    edgeLabelBackground: "#FFFFFF",
    clusterBkg: "#F6F8FC",
    clusterBorder: "#CCD6E5"
  },
  flowchart: {
    curve: "stepAfter",
    nodeSpacing: 56,
    rankSpacing: 64,
    padding: 22,
    diagramPadding: 24,
    subGraphTitleMargin: { top: 12, bottom: 48 },
    wrappingWidth: 260
  },
  themeCSS: `
    .node rect { rx: 6px; ry: 6px; stroke-width: 1.4px; }
    .node polygon { stroke-width: 1.6px; }
    .nodeLabel { font-weight: 500; }
    .nodeLabel p { text-wrap: balance; line-height: 1.6; }
    .edgePath path, .flowchart-link { stroke-width: 1.35px; }
    .edgeLabel { color: #536784; font-size: 12px; }
    .edgeLabel p { padding: 3px 7px; border-radius: 4px; line-height: 1.5; }
    .cluster rect { rx: 8px; ry: 8px; stroke-width: 1px; }
    .cluster-label { font-weight: 600; font-size: 13px; }
    .cluster-label .nodeLabel { background: #F6F8FC; padding: 0 4px; }
  `
};

/** Keep group captions beside crossing links, without moving nodes or changing edges. */
export const investigationDiagramPlugins = {
  mermaid: {
    ...mermaid,
    getMermaid(config?: MermaidConfig) {
      const renderer = mermaid.getMermaid(config);
      return {
        initialize: (nextConfig: MermaidConfig) => renderer.initialize(nextConfig),
        async render(id: string, source: string) {
          // Streamdown's two-argument renderer uses Mermaid's temporary #d{id}
          // under body. It must remain measurable, but never enter page flow.
          // Scope the rule to this render, including fullscreen and error paths.
          const measurementStyle = document.createElement("style");
          measurementStyle.textContent = `body > [id=${JSON.stringify(`d${id}`)}] {
            position: fixed !important; inset: 0 auto auto 0;
            width: 100%; visibility: hidden; pointer-events: none;
            contain: layout paint;
          }`;
          document.head.append(measurementStyle);
          const result = await renderer.render(id, source).finally(() => measurementStyle.remove());
          // Mermaid's strict renderer emits HTML labels inside foreignObject,
          // including HTML void tags such as <br>. Parse its sanitized output in
          // an inert HTML fragment, not as a standalone XML document.
          const template = document.createElement("template");
          template.innerHTML = result.svg;
          const svg = template.content.firstElementChild;
          if (template.content.childElementCount !== 1 || svg?.localName !== "svg"
            || svg.namespaceURI !== "http://www.w3.org/2000/svg" || svg.querySelector("parsererror")) {
            throw new Error("Mermaid did not return a valid SVG diagram");
          }
          try {
            svg.querySelectorAll(".cluster").forEach((cluster) => {
              const rect = Array.from(cluster.children).find((child) => child.tagName === "rect");
              const caption = cluster.querySelector(".cluster-label");
              if (!rect || !caption) return;
              const x = Number(rect.getAttribute("x"));
              const y = Number(rect.getAttribute("y"));
              if (!Number.isFinite(x) || !Number.isFinite(y)) return;
              caption.setAttribute(
                "transform",
                `translate(${x + 18}, ${y + (config?.flowchart?.subGraphTitleMargin?.top ?? 12)})`
              );
              // Mermaid paints crossing edges after clusters; keep captions above those edges.
              const clusters = cluster.parentElement;
              if (clusters?.classList.contains("clusters")) clusters.parentElement?.append(caption);
            });
            // Streamdown embeds this as HTML, just like Mermaid's original output.
            return { ...result, svg: svg.outerHTML };
          } catch {
            // Caption placement is cosmetic; never discard a usable original graph.
            return result;
          }
        }
      };
    }
  }
};

const roleStyles = [
  "classDef tjAgent fill:#F6F8FC,stroke:#A5B5CC,color:#26334D,stroke-width:1.4px;",
  "classDef tjHuman fill:#FFF6E8,stroke:#C4A168,color:#735018,stroke-width:1.4px;",
  "classDef tjDecision fill:#F6F8FC,stroke:#A5B5CC,color:#26334D,stroke-width:1.5px;",
  "classDef tjEvidence fill:#EDF7F4,stroke:#6B9E8B,color:#2D5E4E,stroke-width:1.6px;",
  "classDef tjAction fill:#FFF6E8,stroke:#C4A168,color:#735018,stroke-width:1.6px;"
];

/** Style supplied relationships; never create nodes from a dangling display reference. */
export function styleInvestigationDiagram(diagram: InvestigationDiagram, running = false) {
  const assignments = diagram.node_refs.flatMap((node) => {
    const escaped = node.node_id.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
    const inSource = new RegExp(`\\b${escaped}(?=\\s*(?:[\\[({@;]|--|==|-\\.|$))`, "m").test(
      diagram.mermaid
    );
    if (!inSource) return [];
    if (running && node.status === "running") return [`class ${node.node_id} tjRunning;`];
    return node.action_refs.length
      ? [`class ${node.node_id} tjAction;`]
      : node.evidence_refs.length
        ? [`class ${node.node_id} tjEvidence;`]
        : [];
  });
  return [diagram.mermaid, ...roleStyles,
    running ? "classDef tjRunning fill:#E8F0FF,stroke:#4263C7,color:#284574,stroke-width:2.5px;" : "classDef tjRunning fill:#F6F8FC,stroke:#A5B5CC,color:#26334D,stroke-width:1.4px;",
    ...assignments].join("\n");
}