next-ai-draw-io/lib/drawio/editor-bridge.ts

/**
 * Direct access to the draw.io editor running in our same-origin iframe.
 *
 * draw.io has no public JavaScript API for embedders; everything here uses
 * its internal objects (EditorUi, mxGraph). Every call checks that what it
 * needs exists, so a draw.io update that renames something turns the feature
 * off (with a console warning) instead of breaking the page.
 *
 * With an external draw.io (cross-origin) none of this is available and the
 * app uses the postMessage protocol only.
 */
import { hasCells } from "@/packages/mcp-server/src/pages.ts"
import { type SelectedCell, useCanvasStore } from "@/stores/canvas-store"

type EditorUi = any
type FrameWindow = any

let ui: EditorUi | null = null
let win: FrameWindow | null = null
let frameElement: HTMLIFrameElement | null = null
let detachListeners: (() => void) | null = null
const warned = new Set<string>()
// Keyboard shortcuts of the app that must also work while focus is inside
// the draw.io iframe (keydown events there never reach our window)
let appShortcutHandler: ((event: KeyboardEvent) => boolean) | null = null

/** The handler returns true when it handled the key */
export function setAppShortcutHandler(
    handler: ((event: KeyboardEvent) => boolean) | null,
) {
    appShortcutHandler = handler
}

function warnOnce(key: string, message: string) {
    if (warned.has(key)) return
    warned.add(key)
    console.warn(`[drawio] ${message}`)
}

export function attachEditor(
    editorUi: EditorUi,
    frameWindow: FrameWindow,
    frame: HTMLIFrameElement,
) {
    if (ui === editorUi) return
    detachEditor()
    ui = editorUi
    win = frameWindow
    frameElement = frame
    detachListeners = installListeners()
    keepPanelsOnResize()
    closeSidePanels()
    softenGrid()
    useCanvasStore.getState().set({ hasEditor: true })
    syncAll()
}

// draw.io's simple UI opens or closes its format panel and shape library
// when its window crosses a width; the app's toolbar opens them, so a
// resize (the chat panel sliding, a narrow window) must leave them alone.
// windowResized only toggles them when it knows the previous width.
function keepPanelsOnResize() {
    const original = ui?.windowResized
    if (typeof original !== "function" || original.naiWrapped) return
    const wrapped = function (this: any, ...args: unknown[]) {
        this.lastWindowWidth = null
        return original.apply(this, args)
    }
    wrapped.naiWrapped = true
    ui.windowResized = wrapped
}

// draw.io's grid (#e6e6e6 every 10 px) is busy behind pale shapes
function softenGrid() {
    const view = graph()?.view
    if (!view) return
    try {
        view.gridColor = "#eceef1"
        view.validateBackground()
    } catch {
        // ignore
    }
}

// draw.io opens its format panel (and sometimes the shape library) on start;
// the app opens them on demand from its own toolbar
function closeSidePanels() {
    try {
        if (isFormatPanelOpen()) ui.actions?.get?.("format")?.funct()
    } catch {
        // ignore
    }
    try {
        if (ui.sidebarWindow?.window?.isVisible?.()) {
            ui.sidebarWindow.window.setVisible(false)
        }
    } catch {
        // ignore
    }
}

export function detachEditor() {
    detachListeners?.()
    detachListeners = null
    clearHighlights()
    ui = null
    win = null
    frameElement = null
    previewBase = null
    useCanvasStore.getState().set({
        hasEditor: false,
        selection: [],
        selectionRect: null,
        canUndo: false,
        canRedo: false,
    })
}

function graph() {
    return ui?.editor?.graph ?? null
}

// ---------------------------------------------------------------------------
// Actions (used by the toolbar)
// ---------------------------------------------------------------------------

export function hasAction(name: string): boolean {
    if (!ui) return false
    const action = ui.actions?.get?.(name)
    if (!action) {
        warnOnce(`action:${name}`, `Action "${name}" not found, hiding it`)
        return false
    }
    return true
}

// Actions that open a dialog or start text editing, which takes the focus
const FOCUS_TAKING_ACTIONS = new Set([
    "findReplace",
    "editDiagram",
    "insertImage",
    "insertLink",
    "print",
    "insertText",
])

export function runAction(name: string) {
    const action = ui?.actions?.get?.(name)
    if (!action) return
    try {
        action.funct()
    } catch (error) {
        console.warn(`[drawio] Action "${name}" failed:`, error)
    }
    // The app's buttons are outside the frame: keys go to draw.io again
    if (!FOCUS_TAKING_ACTIONS.has(name)) focusCanvas()
    // Some actions change state without firing an event we listen to
    setTimeout(syncAll, 0)
}

