diagrams-drawio/scripts/src/services/themes/theme-service.ts

import type { DiagramEdge, DiagramIRV2, DiagramNode } from "../../model/diagram-ir.js";

export type BuiltInThemeName = "default" | "corporate" | "dark" | "colorblind-safe" | "handdrawn";

export interface ThemeRoleStyle {
  fillColor: string;
  strokeColor: string;
  fontColor: string;
}

export interface ThemeDefinition {
  name: BuiltInThemeName | (string & {});
  backgroundColor: string;
  node: ThemeRoleStyle;
  container: ThemeRoleStyle;
  edge: Pick<ThemeRoleStyle, "strokeColor" | "fontColor">;
  sketch?: boolean;
}

function deepFreeze<T>(value: T): T {
  if (value && typeof value === "object" && !Object.isFrozen(value)) {
    for (const nested of Object.values(value as Record<string, unknown>)) deepFreeze(nested);
    Object.freeze(value);
  }
  return value;
}

export const BUILT_IN_THEMES: Readonly<Record<BuiltInThemeName, ThemeDefinition>> = deepFreeze({
  default: theme("default", "#FFFFFF", ["#FFFFFF", "#4A5568", "#1A202C"], ["#F7FAFC", "#718096", "#1A202C"], ["#4A5568", "#1A202C"]),
  corporate: theme("corporate", "#FFFFFF", ["#E8F0FE", "#1F5AA6", "#102A43"], ["#F5F8FC", "#526D82", "#102A43"], ["#1F5AA6", "#102A43"]),
  dark: theme("dark", "#111827", ["#243447", "#94A3B8", "#FFFFFF"], ["#1F2937", "#64748B", "#FFFFFF"], ["#CBD5E1", "#FFFFFF"]),
  "colorblind-safe": theme("colorblind-safe", "#FFFFFF", ["#FFF4CC", "#0072B2", "#1A1A1A"], ["#E8F3F8", "#009E73", "#1A1A1A"], ["#0072B2", "#1A1A1A"]),
  handdrawn: { ...theme("handdrawn", "#FFFDF5", ["#FFFDF5", "#444444", "#222222"], ["#F6F1E5", "#555555", "#222222"], ["#444444", "#222222"]), sketch: true },
});

function theme(name: BuiltInThemeName, backgroundColor: string, node: [string, string, string], container: [string, string, string], edge: [string, string]): ThemeDefinition {
  return {
    name,
    backgroundColor,
    node: { fillColor: node[0], strokeColor: node[1], fontColor: node[2] },
    container: { fillColor: container[0], strokeColor: container[1], fontColor: container[2] },
    edge: { strokeColor: edge[0], fontColor: edge[1] },
  };
}

const HEX_COLOR = /^#[0-9A-Fa-f]{6}$/;

export function contrastRatio(left: string, right: string): number {
  const luminance = (color: string): number => {
    if (!HEX_COLOR.test(color)) throw new Error(`Invalid color: ${color}`);
    const values = [1, 3, 5].map((offset) => Number.parseInt(color.slice(offset, offset + 2), 16) / 255)
      .map((value) => value <= 0.04045 ? value / 12.92 : ((value + 0.055) / 1.055) ** 2.4);
    return values[0] * 0.2126 + values[1] * 0.7152 + values[2] * 0.0722;
  };
  const [high, low] = [luminance(left), luminance(right)].sort((a, b) => b - a);
  return (high + 0.05) / (low + 0.05);
}

export function validateTheme(themeDefinition: ThemeDefinition): ThemeDefinition {
  if (!themeDefinition || typeof themeDefinition.name !== "string" || themeDefinition.name.trim() === "") throw new Error("Theme requires a name");
  const colors: Array<[string, string]> = [
    ["backgroundColor", themeDefinition.backgroundColor],
    ["node.fillColor", themeDefinition.node?.fillColor], ["node.strokeColor", themeDefinition.node?.strokeColor], ["node.fontColor", themeDefinition.node?.fontColor],
    ["container.fillColor", themeDefinition.container?.fillColor], ["container.strokeColor", themeDefinition.container?.strokeColor], ["container.fontColor", themeDefinition.container?.fontColor],
    ["edge.strokeColor", themeDefinition.edge?.strokeColor], ["edge.fontColor", themeDefinition.edge?.fontColor],
  ];
  for (const [path, color] of colors) if (typeof color !== "string" || !HEX_COLOR.test(color)) throw new Error(`Theme ${path} must be a #RRGGBB color`);
  for (const role of ["node", "container"] as const) {
    if (contrastRatio(themeDefinition[role].fontColor, themeDefinition[role].fillColor) < 4.5) throw new Error(`Theme ${role} text has insufficient contrast`);
  }
  if (contrastRatio(themeDefinition.edge.fontColor, themeDefinition.backgroundColor) < 4.5) throw new Error("Theme edge text has insufficient contrast");
  return themeDefinition;
}

