zui/src/components/nodes/FunctionNode.tsx

import React, { memo, useCallback, useContext, useMemo, useRef } from 'react'
import CodeMirror from '@uiw/react-codemirror'
import { javascript } from '@codemirror/lang-javascript'
import FlowContext from '../../lib/flowContext'
import { nodePropsAreEqual } from '../../lib/flowUtils'
import { useResizeHeight } from '../../hooks/useResizeHeight'
import { useTheme } from '../../lib/themeContext'
import {
    BaseNode,
    BaseNodeContent,
    BaseNodeFooter,
    BaseNodeHeaderRow,
} from '../base/BaseNode'
import { InputHandle, OutputHandle } from '../base/NodeHandles'
import { NodeFooterEdgeIndicators } from '../base/NodeFooterEdgeIndicators'
import { NodeHeaderTitle } from '../base/NodeHeaderTitle'
import { NodeMenubar } from '../base/NodeMenubar'
import { MenubarItem, MenubarShortcut } from '../ui/menubar'
import { Kbd } from '../ui/kbd'
import { Code2, Variable } from 'lucide-react'

type Props = {
    id: string
    data: { body?: string }
    width?: number
    height?: number
}

export const FunctionNode = memo(function FunctionNode({ id, data, width, height }: Props) {
    const bodyValue = data?.body ?? ''
    const { theme } = useTheme()
    const ctx = useContext(FlowContext)
    const setNodes = ctx?.setNodes
    const editorRef = useRef<unknown>(null)

    const edges = ctx?.edges ?? []
    const nodes = ctx?.nodes ?? []
    const incomingEdges = useMemo(() => edges.filter((e: any) => e.target === id), [edges, id])
    const incomingIds = useMemo(() => incomingEdges.map((e: any) => e.source).sort(), [incomingEdges])
    const connectedVariableNodes = useMemo(
        () => (nodes as any[]).filter((n: any) => incomingIds.includes(n.id) && n.type === 'variable'),
        [nodes, incomingIds]
    )
    const connectedFunctionNodes = useMemo(
        () => (nodes as any[]).filter((n: any) => incomingIds.includes(n.id) && n.type === 'function'),
        [nodes, incomingIds]
    )
    const hasConnectedVariables = connectedVariableNodes.length > 0
    const hasConnectedFunctions = connectedFunctionNodes.length > 0
    const hasConnectedInputs = hasConnectedVariables || hasConnectedFunctions

    const onChange = useCallback(
        (val: string) => {
            if (setNodes) {
                setNodes((nds: any[]) =>
                    nds.map((n) => (n.id === id ? { ...n, data: { ...n.data, body: val } } : n))
                )
            }
        },
        [id, setNodes]
    )

    const insertAt = useCallback(
        (insertText: string, mode: 'prepend' | 'append' | 'cursor') => {
            const ref = editorRef.current as { view: { state: { doc: { length: number }; selection: { main: { from: number } } }; dispatch: (arg: { changes: { from: number; to: number; insert: string } }) => void } } | null
            if (ref?.view) {
                const view = ref.view
                const doc = view.state.doc
                const len = doc.length
                let from: number
                if (mode === 'prepend') from = 0
                else if (mode === 'append') from = len
                else from = view.state.selection.main.from
                view.dispatch({ changes: { from, to: from, insert: insertText } })
                onChange(view.state.doc.toString())
                return
            }
            if (mode === 'prepend') onChange(insertText + bodyValue)
            else onChange(bodyValue + insertText)
        },
        [onChange, bodyValue]
    )

    const insertVariableAtCursor = useCallback(
        (variableNode: any) => {
            insertAt(variableNode.id, 'cursor')
        },
        [insertAt]
    )

    const insertFunctionAtCursor = useCallback(
        (functionNode: any) => {
            insertAt(functionNode.id, 'cursor')
        },
        [insertAt]
    )

    const extensions = useMemo(() => [javascript()], [])
    const [editorHeight, editorContainerRef] = useResizeHeight(120)
    const dimensions =
        width != null && height != null && width > 0 && height > 0
            ? { width, height }
            : undefined

    return (
        <BaseNode className="min-w-72 min-h-[260px]" dimensions={dimensions} resizable nodeId={id} handles={<><InputHandle id="in" nodeId={id} /><OutputHandle id="out" /></>}>
            <BaseNodeHeaderRow icon={<Code2 className="size-4" />} title={<NodeHeaderTitle nodeId={id} displayTitle={id} />} />

            <BaseNodeContent>
                <div className="shrink-0 w-full">
                    <NodeMenubar
                        nodeId={id}
                        nodeType="function"
                        inputsMenuContent={
                            hasConnectedInputs ? (
                                <>
                                    {connectedVariableNodes.map((n: any) => (
                                        <MenubarItem
                                            key={n.id}
                                            className="text-xs flex items-center gap-2 group"
                                            onClick={() => insertVariableAtCursor(n)}
                                        >
                                            <Variable className="size-3.5 shrink-0" />
                                            {n.id}
                                            <MenubarShortcut className="opacity-0 group-hover:opacity-100 transition-opacity"><Kbd>Insert</Kbd></MenubarShortcut>
                                        </MenubarItem>
                                    ))}
                                    {connectedFunctionNodes.map((n: any) => (
                                        <MenubarItem
                                            key={n.id}
                                            className="text-xs flex items-center gap-2 group"
                                            onClick={() => insertFunctionAtCursor(n)}
                                        >
                                            <Code2 className="size-3.5 shrink-0" />
                                            {n.id}
                                            <MenubarShortcut className="opacity-0 group-hover:opacity-100 transition-opacity"><Kbd>Insert</Kbd></MenubarShortcut>
                                        </MenubarItem>
                                    ))}
                                </>
                            ) : (
                                <span className="text-xs text-muted-foreground px-2 py-1">Connect nodes to insert at cursor</span>
                            )
                        }
                    />
                </div>
                <div ref={editorContainerRef} className="min-h-0 flex-1 w-full nodrag nopan overflow-hidden border-t border-input">
                    <CodeMirror
                        // @ts-expect-error ref is { view, state, editor }; package ref type not in our node_modules
                        ref={editorRef}
                        value={bodyValue}
                        height={`${editorHeight}px`}
                        theme={theme}
                        extensions={extensions}
                        onChange={onChange}
                        basicSetup={{ lineNumbers: true, foldGutter: false }}
                        className="text-xs [&_.cm-editor]:outline-none [&_.cm-editor]:cursor-text [&_.cm-gutters]:border-0"
                    />
                </div>
            </BaseNodeContent>

            <BaseNodeFooter>
                <NodeFooterEdgeIndicators nodeId={id} nodeType="function">
                    {bodyValue ? `JavaScript · ${bodyValue.length} chars` : 'JavaScript · none'}
                </NodeFooterEdgeIndicators>
            </BaseNodeFooter>
        </BaseNode>
    )
}, nodePropsAreEqual)

FunctionNode.displayName = 'FunctionNode'

export default FunctionNode