/**
 * For the app's menus: when their item handed the focus to draw.io (its
 * dialog or the canvas), closing the menu must not take it back
 */
export function keepFocusInCanvas(event: Event) {
    if (frameElement && document.activeElement === frameElement) {
        event.preventDefault()
    }
}

/** Whether the grid and the page view are on (for the menu's check marks) */
export function viewToggles(): { grid: boolean; pageView: boolean } {
    const g = graph()
    return { grid: !!g?.isGridEnabled?.(), pageView: !!g?.pageVisible }
}

export function focusCanvas() {
    try {
        frameElement?.focus()
        graph()?.container?.focus()
    } catch {
        // ignore
    }
}

export function stopFreehand() {
    const freehand = graph()?.freehand
    if (freehand?.isDrawing?.()) freehand.stopDrawing()
    syncAll()
}

// ---------------------------------------------------------------------------
// Replacing the diagram (AI changes)
// ---------------------------------------------------------------------------

// Diagram before the current AI turn started streaming. Previews are applied
// without undo history; the final result replaces this base as one undo step.
let previewBase: string | null = null

function currentFileXml(): string | null {
    try {
        const node = ui.getXmlFileData(null, null, true)
        return win.mxUtils.getXml(node)
    } catch {
        return null
    }
}

function withoutUndo(fn: () => void) {
    const manager = ui?.editor?.undoManager
    if (!manager?.undoableEditHappened) {
        fn()
        return
    }
    const original = manager.undoableEditHappened
    manager.undoableEditHappened = () => {}
    try {
        fn()
    } finally {
        manager.undoableEditHappened = original
    }
}

/**
 * replaceDiagramData replaces the current page with one <mxGraphModel>. Only
 * single-page documents qualify; multi-page ones fall back to a full load.
 */
function toSinglePageModel(xml: string): string | null {
    const doc = new DOMParser().parseFromString(xml, "text/xml")
    if (doc.querySelector("parsererror")) return null
    const root = doc.documentElement
    if (root.nodeName === "mxGraphModel") return xml
    if (root.nodeName !== "mxfile") return null
    const diagrams = root.getElementsByTagName("diagram")
    if (diagrams.length !== 1) return null
    const model = diagrams[0].getElementsByTagName("mxGraphModel")[0]
    if (model) return new XMLSerializer().serializeToString(model)
    // Compressed page
    try {
        const text = diagrams[0].textContent?.trim()
        const inflated = text ? win?.Graph?.decompress?.(text) : null
        return typeof inflated === "string" &&
            inflated.includes("<mxGraphModel")
            ? inflated
            : null
    } catch {
        return null
    }
}

function replace(xml: string) {
    const model = toSinglePageModel(xml)
    if (!model || typeof ui?.replaceDiagramData !== "function") {
        throw new Error("Diagram can't be replaced in place")
    }
    ui.replaceDiagramData(model)
}

function isEmptyModel(): boolean {
    const g = graph()
    if (!g) return true
    const root = g.model.getRoot()
    const layer = root ? g.model.getChildAt(root, 0) : null
    return !layer || g.model.getChildCount(layer) === 0
}

// Page settings draw.io applies on a full load only: replacing the page in
// place keeps the old ones
const LOAD_ONLY_SETTINGS =
    /<mxGraphModel\b[^>]*\s(?:backgroundImage=|extFonts=|math="1"|shadow="1")/

function canvasHasLoadOnlySettings(): boolean {
    const g = graph()
    return (
        !!g &&
        (!!g.backgroundImage ||
            !!g.mathEnabled ||
            !!g.shadowVisible ||
            (g.extFonts?.length ?? 0) > 0)
    )
}

/** Can this diagram go through the editor (with undo) instead of a full load? */
export function canReplaceDiagram(xml: string): boolean {
    if (!ui) return false
    if (typeof ui.replaceDiagramData !== "function") {
        warnOnce("replace", "replaceDiagramData not found, using full loads")
        return false
    }
    // A document with them, or replacing one with them, loads in full
    if (LOAD_ONLY_SETTINGS.test(xml) || canvasHasLoadOnlySettings()) {
        return false
    }
    const pageCount = Array.isArray(ui.pages) ? ui.pages.length : 1
    return pageCount <= 1 && toSinglePageModel(xml) !== null
}

/** Show a streaming preview without touching the undo history */
export function previewDiagram(xml: string) {
    if (previewBase === null) previewBase = currentFileXml() ?? ""
    const wasEmpty = isEmptyModel()
    withoutUndo(() => replace(xml))
    if (wasEmpty) fitDiagram()
}

