next-ai-draw-io/components/chat/use-versions.ts

import type { RefObject } from "react"
import { useEffect, useRef } from "react"
import { toast } from "sonner"
import type { DiagramCommit, LoadMode } from "@/contexts/diagram-context"
import { useDictionary } from "@/hooks/use-dictionary"
import { countCells, diffDiagrams } from "@/lib/diagram-diff"
import { clearHighlights, highlightCells } from "@/lib/drawio/editor-bridge"
import { formatMessage } from "@/lib/i18n/utils"
import { BLANK_MXFILE } from "@/packages/mcp-server/src/pages.ts"
import { useVersionsStore } from "@/stores/versions-store"

interface UseVersionsParams {
    loadDiagram: (
        xml: string,
        skipValidation?: boolean,
        mode?: LoadMode,
    ) => string | null
    getVersionSvg: () => Promise<string | null>
    setCommitHandler: (
        handler: ((commit: DiagramCommit) => void) | null,
    ) => void
    /** Index of the user message whose turn is running */
    currentTurnRef: RefObject<number>
    /** The document on the canvas now */
    chartXMLRef: RefObject<string>
    isBusy: boolean
    /** Live busy state, for the restore toast's Undo that runs later */
    isTurnActive: () => boolean
    /** Changes when another chat comes on screen */
    getChatGeneration: () => number
}

/**
 * Turns every committed AI change into a version (card in the chat), marks
 * the changed shapes on the canvas, and restores or undoes versions.
 */
export function useVersions({
    loadDiagram,
    getVersionSvg,
    setCommitHandler,
    currentTurnRef,
    chartXMLRef,
    isBusy,
    isTurnActive,
    getChatGeneration,
}: UseVersionsParams) {
    const dict = useDictionary()
    // A restore commits too; that path must not add a version
    const restoringRef = useRef(false)

    // Registered once: re-registering on every render would leave a gap
    // (effect cleanups run before the streaming hook commits). The handler
    // only reads refs and stable functions.
    useEffect(() => {
        setCommitHandler(({ beforeXml, afterXml, toolCallId }) => {
            if (restoringRef.current) return
            const { summary, touchedIds } = diffDiagrams(beforeXml, afterXml)
            const before = countCells(beforeXml)
            const fromScratch = before.shapes + before.edges === 0
            const id = useVersionsStore.getState().addVersion({
                xml: afterXml,
                beforeXml,
                turnIndex: currentTurnRef.current,
                toolCallId,
                summary,
                fromScratch,
            })
            // Mark what changed, unless the whole diagram is new
            if (!fromScratch) {
                setTimeout(() => {
                    const marker = getComputedStyle(document.documentElement)
                        .getPropertyValue("--marker")
                        .trim()
                    highlightCells(touchedIds, marker || "#ffd84d")
                }, 60)
            }
            // The picture is of the canvas: take it right away, before a
            // newer version can be on it; once one is, this one gets none
            // rather than the newer one's
            const isNewest = () =>
                useVersionsStore.getState().versions.at(-1)?.id === id
            setTimeout(async () => {
                const svg = isNewest() ? await getVersionSvg() : null
                useVersionsStore
                    .getState()
                    .updateVersion(id, { svg: svg && isNewest() ? svg : "" })
            }, 0)
        })
        return () => setCommitHandler(null)
    }, [])

    const apply = (xml: string) => {
        restoringRef.current = true
        try {
            loadDiagram(xml, true, "commit")
        } finally {
            restoringRef.current = false
        }
        clearHighlights()
    }

    const restoreVersion = (versionId: string) => {
        if (isBusy) return
        const store = useVersionsStore.getState()
        const index = store.versions.findIndex((v) => v.id === versionId)
        if (index < 0) return
        // The canvas before the restore, to go back to it: a multi-page
        // document is loaded in full, which Ctrl+Z cannot undo
        const before = chartXMLRef.current
        const wasOnCanvas = store.onCanvasVersionId
        const wasUndone = store.undoneVersionId
        const generation = getChatGeneration()
        apply(store.versions[index].xml)
        store.setOnCanvas(versionId)
        toast.success(
            formatMessage(dict.versions.restoredVersion, { n: index + 1 }),
            {
                id: "version-restored",
                duration: 8000,
                action: {
                    label: dict.versions.undoRestore,
                    onClick: () => {
                        // Another chat is on screen, or an answer is
                        // changing the canvas: the old canvas is not theirs
                        if (!before || isTurnActive()) return
                        if (getChatGeneration() !== generation) return
                        apply(before)
                        const now = useVersionsStore.getState()
                        if (wasUndone) now.setUndone(wasUndone)
                        else if (wasOnCanvas) now.setOnCanvas(wasOnCanvas)
                    },
                },
            },
        )
    }

    /** Undo the change of the latest version, or redo it when undone */
    const undoVersion = (versionId: string) => {
        if (isBusy) return
        const store = useVersionsStore.getState()
        const version = store.versions.find((v) => v.id === versionId)
        if (!version) return
        if (store.undoneVersionId === versionId) {
            apply(version.xml)
            store.setOnCanvas(versionId)
            return
        }
        if (version.beforeXml === undefined) return
        apply(version.beforeXml || BLANK_MXFILE)
        store.setUndone(versionId)
    }

    return { restoreVersion, undoVersion }
}