interface StyleEntry { key: string; value?: string }

function parseStyle(style = ""): StyleEntry[] {
  return style.split(";").filter(Boolean).map((token) => {
    const separator = token.indexOf("=");
    return separator < 0 ? { key: token } : { key: token.slice(0, separator), value: token.slice(separator + 1) };
  });
}

function themedStyle(style: string | undefined, colors: Partial<ThemeRoleStyle>, sketch: boolean, backgroundColor?: string, labelBackgroundColor?: string): string {
  const entries = parseStyle(style);
  const originalFillNone = entries.some(({ key, value }) => key === "fillColor" && value?.toLowerCase() === "none");
  const visual = new Map<string, string>();
  if (colors.fillColor) visual.set("fillColor", originalFillNone ? "none" : colors.fillColor);
  // Edge labels are drawn on the canvas, so their backdrop must follow the canvas colour or they render as boxes.
  if (labelBackgroundColor) visual.set("labelBackgroundColor", labelBackgroundColor);
  if (colors.strokeColor) visual.set("strokeColor", colors.strokeColor);
  if (colors.fontColor) {
    let fontColor = colors.fontColor;
    if (originalFillNone && backgroundColor && contrastRatio(fontColor, backgroundColor) < 4.5) {
      fontColor = contrastRatio("#000000", backgroundColor) >= 4.5 ? "#000000" : "#FFFFFF";
    }
    visual.set("fontColor", fontColor);
  }
  visual.set("sketch", sketch ? "1" : "0");
  const visualKeys = new Set(["fillColor", "strokeColor", "fontColor", "labelBackgroundColor", "sketch"]);
  const handled = new Set<string>();
  const output = entries.flatMap((entry) => {
    if (!visualKeys.has(entry.key)) return [entry];
    if (handled.has(entry.key)) return [];
    handled.add(entry.key);
    const replacement = visual.get(entry.key);
    if (replacement === undefined) return [entry];
    visual.delete(entry.key);
    return [{ key: entry.key, value: replacement }];
  });
  for (const key of ["fillColor", "strokeColor", "fontColor", "labelBackgroundColor", "sketch"]) {
    const value = visual.get(key);
    if (value !== undefined) output.push({ key, value });
  }
  return output.map(({ key, value }) => value === undefined ? key : `${key}=${value}`).join(";") + ";";
}

export function applyTheme(ir: DiagramIRV2, selected: BuiltInThemeName | ThemeDefinition): DiagramIRV2 {
  const selectedTheme = typeof selected === "string" ? BUILT_IN_THEMES[selected] : selected;
  if (!selectedTheme) throw new Error(`Unknown built-in theme: ${selected}`);
  validateTheme(selectedTheme);
  const isolated = structuredClone(ir);
  const applyNode = (node: DiagramNode): DiagramNode => ({
    ...node,
    style: themedStyle(node.style, node.kind === "container" ? selectedTheme.container : selectedTheme.node, selectedTheme.sketch === true, selectedTheme.backgroundColor),
  });
  const applyEdge = (edge: DiagramEdge): DiagramEdge => ({ ...edge, style: themedStyle(edge.style, selectedTheme.edge, selectedTheme.sketch === true, undefined, selectedTheme.backgroundColor) });
  return {
    ...isolated,
    theme: selectedTheme.name,
    pages: isolated.pages.map((page) => {
      const extensions = page.extensions ?? {};
      const drawio = extensions.$drawio && typeof extensions.$drawio === "object" && !Array.isArray(extensions.$drawio)
        ? extensions.$drawio as Record<string, unknown>
        : {};
      const attributes = drawio.attributes && typeof drawio.attributes === "object" && !Array.isArray(drawio.attributes)
        ? drawio.attributes as Record<string, unknown>
        : {};
      return {
        ...page,
        extensions: {
          ...extensions,
          $drawio: { ...drawio, attributes: { ...attributes, "model:background": selectedTheme.backgroundColor } },
        },
        nodes: page.nodes.map(applyNode),
        edges: page.edges.map(applyEdge),
      };
    }),
  };
}