/** Name of the page in a single-page mxfile, if it has one */
function pageName(xml: string): string | null {
    const doc = new DOMParser().parseFromString(xml, "text/xml")
    if (doc.documentElement?.nodeName !== "mxfile") return null
    return doc.getElementsByTagName("diagram")[0]?.getAttribute("name") || null
}

/**
 * Apply the final AI result (or a restored version) as a single undo step,
 * page name included. The view is fitted only when the canvas was empty, so
 * a restore keeps the user's zoom.
 */
export function commitDiagram(xml: string) {
    const wasEmpty =
        isEmptyModel() || (previewBase !== null && !hasCells(previewBase))
    if (previewBase !== null) {
        const base = previewBase
        previewBase = null
        // Undo then goes back to the diagram before streaming started
        if (base && toSinglePageModel(base)) withoutUndo(() => replace(base))
    }
    const model = graph()?.model
    const name = pageName(xml)
    const page = ui?.currentPage
    model?.beginUpdate()
    try {
        replace(xml)
        if (name && page && win?.RenamePage && page.getName?.() !== name) {
            model.execute(new win.RenamePage(ui, page, name))
        }
    } finally {
        model?.endUpdate()
    }
    if (wasEmpty) fitDiagram()
}

/** Throw away the preview and go back to the given diagram */
export function revertPreview(targetXml: string) {
    previewBase = null
    withoutUndo(() => replace(targetXml))
}

/** Forget the preview base (new user turn) without changing the canvas */
export function resetPreview() {
    previewBase = null
}

// When the app last fitted the diagram on its own; a canvas resize right
// after (the chat panel sliding in) fits again, unless the user zoomed since
let lastAutoFitAt = 0

// True while fitDiagram runs: its own zoom change is not the user's
let fitting = false

/** The zoom changed: if the user did it, a later resize keeps their view */
function zoomChanged() {
    if (!fitting) lastAutoFitAt = 0
}

/** Fit again if the last automatic fit just happened (layout still moving) */
export function refitIfRecent(windowMs = 1500) {
    if (Date.now() - lastAutoFitAt < windowMs) {
        fitDiagram()
        return true
    }
    return false
}

// Room around a fitted diagram: the app's toolbar sits over the top of the
// canvas and the page tabs over the bottom (half of this on each side)
const FIT_BORDER = 128

/**
 * Fit the shapes (not the whole page) into the canvas, at most at 100%,
 * clear of the toolbar and the page tabs
 */
export function fitDiagram() {
    const g = graph()
    if (!g) return
    fitting = true
    try {
        const b = g.getGraphBounds()
        const { scale, translate } = g.view
        if (b && b.width > 0 && b.height > 0 && win?.mxRectangle) {
            const model = new win.mxRectangle(
                b.x / scale - translate.x,
                b.y / scale - translate.y,
                b.width / scale,
                b.height / scale,
            )
            g.fitWindow(model, FIT_BORDER, 1)
        } else if (ui?.actions?.get?.("fitWindow")) {
            ui.actions.get("fitWindow").funct()
            if (g.view.scale > 1) g.zoomTo(1, true)
        }
    } catch {
        // ignore
    } finally {
        fitting = false
    }
    // After the zoom events of the fit itself
    lastAutoFitAt = Date.now()
    syncZoom()
}

/** The canvas changed width: keep what was in the middle in the middle */
export function keepCenter(oldWidth: number, newWidth: number) {
    const container = graph()?.container
    if (!container || !oldWidth || !newWidth) return
    container.scrollLeft += (oldWidth - newWidth) / 2
}

// ---------------------------------------------------------------------------
// Highlighting what the AI changed
// ---------------------------------------------------------------------------

let highlights: any[] = []

export function clearHighlights() {
    for (const h of highlights) {
        try {
            h.destroy()
        } catch {
            // ignore
        }
    }
    highlights = []
}

// Outline of changed shapes: dark enough to show on white paper and on
// draw.io's pale fills (yellow ones included)
const HIGHLIGHT_OUTLINE = "#a86b00"

/**
 * Marks changed cells: shapes get a soft halo plus a thin outline,
 * connectors only the outline (a halo would cover their labels and arrows)
 */
