Customize Mermaid diagrams

import { useCallback, useState } from 'react'
import { TLComponents, Tldraw, TldrawUiButton, useEditor, useValue } from 'tldraw'
import 'tldraw/tldraw.css'
import { FlowchartShapeUtil } from './customMermaidShapeUtil'
import './custom-shape-mermaid.css'
import { getFlowchartSourceError } from './flowchartSourceGuard'
import { mapNodeToRenderSpec } from './mermaidPipelineBlueprint'
import { type StepStatus, pipelineStateAtom, runFullPipeline } from './mermaidPipelineState'
import { applyPipelineStepIndices, extractFlowchartPipelineFromEditor } from './pipelineFromEditor'

const components: TLComponents = {
	TopPanel: TopPanel,
}

const customShapes = [FlowchartShapeUtil]

const DEFAULT_MERMAID = `flowchart LR
  s1[Checkout] --> s2[Build]
  s2 --> s3[Unit tests]
  s2 --> s4[Integration tests]
  s3 --> s5[Deploy]
  s4 --> s5`

export default function MermaidDiagramsCustomShapes() {
	return (
		<div className="tldraw__editor">
			<Tldraw components={components} shapeUtils={customShapes} />
		</div>
	)
}

function TopPanel() {
	const editor = useEditor()
	const [mermaidText, setMermaidText] = useState(DEFAULT_MERMAID)
	const [isApplying, setIsApplying] = useState(false)
	const pipeline = useValue(pipelineStateAtom)

	const canRun =
		!pipeline.parseError && pipeline.nodeIds.length > 0 && !pipeline.isRunning && !isApplying

	const applyWorkflow = useCallback(async () => {
		if (isApplying) return
		setIsApplying(true)

		try {
			const sourceError = getFlowchartSourceError(mermaidText)
			if (sourceError) {
				pipelineStateAtom.set({
					nodeIds: [],
					edges: [],
					statusByNodeId: {},
					parseError: sourceError,
					isRunning: false,
				})
				return
			}

			// [1]
			const { createMermaidDiagram } = await import('@tldraw/mermaid')

			editor.deleteShapes([...editor.getCurrentPageShapeIds()])

			try {
				// [2]
				await createMermaidDiagram(editor, mermaidText, {
					blueprintRender: {
						position: { x: 200, y: 400 },
						centerOnPosition: false,
						mapNodeToRenderSpec,
					},
				})
			} catch {
				pipelineStateAtom.set({
					nodeIds: [],
					edges: [],
					statusByNodeId: {},
					parseError: `An error occurred; please make sure your diagram is valid.`,
					isRunning: false,
				})
				return
			}

			// [3]
			const parsed = extractFlowchartPipelineFromEditor(editor)
			if (!parsed.ok) {
				pipelineStateAtom.set({
					nodeIds: [],
					edges: [],
					statusByNodeId: {},
					parseError: parsed.error,
					isRunning: false,
				})
			} else {
				pipelineStateAtom.set({
					nodeIds: parsed.nodeIds,
					edges: parsed.edges,
					statusByNodeId: Object.fromEntries(
						parsed.nodeIds.map((id) => [id, 'pending' as const])
					) as Record<string, StepStatus>,
					parseError: null,
					isRunning: false,
				})
				applyPipelineStepIndices(editor, parsed.stepIndexByNodeId)
			}

			editor.selectNone()
		} catch {
			pipelineStateAtom.update((s) => ({
				...s,
				parseError: `An error occurred; please make sure your diagram is valid.`,
			}))
		} finally {
			setIsApplying(false)
		}
	}, [editor, isApplying, mermaidText])

	const runPipeline = useCallback(() => {
		void runFullPipeline()
	}, [])

	return (
		<div className="custom-shape-mermaid">
			<div>
				Paste a Mermaid <strong>flowchart</strong> or <strong>graph</strong> (branching is ok; merge
				nodes run after <strong>all</strong> incoming steps pass). Apply runs Mermaid import, then
				builds the DAG from <strong>arrows on the canvas</strong>. Run simulates steps; failures can
				be retried on the shape. Step badges are Kahn layers, not a global sequence. Status is only
				in memory for this demo.
			</div>
			<textarea
				value={mermaidText}
				onChange={(e) => setMermaidText(e.target.value)}
				rows={7}
				spellCheck={false}
				className="custom-shape-mermaid__textarea"
			/>
			{pipeline.parseError && (
				<div className="custom-shape-mermaid__error">{pipeline.parseError}</div>
			)}
			<div className="custom-shape-mermaid__controls">
				<TldrawUiButton type="normal" onClick={applyWorkflow} disabled={isApplying}>
					{isApplying ? 'Applying…' : 'Apply workflow'}
				</TldrawUiButton>
				<TldrawUiButton type="low" onClick={runPipeline} disabled={!canRun}>
					Run pipeline
				</TldrawUiButton>
			</div>
			{pipeline.isRunning && (
				<div className="custom-shape-mermaid__notice">Running simulated steps…</div>
			)}
		</div>
	)
}

/*
[1]
`@tldraw/mermaid` pulls in the (large) mermaid library, so it is imported lazily the first
time the user applies a workflow rather than on page load.

[2]
`blueprintRender.mapNodeToRenderSpec` lets you decide which shape each parsed Mermaid node
becomes. Our mapper (see mermaidPipelineBlueprint.ts) turns every flowchart vertex into a
`flowchart-util` shape and stores the Mermaid node id in its props so we can find it later.

[3]
The pipeline graph is not read from the Mermaid text. Instead it is rebuilt from the arrows
and arrow bindings that the import created on the canvas (see pipelineFromEditor.ts), so if
you reconnect arrows by hand and apply again the pipeline follows the drawing.
*/