next-ai-draw-io/components/chat/chat-engine.tsx

"use client"

import { useChat } from "@ai-sdk/react"
import { DefaultChatTransport, isToolUIPart, type UIMessage } from "ai"
import { usePathname, useRouter, useSearchParams } from "next/navigation"
import type React from "react"
import {
    createContext,
    useCallback,
    useContext,
    useEffect,
    useLayoutEffect,
    useMemo,
    useRef,
    useState,
} from "react"
import { flushSync } from "react-dom"
import { toast } from "sonner"
import type { ValidationState } from "@/components/chat/ValidationCard"
import { useDiagram } from "@/contexts/diagram-context"
import { useDiagramToolHandlers } from "@/hooks/use-diagram-tool-handlers"
import { useDictionary } from "@/hooks/use-dictionary"
import {
    getSelectedAIConfig,
    type UseModelConfigReturn,
    useModelConfig,
} from "@/hooks/use-model-config"
import {
    type SessionData,
    useSessionManager,
} from "@/hooks/use-session-manager"
import { useValidateDiagram } from "@/hooks/use-validate-diagram"
import { getApiEndpoint } from "@/lib/base-path"
import { findCachedResponse } from "@/lib/cached-responses"
import { buildChatHeaders } from "@/lib/chat-request"
import { inflatePages } from "@/lib/diagram-diff"
import { formatMessage } from "@/lib/i18n/utils"
import { isPdfFile, isTextFile } from "@/lib/pdf-utils"
import type { SessionMetadata } from "@/lib/session-storage"
import { sanitizeMessages } from "@/lib/session-storage"
import { STORAGE_KEYS } from "@/lib/storage"
import type { UrlData } from "@/lib/url-utils"
import { type FileData, useFileProcessor } from "@/lib/use-file-processor"
import { useQuotaManager } from "@/lib/use-quota-manager"
import { formatXML, isRealDiagram } from "@/lib/utils"
import { prepareNewDiagram } from "@/packages/mcp-server/src/new-diagram.ts"
import { BLANK_MXFILE, hasCells } from "@/packages/mcp-server/src/pages.ts"
import { type SelectedCell, useCanvasStore } from "@/stores/canvas-store"
import { useSettingsStore } from "@/stores/settings-store"
import { useUiStore } from "@/stores/ui-store"
import {
    useVersionsStore,
    versionsFromLegacyHistory,
} from "@/stores/versions-store"
import {
    APPENDED_FILE_SECTIONS_PATTERN,
    getUserOriginalText,
} from "./message-text"
import { useDiagramStreaming } from "./use-diagram-streaming"
import { useVersions } from "./use-versions"

// sessionStorage key for the unsent input (survives remounts and reloads)
const SESSION_STORAGE_INPUT_KEY = "next-ai-draw-io-input"

const TOOL_ERROR_STATE = "output-error" as const
// 3 to support VLM validation retries (matches MAX_VALIDATION_RETRIES)
const MAX_AUTO_RETRY_COUNT = 3
// Limit for truncation continuation retries
const MAX_CONTINUATION_RETRY_COUNT = 2
// Save the session at most once per second
const SAVE_DEBOUNCE_MS = 1000

export type ChatStatus = "submitted" | "streaming" | "ready" | "error"

/**
 * Check if auto-resubmit should happen based on tool errors.
 * Only checks the LAST tool part (most recent tool call), not all tool parts.
 */
function hasToolErrors(messages: UIMessage[]): boolean {
    const lastMessage = messages[messages.length - 1]
    if (lastMessage?.role !== "assistant") return false
    const lastToolPart = lastMessage.parts.filter(isToolUIPart).at(-1)
    return lastToolPart?.state === TOOL_ERROR_STATE
}

/**
 * Snapshots keep the full multi-page document, but the model only sees and
 * edits the first page, so give it the first page's mxGraphModel.
 * Older snapshots already hold a single mxGraphModel and are returned as is.
 */
function getFirstPageXml(xml: string): string {
    if (!xml.includes("<mxfile")) return xml
    const doc = new DOMParser().parseFromString(xml, "text/xml")
    const model = doc.querySelector("diagram")?.querySelector("mxGraphModel")
    return model ? formatXML(new XMLSerializer().serializeToString(model)) : xml
}

// Shapes sent with a user message (also kept in its metadata)
type MessageSelection = Pick<SelectedCell, "id" | "label">[]

function selectionOf(message: UIMessage): MessageSelection {
    const meta = message.metadata as
        | { selectedCells?: MessageSelection }
        | undefined
    return Array.isArray(meta?.selectedCells) ? meta.selectedCells : []
}