export function highlightCells(ids: string[], color: string) {
    clearHighlights()
    const g = graph()
    if (!g || !win?.mxCellHighlight || ids.length === 0) return
    const add = (
        state: any,
        stroke: string,
        width: number,
        opacity: number,
    ) => {
        const h = new win.mxCellHighlight(g, stroke, width)
        h.opacity = opacity
        h.highlight(state)
        highlights.push(h)
    }
    for (const id of ids.slice(0, 200)) {
        const cell = g.model.getCell(id)
        const state = cell ? g.view.getState(cell) : null
        if (!state) continue
        if (!g.model.isEdge(cell)) add(state, color, 8, 18)
        add(state, HIGHLIGHT_OUTLINE, 2, 100)
    }
}

// ---------------------------------------------------------------------------
// Selection
// ---------------------------------------------------------------------------

function cellLabel(cell: any): string {
    const g = graph()
    let text = ""
    try {
        text = g?.convertValueToString(cell) ?? ""
    } catch {
        text = ""
    }
    // Labels can contain HTML
    text = text
        .replace(/<br\s*\/?>/gi, " ")
        .replace(/<[^>]+>/g, "")
        .replace(/&nbsp;/g, " ")
        .replace(/&amp;/g, "&")
        .replace(/&lt;/g, "<")
        .replace(/&gt;/g, ">")
        .replace(/\s+/g, " ")
        .trim()
    return text.length > 80 ? `${text.slice(0, 80)}…` : text
}

function readSelection(): SelectedCell[] {
    const g = graph()
    if (!g) return []
    return (g.getSelectionCells() as any[])
        .filter((cell) => cell?.id)
        .map((cell) => ({
            id: String(cell.id),
            label: cellLabel(cell),
            isEdge: !!g.model.isEdge(cell),
        }))
}

function readSelectionRect() {
    const g = graph()
    if (!g || g.isSelectionEmpty()) return null
    const cells = g.getSelectionCells()
    const bounds = g.view.getBounds(cells)
    if (!bounds) return null
    const container = g.container as HTMLElement
    const rect = container.getBoundingClientRect()
    return {
        x: rect.left + bounds.x - container.scrollLeft,
        y: rect.top + bounds.y - container.scrollTop,
        width: bounds.width,
        height: bounds.height,
    }
}

// ---------------------------------------------------------------------------
// Pages
// ---------------------------------------------------------------------------

function readPages() {
    const pages = (ui?.pages ?? []) as any[]
    return pages.map((page, index) => ({
        id: String(page.getId?.() ?? index),
        name: String(page.getName?.() ?? `Page-${index + 1}`),
    }))
}

function findPage(id: string) {
    return ((ui?.pages ?? []) as any[]).find(
        (page) => String(page.getId?.()) === id,
    )
}

export function selectPage(id: string) {
    const page = findPage(id)
    if (page && typeof ui.selectPage === "function") ui.selectPage(page)
    focusCanvas()
    syncAll()
}

export function addPage() {
    runAction("insertPage")
}

export function renamePage(id: string) {
    const page = findPage(id)
    if (page && typeof ui.renamePage === "function") ui.renamePage(page)
}

export function duplicatePage(id: string) {
    const page = findPage(id)
    if (page && typeof ui.duplicatePage === "function") {
        // Named as draw.io's own menu names it ("Copy of …")
        const name = win?.mxResources?.get?.("copyOf", [page.getName()])
        ui.duplicatePage(page, name)
    }
    focusCanvas()
    syncAll()
}

export function removePage(id: string) {
    const page = findPage(id)
    if (page && typeof ui.removePage === "function") ui.removePage(page)
    syncAll()
}

export function canManagePages(): boolean {
    return (
        !!ui && Array.isArray(ui.pages) && typeof ui.selectPage === "function"
    )
}

// ---------------------------------------------------------------------------
// Keeping the store in sync
// ---------------------------------------------------------------------------

function syncZoom() {
    const g = graph()
    if (!g) return
    useCanvasStore.getState().set({ zoom: Math.round(g.view.scale * 100) })
}

function syncSelectionRect() {
    useCanvasStore.getState().set({ selectionRect: readSelectionRect() })
}

function syncAll() {
    if (!ui) return
    const g = graph()
    const manager = ui.editor?.undoManager
    useCanvasStore.getState().set({
        canUndo: !!manager?.canUndo?.(),
        canRedo: !!manager?.canRedo?.(),
        zoom: g ? Math.round(g.view.scale * 100) : 100,
        pages: readPages(),
        currentPageId: ui.currentPage?.getId
            ? String(ui.currentPage.getId())
            : null,
        selection: readSelection(),
        selectionRect: readSelectionRect(),
        isFreehand: !!g?.freehand?.isDrawing?.(),
        isFormatOpen: isFormatPanelOpen(),
        isShapesOpen: isShapesPanelOpen(),
    })
}

// draw.io keeps its docked shape library "visible" from the start; the app's
// CSS hides it unless the frame's body has this class
const SHAPES_OPEN_CLASS = "nai-shapes-open"

function isShapesPanelOpen(): boolean {
    return !!win?.document?.body?.classList.contains(SHAPES_OPEN_CLASS)
}

/** Show or hide draw.io's shape library, floating over the canvas */
export function toggleShapesPanel() {
    const body = win?.document?.body
    if (!body) return
    // Only the class: draw.io treats its docked library as open from the
    // start, and lays the canvas out next to it while it shows
    body.classList.toggle(SHAPES_OPEN_CLASS, !isShapesPanelOpen())
    syncAll()
}

function isFormatPanelOpen(): boolean {
    try {
        if (typeof ui.isFormatPanelVisible === "function") {
            return !!ui.isFormatPanelVisible()
        }
        return (ui.formatWidth ?? 0) > 0
    } catch {
        return false
    }
}

function installListeners(): () => void {
    const g = graph()
    if (!g || !win?.mxEvent) return () => {}
    const mxEvent = win.mxEvent
    const cleanups: (() => void)[] = []

    const listen = (source: any, name: string, handler: () => void) => {
        if (!source?.addListener) return
        source.addListener(name, handler)
        cleanups.push(() => source.removeListener(handler))
    }

    // Any user action ends the "what the AI just changed" highlight
    const onSelection = () => {
        clearHighlights()
        syncAll()
    }
    // Drawing stopped (Escape, X or the Select button): the toolbar shows
    // it, and draw.io's freehand settings window goes away
    const onFreehand = () => {
        if (!g.freehand?.isDrawing?.()) {
            ui?.freehandWindow?.window?.setVisible?.(false)
        }
        syncAll()
    }
    const onView = () => {
        syncZoom()
        syncSelectionRect()
    }
    const onScale = () => {
        zoomChanged()
        onView()
    }

    listen(g.getSelectionModel(), mxEvent.CHANGE, onSelection)
    listen(g.model, mxEvent.CHANGE, syncAll)
    listen(ui.editor?.undoManager, mxEvent.ADD, syncAll)
    listen(ui.editor?.undoManager, mxEvent.UNDO, syncAll)
    listen(ui.editor?.undoManager, mxEvent.REDO, syncAll)
    listen(ui.editor?.undoManager, mxEvent.CLEAR, syncAll)
    listen(g, "freehandStateChanged", onFreehand)
    listen(g.view, mxEvent.SCALE, onScale)
    listen(g.view, mxEvent.TRANSLATE, onView)
    listen(g.view, mxEvent.SCALE_AND_TRANSLATE, onScale)
    listen(ui.editor, "pageSelected", syncAll)
    listen(ui, "formatWidthChanged", syncAll)

    // Capture phase, so the app shortcut wins before draw.io's key handler
    const onKeyDown = (event: KeyboardEvent) => {
        if (appShortcutHandler?.(event)) {
            event.preventDefault()
            event.stopPropagation()
        }
    }
    const frameDocument = win.document as Document | undefined
    frameDocument?.addEventListener("keydown", onKeyDown, true)
    cleanups.push(() =>
        frameDocument?.removeEventListener("keydown", onKeyDown, true),
    )

    const container = g.container as HTMLElement | undefined
    container?.addEventListener("scroll", syncSelectionRect, { passive: true })
    cleanups.push(() =>
        container?.removeEventListener("scroll", syncSelectionRect),
    )

    return () => {
        for (const cleanup of cleanups) cleanup()
    }
}

/**
 * Finds the EditorUi instance inside a same-origin draw.io iframe.
 *
 * draw.io keeps no global reference to it, so we wrap prototype methods it
 * calls on startup and after every load; the first call hands us `this`.
 * Returns false when the frame is cross-origin.
 */
export function watchForEditorUi(
    frameWindow: Window,
    onFound: (editorUi: EditorUi) => void,
): boolean {
    let w: any
    try {
        w = frameWindow as any
        // Throws for a cross-origin frame
        void w.document
    } catch {
        return false
    }
    const hook = (proto: any, method: string) => {
        if (!proto || typeof proto[method] !== "function") return
        if (proto[method].__naiWrapped) return
        const original = proto[method]
        const wrapped = function (this: any, ...args: any[]) {
            onFound(this)
            return original.apply(this, args)
        }
        ;(wrapped as any).__naiWrapped = true
        proto[method] = wrapped
    }
    hook(w.EditorUi?.prototype, "updateActionStates")
    hook(w.App?.prototype, "fileLoaded")
    hook(w.EditorUi?.prototype, "fileLoaded")
    return true
}