function newLangfuseSessionId() {
    return `session-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
}

/** Read the diagram document out of a .drawio, .xml or .drawio.svg file */
async function readDiagramFile(file: File): Promise<string | null> {
    const text = await file.text()
    if (text.includes("<mxfile") || text.includes("<mxGraphModel")) {
        if (!text.trimStart().startsWith("<svg") && !text.includes("<svg")) {
            return text
        }
    }
    // Editable SVG: the diagram is stored in the root's `content` attribute
    const doc = new DOMParser().parseFromString(text, "image/svg+xml")
    const content = doc.documentElement?.getAttribute("content")
    if (content?.includes("<mxfile")) return content
    return null
}

export interface ChatEngine {
    messages: UIMessage[]
    status: ChatStatus
    error: Error | undefined
    isBusy: boolean
    setMessages: (messages: UIMessage[]) => void
    stop: () => void

    input: string
    setInput: (value: string) => void
    files: File[]
    pdfData: Map<File, FileData>
    setFiles: (files: File[]) => void
    urlData: Map<string, UrlData>
    setUrlData: React.Dispatch<React.SetStateAction<Map<string, UrlData>>>
    isExtractingAttachments: boolean

    submit: (overrideText?: string) => Promise<void>
    sendTemplate: (prompt: string) => void
    improveWithSuggestions: (feedback: string) => Promise<void>
    regenerate: (messageIndex: number) => Promise<void>
    retryLastMessage: () => void
    editMessage: (messageIndex: number, newText: string) => Promise<void>

    /** Selected shapes that will be sent with the next message */
    chatSelection: SelectedCell[]
    dismissSelection: () => void

    sessions: SessionMetadata[]
    sessionsAvailable: boolean
    currentSessionId: string | null
    currentTitle: string | null
    isRestored: boolean
    newChat: () => Promise<void>
    selectSession: (id: string) => Promise<void>
    deleteSession: (id: string) => Promise<void>
    renameSession: (title: string) => Promise<void>
    openDiagramFile: (file: File) => Promise<void>

    restoreVersion: (versionId: string) => void
    undoVersion: (versionId: string) => void

    langfuseSessionId: string
    loadedMessageIdsRef: React.RefObject<Set<string>>
    validationStates: Record<string, ValidationState>
    modelConfig: UseModelConfigReturn
}

const ChatEngineContext = createContext<ChatEngine | null>(null)

export function useChatEngine(): ChatEngine {
    const engine = useContext(ChatEngineContext)
    if (!engine) {
        throw new Error("useChatEngine must be used inside ChatEngineProvider")
    }
    return engine
}

export function ChatEngineProvider({
    children,
}: {
    children: React.ReactNode
}) {
    const {
        loadDiagram: onDisplayChart,
        handleExport: onExport,
        exportResolversRef,
        chartXML,
        chartXMLRef: liveChartXMLRef,
        latestSvg,
        clearDiagram,
        getThumbnailSvg,
        getVersionSvg,
        captureValidationPng,
        setCommitHandler,
        startTurn,
    } = useDiagram()

    const dict = useDictionary()
    const router = useRouter()
    const pathname = usePathname()
    const searchParams = useSearchParams()
    const urlSessionId = searchParams.get("session")

    const settings = useSettingsStore()
    const openSettings = useUiStore((s) => s.openSettings)
    const focusComposer = useUiStore((s) => s.focusComposer)

    const onFetchChart = () => {
        // Waits for the reply to its own export, by its tag
        const tag = onExport()
        return Promise.race([
            new Promise<string>((resolve) => {
                if (tag) exportResolversRef.current[tag] = resolve
            }),
            new Promise<string>((_, reject) => {
                setTimeout(() => {
                    delete exportResolversRef.current[tag]
                    reject(new Error("Chart export timed out after 10 seconds"))
                }, 10000)
            }),
        ])
    }

    const { files, pdfData, handleFileChange, setFiles } = useFileProcessor({
        tooLong: dict.errors.fileTooLong,
        readFailed: dict.errors.fileReadFailed,
    })
    const [urlData, setUrlData] = useState<Map<string, UrlData>>(new Map())
    const filesRef = useRef(files)
    filesRef.current = files

    const modelConfig = useModelConfig()
    const sessionManager = useSessionManager({ initialSessionId: urlSessionId })

    const [input, setInputState] = useState("")
    const [dailyRequestLimit, setDailyRequestLimit] = useState(0)
    const [dailyTokenLimit, setDailyTokenLimit] = useState(0)
    const [tpmLimit, setTpmLimit] = useState(0)

    // Restore the unsent input from sessionStorage
    useEffect(() => {
        const savedInput = sessionStorage.getItem(SESSION_STORAGE_INPUT_KEY)
        if (savedInput) {
            inputRef.current = savedInput
            setInputState(savedInput)
        }
    }, [])

    // Latest composer contents, for code that runs after an await
    const inputRef = useRef("")
    const setInput = useCallback((value: string) => {
        inputRef.current = value
        setInputState(value)
        if (value) sessionStorage.setItem(SESSION_STORAGE_INPUT_KEY, value)
        else sessionStorage.removeItem(SESSION_STORAGE_INPUT_KEY)
    }, [])

    useEffect(() => {
        fetch(getApiEndpoint("/api/config"))
            .then((res) => res.json())
            .then((data) => {
                setDailyRequestLimit(data.dailyRequestLimit || 0)
                setDailyTokenLimit(data.dailyTokenLimit || 0)
                setTpmLimit(data.tpmLimit || 0)
            })
            .catch(() => {})
    }, [])

    const quotaManager = useQuotaManager({
        dailyRequestLimit,
        dailyTokenLimit,
        tpmLimit,
        onConfigModel: () => openSettings("models"),
    })

    // Session ID for Langfuse tracing (restored from localStorage if available)
    const [langfuseSessionId, setLangfuseSessionId] = useState(() => {
        if (typeof window !== "undefined") {
            const saved = localStorage.getItem(STORAGE_KEYS.sessionId)
            if (saved) return saved
        }
        return newLangfuseSessionId()
    })
    useEffect(() => {
        localStorage.setItem(STORAGE_KEYS.sessionId, langfuseSessionId)
    }, [langfuseSessionId])

    // Title for a session that is not saved yet (opened file, early rename)
    const [pendingTitle, setPendingTitleState] = useState<string | null>(null)
    const pendingTitleRef = useRef<string | null>(null)
    const setPendingTitle = (title: string | null) => {
        pendingTitleRef.current = title
        setPendingTitleState(title)
    }

    // XML snapshot taken before each user message (keyed by message index)
    const xmlSnapshotsRef = useRef<Map<number, string>>(new Map())
    // Index of the user message whose turn is running (versions belong to it)
    const currentTurnRef = useRef(0)

    const hasRestoredRef = useRef(false)
    const [isRestored, setIsRestored] = useState(false)

    // Latest chartXML for callbacks (avoids stale closures)
    const chartXMLRef = useRef(chartXML)
    // Session loaded without a diagram (prevents thumbnail contamination)
    const justLoadedSessionIdRef = useRef<string | null>(null)
    useEffect(() => {
        chartXMLRef.current = chartXML
        if (chartXML) justLoadedSessionIdRef.current = null
    }, [chartXML])

    const latestSvgRef = useRef(latestSvg)
    useEffect(() => {
        latestSvgRef.current = latestSvg
    }, [latestSvg])

    const autoRetryCountRef = useRef(0)
    const continuationRetryCountRef = useRef(0)
    // Partial XML when output was truncated; non-empty means continuation mode
    const partialXmlRef = useRef<string>("")
    // Tool calls already applied, so remounts never replay old outputs
    const processedToolCallsRef = useRef<Set<string>>(new Set())
    // Set by Stop until the user sends the next message
    const stoppedRef = useRef(false)
    const preparingSendRef = useRef(false)
    // Presses of Stop: a check that began before one still knows of it after
    // the next message clears stoppedRef
    const stopCountRef = useRef(0)
    // Diagram before each streamed display_diagram and edit_diagram preview
    // (key: toolCallId), shared between the streaming preview and the tool
    // handler
    const editDiagramOriginalXmlRef = useRef<Map<string, string>>(new Map())
    const saveDebounceRef = useRef<ReturnType<typeof setTimeout> | null>(null)

    const [validationStates, setValidationStates] = useState<
        Record<string, ValidationState>
    >({})
    const handleValidationStateChange = useCallback(
        (toolCallId: string, state: ValidationState) => {
            setValidationStates((prev) => ({ ...prev, [toolCallId]: state }))
        },
        [],
    )

    // Failed VLM validations in the current user turn (reset on user action)
    const validationRetryCountRef = useRef(0)
    const { validateWithFallback, cancel: cancelValidation } =
        useValidateDiagram()

    const { handleToolCall } = useDiagramToolHandlers({
        partialXmlRef,
        editDiagramOriginalXmlRef,
        processedToolCallsRef,
        validationRetryCountRef,
        // A preview undone just before the tool call is in this one already
        chartXMLRef: liveChartXMLRef,
        onDisplayChart,
        onFetchChart,
        captureValidationPng,
        validateDiagram: validateWithFallback,
        enableVlmValidation: settings.vlmValidationEnabled,
        sessionId: langfuseSessionId,
        watchStop: () => {
            const stopsBefore = stopCountRef.current
            return () =>
                stoppedRef.current || stopCountRef.current !== stopsBefore
        },
        onValidationStateChange: handleValidationStateChange,
    })

    const transport = useMemo(
        () => new DefaultChatTransport({ api: getApiEndpoint("/api/chat") }),
        [],
    )

    const {
        messages,
        sendMessage,
        addToolOutput,
        status,
        error,
        setMessages,
        stop,
    } = useChat({
        transport,
        onToolCall: async ({ toolCall }) => {
            await handleToolCall({ toolCall }, addToolOutput)
        },
        onFinish: ({ message, isAbort, isError }) => {
            // Stopped or failed: tool calls still streaming never reach the
            // tool handler. Mark them handled so a later render of the
            // stream does not draw their preview again.
            if (!isAbort && !isError) return
            for (const part of message.parts as any[]) {
                if (part.state === "input-streaming" && part.toolCallId) {
                    processedToolCallsRef.current.add(part.toolCallId)
                }
            }
        },
        onError: (error) => {
            // A diagram still streaming when the request failed never
            // reaches the tool handler: undo its preview. Only previews not
            // handled yet are stored, and the first one holds the diagram
            // before any of them.
            const [originalXml] = editDiagramOriginalXmlRef.current.values()
            if (originalXml) onDisplayChart(originalXml, true, "revert")
            editDiagramOriginalXmlRef.current.clear()

            // Server errors are JSON: a quota limit ({type: request, token or
            // tpm}), a provider error ({type: "provider", code, message}) or
            // {error}. The SDK puts the response body in error.message.
            let data: any = null
            try {
                data = JSON.parse(error.message)
            } catch {
                // Plain text, e.g. a network failure in the browser
            }
            if (data?.type === "request") {
                quotaManager.showQuotaLimitToast(data.used, data.limit)
                return
            }
            if (data?.type === "token") {
                quotaManager.showTokenLimitToast(data.used, data.limit)
                return
            }
            if (data?.type === "tpm") {
                quotaManager.showTPMLimitToast(data.limit)
                return
            }

            const isAccessCodeError = String(
                data?.error ?? error.message,
            ).includes("Invalid or missing access code")
            if (!isAccessCodeError) console.error("Chat error:", error)

            // A hint the user can act on, then the provider's own words
            let text: string = error.message
            let openModelConfig = false
            if (data?.type === "provider") {
                const hints = dict.errors.llm as Record<string, string>
                const hint = hints[data.code]
                text =
                    hint && data.message
                        ? `${hint}\n\n${data.message}`
                        : hint || data.message
                openModelConfig = [
                    "invalid_api_key",
                    "forbidden",
                    "model_not_found",
                    "server_key_forbidden",
                ].includes(data.code)
            } else if (typeof data?.error === "string") {
                text = data.error
            } else if (error.message === "Failed to fetch") {
                text = dict.errors.networkError
            }

            // A system message, so it can be cleared with the conversation
            setMessages((currentMessages) => [
                ...currentMessages,
                {
                    id: `error-${Date.now()}`,
                    role: "system" as const,
                    parts: [{ type: "text" as const, text }],
                    // The message shows a button that opens model settings
                    ...(openModelConfig && {
                        metadata: { openModelConfig: true },
                    }),
                },
            ])

            if (isAccessCodeError) openSettings("advanced")
        },
        // Re-render streamed messages at most every 150 ms. The streaming
        // diagram preview draws on each update, so this also limits redraws
        experimental_throttle: 150,
        sendAutomaticallyWhen: ({ messages }) => {
            // The user stopped: a tool result that arrives later (a VLM
            // check still running) must not start a new request
            if (stoppedRef.current) return false

            const isInContinuationMode = partialXmlRef.current.length > 0
            const shouldRetry = hasToolErrors(messages)

            if (!shouldRetry) {
                autoRetryCountRef.current = 0
                continuationRetryCountRef.current = 0
                partialXmlRef.current = ""
                return false
            }

            if (isInContinuationMode) {
                if (
                    continuationRetryCountRef.current >=
                    MAX_CONTINUATION_RETRY_COUNT
                ) {
                    toast.error(
                        formatMessage(dict.errors.continuationRetryLimit, {
                            max: MAX_CONTINUATION_RETRY_COUNT,
                        }),
                    )
                    continuationRetryCountRef.current = 0
                    partialXmlRef.current = ""
                    return false
                }
                continuationRetryCountRef.current++
            } else {
                if (autoRetryCountRef.current >= MAX_AUTO_RETRY_COUNT) {
                    toast.error(
                        formatMessage(dict.errors.retryLimit, {
                            max: MAX_AUTO_RETRY_COUNT,
                        }),
                    )
                    autoRetryCountRef.current = 0
                    partialXmlRef.current = ""
                    return false
                }
                autoRetryCountRef.current++
            }
            return true
        },
    })

    const isBusy = status === "streaming" || status === "submitted"
    // isBusy of the latest render, for code that runs after an await
    const busyRef = useRef(isBusy)
    busyRef.current = isBusy
    // While an answer is sent or streams in, the chat must stay on screen:
    // the rest of the answer and its diagram would land in another chat
    const isTurnActive = () => busyRef.current || preparingSendRef.current

    // Message IDs loaded from a session (skip animations and replays)
    const loadedMessageIdsRef = useRef<Set<string>>(new Set())

    useDiagramStreaming({
        messages,
        processedToolCallsRef,
        editDiagramOriginalXmlRef,
        loadedMessageIdsRef,
    })

    const { restoreVersion, undoVersion } = useVersions({
        loadDiagram: onDisplayChart,
        getVersionSvg,
        setCommitHandler,
        currentTurnRef,
        chartXMLRef: liveChartXMLRef,
        isBusy,
        isTurnActive,
        getChatGeneration: sessionManager.getChatGeneration,
    })

    // ---------------------------------------------------------------------
    // Selection sent with the next message
    // ---------------------------------------------------------------------

    const selection = useCanvasStore((s) => s.selection)
    // The model sees and edits the first page only
    const onFirstPage = useCanvasStore(
        (s) => s.pages.length === 0 || s.pages[0]?.id === s.currentPageId,
    )
    const [dismissedSelectionKey, setDismissedSelectionKey] = useState("")
    const selectionKey = selection.map((c) => c.id).join(",")
    const chatSelection =
        !onFirstPage || selectionKey === dismissedSelectionKey ? [] : selection
    // Once nothing is selected, picking the same shapes again attaches them
    useEffect(() => {
        if (selection.length === 0) setDismissedSelectionKey("")
    }, [selection.length])
    const dismissSelection = useCallback(
        () => setDismissedSelectionKey(selectionKey),
        [selectionKey],
    )

    // ---------------------------------------------------------------------
    // Sessions
    // ---------------------------------------------------------------------

    const messagesRef = useRef(messages)
    useEffect(() => {
        messagesRef.current = messages
    }, [messages])

    // Last synced session ID, to detect external changes (URL back/forward)
    const lastSyncedSessionIdRef = useRef<string | null>(null)
    // Message arrays of our own saves. A session holding one of them was
    // created by our own save, so it must not be treated as an external
    // switch (with two saves of a new chat at once, the first creates it).
    const savedMessagesRef = useRef(new WeakSet<object>())
    // Session was just loaded (skip the auto-save that would reorder it)
    const justLoadedSessionRef = useRef(false)

    const syncUIWithSession = useCallback(
        (
            data:
                | (SessionData & {
                      diagramHistory?: { svg: string; xml: string }[]
                  })
                | null,
        ) => {
            const hasRealDiagram = isRealDiagram(data?.diagramXml)
            // Its Undo would put the previous chat's canvas here
            toast.dismiss("version-restored")
            if (data) {
                loadedMessageIdsRef.current = new Set(
                    (data.messages as { id: string }[]).map((m) => m.id),
                )
                setMessages(data.messages as unknown as UIMessage[])
                xmlSnapshotsRef.current = new Map(data.xmlSnapshots)
                if (hasRealDiagram) {
                    onDisplayChart(data.diagramXml, true)
                    chartXMLRef.current = data.diagramXml
                } else {
                    clearDiagram()
                    chartXMLRef.current = ""
                    latestSvgRef.current = ""
                }
                useVersionsStore
                    .getState()
                    .setVersions(
                        data.versions ??
                            versionsFromLegacyHistory(data.diagramHistory),
                    )
            } else {
                loadedMessageIdsRef.current = new Set()
                setMessages([])
                xmlSnapshotsRef.current.clear()
                clearDiagram()
                chartXMLRef.current = ""
                latestSvgRef.current = ""
                useVersionsStore.getState().clear()
            }
        },
        [setMessages, onDisplayChart, clearDiagram],
    )

    const buildSessionData = useCallback(
        async (
            options: { withThumbnail?: boolean } = {},
        ): Promise<SessionData> => {
            const currentDiagramXml = chartXMLRef.current || ""
            const hasRealDiagram = isRealDiagram(currentDiagramXml)
            let thumbnailDataUrl: string | undefined
            if (hasRealDiagram && options.withThumbnail) {
                const freshThumb = await getThumbnailSvg()
                if (freshThumb) {
                    latestSvgRef.current = freshThumb
                    thumbnailDataUrl = freshThumb
                } else if (latestSvgRef.current) {
                    thumbnailDataUrl = latestSvgRef.current
                }
            }
            const messages = sanitizeMessages(messagesRef.current)
            savedMessagesRef.current.add(messages)
            return {
                messages,
                xmlSnapshots: Array.from(xmlSnapshotsRef.current.entries()),
                diagramXml: currentDiagramXml,
                thumbnailDataUrl,
                versions: useVersionsStore.getState().versions,
                title: pendingTitleRef.current ?? undefined,
            }
        },
        [getThumbnailSvg],
    )

    // Restore from the session manager once it has loaded
    useLayoutEffect(() => {
        if (hasRestoredRef.current) return
        if (sessionManager.isLoading) return
        hasRestoredRef.current = true
        try {
            const currentSession = sessionManager.currentSession
            if (currentSession) {
                justLoadedSessionRef.current = true
                syncUIWithSession(currentSession)
            }
            lastSyncedSessionIdRef.current = sessionManager.currentSessionId
        } catch (error) {
            console.error("Failed to restore session:", error)
            toast.error(dict.errors.sessionCorrupted)
        } finally {
            setIsRestored(true)
        }
    }, [
        sessionManager.isLoading,
        sessionManager.currentSession,
        syncUIWithSession,
        dict.errors.sessionCorrupted,
    ])

    // Sync UI when the session changes externally (URL back/forward)
    useEffect(() => {
        if (!isRestored) return
        if (!sessionManager.isAvailable) return
        const newSessionId = sessionManager.currentSessionId
        const newSession = sessionManager.currentSession
        if (newSessionId === lastSyncedSessionIdRef.current) return
        const isOwnNewSession =
            !!newSession && savedMessagesRef.current.has(newSession.messages)
        lastSyncedSessionIdRef.current = newSessionId
        if (isOwnNewSession) {
            setPendingTitle(null)
            return
        }
        justLoadedSessionRef.current = true
        syncUIWithSession(newSession)
    }, [
        isRestored,
        sessionManager.isAvailable,
        sessionManager.currentSessionId,
        sessionManager.currentSession,
        syncUIWithSession,
    ])

    const {
        isAvailable: sessionIsAvailable,
        currentSessionId,
        saveCurrentSession,
        getChatGeneration,
        getSaveTicket,
    } = sessionManager
    // Ref avoids re-running the save effect after every save
    const saveCurrentSessionRef = useRef(saveCurrentSession)
    saveCurrentSessionRef.current = saveCurrentSession

    const versions = useVersionsStore((s) => s.versions)

    // Debounced auto-save, only when not streaming
    useEffect(() => {
        if (!hasRestoredRef.current) return
        if (!sessionIsAvailable) return
        if (status === "streaming" || status === "submitted") return
        if (justLoadedSessionRef.current) {
            justLoadedSessionRef.current = false
            return
        }
        if (saveDebounceRef.current) clearTimeout(saveDebounceRef.current)

        // Capture the chat on screen at schedule time; the save is dropped
        // if another chat is on screen by the time it runs
        const scheduledForChat = getChatGeneration()
        const hasDiagramNow = isRealDiagram(chartXMLRef.current)
        const isNoDiagramSession =
            justLoadedSessionIdRef.current === currentSessionId

        saveDebounceRef.current = setTimeout(async () => {
            try {
                if (messagesRef.current.length > 0 || hasDiagramNow) {
                    // Taken before the data is read, for the chat it was
                    // scheduled for
                    const ticket = {
                        ...getSaveTicket(),
                        generation: scheduledForChat,
                    }
                    const sessionData = await buildSessionData({
                        withThumbnail: hasDiagramNow && !isNoDiagramSession,
                    })
                    await saveCurrentSessionRef.current(sessionData, ticket)
                }
            } catch (error) {
                console.error("Failed to save session:", error)
            }
        }, SAVE_DEBOUNCE_MS)

        return () => {
            if (saveDebounceRef.current) clearTimeout(saveDebounceRef.current)
        }
    }, [
        chartXML,
        messages,
        versions,
        status,
        sessionIsAvailable,
        currentSessionId,
        getChatGeneration,
        getSaveTicket,
        buildSessionData,
    ])

    // Put the session in the URL once it exists
    useEffect(() => {
        if (sessionManager.currentSessionId && !urlSessionId) {
            router.replace(`?session=${sessionManager.currentSessionId}`, {
                scroll: false,
            })
        }
    }, [sessionManager.currentSessionId, urlSessionId, router])

    // Save when the page is hidden (more reliable than beforeunload for IndexedDB)
    useEffect(() => {
        if (!sessionManager.isAvailable) return
        const handleVisibilityChange = async () => {
            if (
                document.visibilityState === "hidden" &&
                (messagesRef.current.length > 0 ||
                    isRealDiagram(chartXMLRef.current))
            ) {
                try {
                    const ticket = sessionManager.getSaveTicket()
                    const sessionData = await buildSessionData({
                        withThumbnail: false,
                    })
                    await sessionManager.saveCurrentSession(sessionData, ticket)
                } catch (error) {
                    console.error(
                        "Failed to save session on visibility change:",
                        error,
                    )
                }
            }
        }
        document.addEventListener("visibilitychange", handleVisibilityChange)
        return () =>
            document.removeEventListener(
                "visibilitychange",
                handleVisibilityChange,
            )
    }, [sessionManager, buildSessionData])

    // ---------------------------------------------------------------------
    // Sending
    // ---------------------------------------------------------------------

    const isExtractingAttachments =
        files.some((f) => pdfData.get(f)?.isExtracting) ||
        Array.from(urlData.values()).some((d) => d.isExtracting)

    // Append PDF, text file and URL content to the user's text; images
    // become file parts when imageParts is given
    const processFilesAndAppendContent = async (
        baseText: string,
        imageParts?: any[],
    ): Promise<string> => {
        let userText = baseText
        for (const file of files) {
            if (isPdfFile(file)) {
                const extracted = pdfData.get(file)
                if (extracted?.text) {
                    userText += `\n\n[PDF: ${file.name}]\n${extracted.text}`
                }
            } else if (isTextFile(file)) {
                const extracted = pdfData.get(file)
                if (extracted?.text) {
                    userText += `\n\n[File: ${file.name}]\n${extracted.text}`
                }
            } else if (imageParts) {
                const reader = new FileReader()
                const dataUrl = await new Promise<string>((resolve) => {
                    reader.onload = () => resolve(reader.result as string)
                    reader.readAsDataURL(file)
                })
                imageParts.push({
                    type: "file",
                    url: dataUrl,
                    mediaType: file.type,
                })
            }
        }
        for (const [url, data] of urlData) {
            if (data.content) {
                userText += `\n\n[URL: ${url}]\nTitle: ${data.title}\n\n${data.content}`
            }
        }
        return userText
    }

    const getPreviousXml = (beforeIndex: number): string => {
        const snapshotKeys = Array.from(xmlSnapshotsRef.current.keys())
            .filter((k) => k < beforeIndex)
            .sort((a, b) => b - a)
        return snapshotKeys.length > 0
            ? getFirstPageXml(
                  xmlSnapshotsRef.current.get(snapshotKeys[0]) || "",
              )
            : ""
    }

    const sendChatMessage = (
        parts: any,
        xml: string,
        previousXml: string,
        turnIndex: number,
        selectedCells: MessageSelection,
    ) => {
        autoRetryCountRef.current = 0
        continuationRetryCountRef.current = 0
        validationRetryCountRef.current = 0
        partialXmlRef.current = ""
        stoppedRef.current = false
        currentTurnRef.current = turnIndex
        // Busy from now on, before the next render says so
        busyRef.current = true
        startTurn()

        const config = getSelectedAIConfig()
        const { customSystemMessage, minimalStyle, maxOutputTokens } =
            useSettingsStore.getState()

        const selected = selectedCells.map(({ id, label }) => ({ id, label }))
        const options = {
            body: {
                xml,
                previousXml,
                sessionId: langfuseSessionId,
                customSystemMessage,
                ...(selected.length > 0 && { selectedCells: selected }),
            },
            headers: buildChatHeaders(config, {
                minimalStyle,
                maxOutputTokens,
            }),
        }
        // The selection stays with the message, for regenerate and retry
        sendMessage(
            {
                parts,
                ...(selected.length > 0 && {
                    metadata: { selectedCells: selected },
                }),
            },
            options,
        )
    }

    // Export the current diagram, snapshot it for this message, and send
    const sendWithCurrentDiagram = async (
        parts: any[],
        selectedCells: MessageSelection = [],
    ) => {
        const chartXml = formatXML(await onFetchChart())
        const turnIndex = messagesRef.current.length
        const previousXml = getPreviousXml(turnIndex)
        // Snapshot the full multi-page document (kept fresh by autosave) so
        // regenerate/edit can restore every page; the model gets page 1 only
        xmlSnapshotsRef.current.set(turnIndex, chartXMLRef.current || chartXml)
        sendChatMessage(parts, chartXml, previousXml, turnIndex, selectedCells)
    }

    const clearComposer = () => {
        setInput("")
        setFiles([])
        setUrlData(new Map())
    }

    const submitInput = async (overrideText?: string) => {
        const text = overrideText ?? input
        if (!text.trim() || isBusy || isExtractingAttachments) return

        // Cached example: only with no messages and an empty canvas (same
        // rule as the server)
        if (
            messagesRef.current.length === 0 &&
            !hasCells(chartXMLRef.current || "")
        ) {
            // The file name keeps a user's own file from matching an example
            const cached = findCachedResponse(
                text.trim(),
                files.length > 0,
                files.length === 1 ? files[0].name : undefined,
            )
            if (cached) {
                const toolCallId = `cached-${Date.now()}`
                const userText = await processFilesAndAppendContent(text)
                currentTurnRef.current = 0
                startTurn()
                setMessages([
                    {
                        id: `user-${Date.now()}`,
                        role: "user" as const,
                        parts: [{ type: "text" as const, text: userText }],
                    },
                    {
                        id: `assistant-${Date.now()}`,
                        role: "assistant" as const,
                        parts: [
                            {
                                type: "tool-display_diagram" as const,
                                toolCallId,
                                state: "output-available" as const,
                                input: { xml: cached.xml },
                                output: "Successfully displayed the diagram.",
                            },
                        ],
                    },
                ] as any)
                // Snapshot the canvas before the example so editing this message works
                xmlSnapshotsRef.current.set(
                    0,
                    chartXMLRef.current || BLANK_MXFILE,
                )
                // Load its diagram here: these messages never reach the
                // tool handler
                const prepared = prepareNewDiagram(cached.xml, {
                    pageId: "page-1",
                    pageName: "Page-1",
                })
                if (prepared.ok) {
                    onDisplayChart(prepared.xml, true, "commit", {
                        toolCallId,
                    })
                }
                clearComposer()
                return
            }
        }

        const parts: any[] = []
        // The backend only reads the first text part, so combine them
        const userText = await processFilesAndAppendContent(text, parts)
        parts.unshift({ type: "text", text: userText })
        // Clear right away, so text typed while the diagram exports is kept
        const sentFiles = files
        const sentUrls = urlData
        clearComposer()
        setDismissedSelectionKey(selectionKey)
        try {
            await sendWithCurrentDiagram(parts, chatSelection)
        } catch (error) {
            console.error("Error fetching chart data:", error)
            toast.error(dict.errors.failedToExport)
            // Nothing was sent: give the message back, unless something new
            // was typed or attached meanwhile
            if (!inputRef.current && filesRef.current.length === 0) {
                setInput(text)
                setFiles(sentFiles)
                setUrlData(sentUrls)
            }
        }
    }

    const submit = async (overrideText?: string) => {
        // While a send is prepared (attachments read, diagram exported) the
        // status is still "ready": a second Enter or click would send the
        // message again
        if (preparingSendRef.current) return
        preparingSendRef.current = true
        try {
            await submitInput(overrideText)
        } finally {
            preparingSendRef.current = false
        }
    }

    // Templates fill the input and send it right away (attachments included)
    const sendTemplate = (prompt: string) => {
        setInput(prompt)
        submit(prompt)
    }

    const improveWithSuggestions = async (feedback: string) => {
        if (isBusy) return
        try {
            await sendWithCurrentDiagram([{ type: "text", text: feedback }])
        } catch (error) {
            console.error("Error fetching chart data:", error)
            toast.error(dict.errors.failedToExport)
        }
    }

    const restoreDiagramFromSnapshot = (savedXml: string) => {
        onDisplayChart(savedXml, true)
        chartXMLRef.current = savedXml
    }

    const cleanupAfterTurn = (messageIndex: number, keepVersions = false) => {
        for (const key of xmlSnapshotsRef.current.keys()) {
            if (key > messageIndex) xmlSnapshotsRef.current.delete(key)
        }
        // The canvas goes back to the snapshot: no kept version is on it
        if (keepVersions) useVersionsStore.getState().clearCanvasFlags()
        else useVersionsStore.getState().removeFromTurn(messageIndex)
    }

    const regenerate = async (messageIndex: number) => {
        if (isBusy) return
        let userMessageIndex = messageIndex - 1
        while (
            userMessageIndex >= 0 &&
            messages[userMessageIndex].role !== "user"
        ) {
            userMessageIndex--
        }
        if (userMessageIndex < 0) return

        const userMessage = messages[userMessageIndex]
        const userParts = userMessage.parts
        if (!userParts?.some((p: any) => p.type === "text")) return

        const savedXml = xmlSnapshotsRef.current.get(userMessageIndex)
        if (!savedXml) {
            console.error(
                "No saved XML snapshot for message index:",
                userMessageIndex,
            )
            return
        }

        const previousXml = getPreviousXml(userMessageIndex)
        restoreDiagramFromSnapshot(savedXml)
        cleanupAfterTurn(userMessageIndex)

        // Remove the user message onwards (sendMessage re-adds it)
        flushSync(() => setMessages(messages.slice(0, userMessageIndex)))
        sendChatMessage(
            userParts,
            getFirstPageXml(savedXml),
            previousXml,
            userMessageIndex,
            selectionOf(userMessage),
        )
    }

    // keepVersions: a retry keeps what the failed turn drew (in the strip)
    const editMessage = async (
        messageIndex: number,
        newText: string,
        keepVersions = false,
    ) => {
        if (isBusy) return
        const message = messages[messageIndex]
        if (message?.role !== "user") return

        const savedXml = xmlSnapshotsRef.current.get(messageIndex)
        if (!savedXml) {
            console.error(
                "No saved XML snapshot for message index:",
                messageIndex,
            )
            return
        }

        const previousXml = getPreviousXml(messageIndex)
        restoreDiagramFromSnapshot(savedXml)
        cleanupAfterTurn(messageIndex, keepVersions)

        // The edit box only shows the typed text; keep appended file content
        const newParts = message.parts?.map((part: any) => {
            if (part.type === "text") {
                const appended =
                    part.text.match(APPENDED_FILE_SECTIONS_PATTERN)?.[0] ?? ""
                return { ...part, text: newText + appended }
            }
            return part
        }) || [{ type: "text", text: newText }]

        flushSync(() => setMessages(messages.slice(0, messageIndex)))
        sendChatMessage(
            newParts,
            getFirstPageXml(savedXml),
            previousXml,
            messageIndex,
            selectionOf(message),
        )
    }

    /** Send the last user message again (after an error) */
    const retryLastMessage = () => {
        const index = messages.map((m) => m.role).lastIndexOf("user")
        if (index < 0) return
        editMessage(index, getUserOriginalText(messages[index]), true)
    }

    const handleStop = useCallback(() => {
        stoppedRef.current = true
        stopCountRef.current++
        // A running screenshot check holds up the chat (the SDK waits for
        // the tool handler): end it, so the call gets its result now
        cancelValidation()
        const lastMessage = messages[messages.length - 1]
        // Calls the tool handler already took can still show as streaming:
        // the messages update at most every 150 ms (useChat throttle)
        const toolParts = lastMessage?.parts?.filter(
            (part: any) =>
                part.type?.startsWith("tool-") &&
                part.state === "input-streaming" &&
                !processedToolCallsRef.current.has(part.toolCallId),
        )
        for (const part of (toolParts ?? []) as any[]) {
            if (part.toolCallId) {
                addToolOutput({
                    tool: part.type.replace("tool-", ""),
                    toolCallId: part.toolCallId,
                    state: "output-error",
                    errorText: "Stopped by user",
                })
            }
        }
        stop()
    }, [messages, addToolOutput, stop, cancelValidation])

    // ---------------------------------------------------------------------
    // Session actions
    // ---------------------------------------------------------------------

    // The current chat could not be saved (storage full). The list where
    // old chats can be deleted shows only in an empty chat, so let the user
    // go on without saving. It replaces the plain message, and has its own
    // id so a later failed auto-save does not take its button away.
    const offerToContinueUnsaved = useCallback(
        (proceed: () => void) => {
            toast.dismiss("session-save-failed")
            toast.error(dict.errors.sessionSaveFailedLeave, {
                id: "session-save-leave",
                duration: 15000,
                action: {
                    label: dict.errors.continueWithoutSaving,
                    onClick: proceed,
                },
            })
        },
        [dict],
    )

    // A new turn makes the offer stale: going on would clear the chat while
    // the answer streams in
    useEffect(() => {
        if (isBusy) toast.dismiss("session-save-leave")
    }, [isBusy])

    // Save the chat on screen before leaving it (also a diagram drawn
    // without messages). If that failed (storage full), stay on it unless
    // the user goes on without saving: then proceed runs later. Resolves to
    // whether the caller can leave now.
    const saveBeforeLeaving = async (proceed: () => void) => {
        if (!sessionManager.isAvailable) return true
        if (
            messagesRef.current.length === 0 &&
            !isRealDiagram(chartXMLRef.current)
        ) {
            return true
        }
        // Of the chat on screen now, also if another one comes on screen
        // while the thumbnail is taken
        const ticket = sessionManager.getSaveTicket()
        const sessionData = await buildSessionData({ withThumbnail: true })
        if (!(await sessionManager.saveCurrentSession(sessionData, ticket))) {
            offerToContinueUnsaved(proceed)
            return false
        }
        return true
    }

    // The checks run again after saving: a message can be sent meanwhile

    const selectSession = async (id: string) => {
        if (!sessionManager.isAvailable || isTurnActive()) return
        const open = async () => {
            if (isTurnActive()) return
            const sessionData = await sessionManager.switchSession(id)
            if (!sessionData) return
            const hasRealDiagram = isRealDiagram(sessionData.diagramXml)
            justLoadedSessionRef.current = true
            // Use the new session's thumbnail, never the previous one's
            latestSvgRef.current = sessionData.thumbnailDataUrl || ""
            justLoadedSessionIdRef.current = hasRealDiagram ? null : id
            setValidationStates({})
            syncUIWithSession(sessionData)
            router.replace(`?session=${id}`, { scroll: false })
        }
        if (await saveBeforeLeaving(open)) await open()
    }

    const deleteSession = useCallback(
        async (id: string) => {
            if (!sessionManager.isAvailable) return
            if (id === sessionManager.currentSessionId && isTurnActive()) return
            const result = await sessionManager.deleteSession(id)
            if (result.wasCurrentSession) {
                syncUIWithSession(null)
                router.replace(pathname, { scroll: false })
            }
        },
        [sessionManager, syncUIWithSession, router, pathname, isBusy],
    )

    // Clear everything for a new conversation (the old one is saved)
    const startNewChat = () => {
        // Clear session state BEFORE the URL, so the URL effect does not
        // bring the old session back
        sessionManager.clearCurrentSession()
        toast.dismiss("version-restored")
        setMessages([])
        setInput("")
        clearDiagram()
        useVersionsStore.getState().clear()
        setValidationStates({})
        handleFileChange([])
        setUrlData(new Map())
        setLangfuseSessionId(newLangfuseSessionId())
        xmlSnapshotsRef.current.clear()
        setPendingTitle(null)
        router.replace(pathname, { scroll: false })
    }

    const newChat = async () => {
        if (isTurnActive()) return
        const start = () => {
            if (isTurnActive()) return
            startNewChat()
            toast.success(dict.dialogs.clearSuccess)
            focusComposer()
        }
        if (!(await saveBeforeLeaving(start))) return
        // The list shows the chat just saved
        if (sessionManager.isAvailable) await sessionManager.refreshSessions()
        start()
    }

    const renameSession = useCallback(
        async (title: string) => {
            const id = sessionManager.currentSessionId
            if (!id) {
                // Not saved yet: use the title when the session is created
                setPendingTitle(title.trim().slice(0, 100) || null)
                return
            }
            await sessionManager.renameSession(id, title)
        },
        [sessionManager],
    )

    const openDiagramFile = async (file: File) => {
        if (isTurnActive()) return
        const xml = await readDiagramFile(file).catch(() => null)
        if (!xml) {
            toast.error(
                formatMessage(dict.errors.notADiagramFile, {
                    name: file.name,
                }),
            )
            return
        }
        // A loaded file starts a fresh conversation named after it
        const open = () => {
            if (isTurnActive()) return
            startNewChat()
            // This empty chat is ours: the effect that follows session
            // changes must not clear the file it is about to show
            lastSyncedSessionIdRef.current = null
            setPendingTitle(
                file.name.replace(/\.(drawio\.svg|drawio|xml|svg)$/i, ""),
            )
            // Pages inflated: the previews and the model read the XML before
            // draw.io saves it again
            onDisplayChart(inflatePages(xml), true)
        }
        if (await saveBeforeLeaving(open)) open()
    }

    const engine: ChatEngine = {
        messages,
        status: status as ChatStatus,
        error,
        isBusy,
        setMessages: setMessages as (messages: UIMessage[]) => void,
        stop: handleStop,
        input,
        setInput,
        files,
        pdfData,
        setFiles: handleFileChange,
        urlData,
        setUrlData,
        isExtractingAttachments,
        submit,
        sendTemplate,
        improveWithSuggestions,
        regenerate,
        retryLastMessage,
        editMessage,
        chatSelection,
        dismissSelection,
        sessions: sessionManager.sessions,
        sessionsAvailable: sessionManager.isAvailable,
        currentSessionId: sessionManager.currentSessionId,
        currentTitle: sessionManager.currentSession?.title ?? pendingTitle,
        isRestored,
        newChat,
        selectSession,
        deleteSession,
        renameSession,
        openDiagramFile,
        restoreVersion,
        undoVersion,
        langfuseSessionId,
        loadedMessageIdsRef,
        validationStates,
        modelConfig,
    }

    return (
        <ChatEngineContext.Provider value={engine}>
            {children}
        </ChatEngineContext.Provider>
    )
}