supabase/apps/studio/components/ui/AIEditor/index.tsx

import { Monaco, OnMount } from '@monaco-editor/react'
import { AnimatePresence, motion } from 'framer-motion'
import type { editor as monacoEditor } from 'monaco-editor'
import { useCallback, useEffect, useRef, useState } from 'react'
import { toast } from 'sonner'
import { KeyboardShortcut } from 'ui'
import { useSetCommandMenuOpen } from 'ui-patterns/CommandMenu'

import { CodeEditor, type ValidLanguages } from '../CodeEditor/CodeEditor'
import { DiffEditor } from '../DiffEditor'
import { ResizableAIWidget } from './ResizableAIWidget'
import { getEditorSelectionParts } from './utils'
import { SIDEBAR_KEYS } from '@/components/layouts/ProjectLayout/LayoutSidebar/LayoutSidebarProvider'
import { constructHeaders } from '@/data/fetchers'
import { useLatest } from '@/hooks/misc/useLatest'
import { SHORTCUT_IDS } from '@/state/shortcuts/registry'
import { useIsShortcutEnabled } from '@/state/shortcuts/useIsShortcutEnabled'
import { useSidebarManagerSnapshot } from '@/state/sidebar-manager-state'

interface AIEditorProps {
  id?: string
  language: ValidLanguages
  value?: string
  defaultValue?: string
  aiEndpoint?: string
  aiMetadata?: {
    projectRef?: string
    connectionString?: string | null
    orgSlug?: string
    language?: string
  }
  initialPrompt?: string
  readOnly?: boolean
  autoFocus?: boolean
  className?: string
  options?: monacoEditor.IStandaloneEditorConstructionOptions
  onChange?: (value: string) => void
  onClose?: () => void
  closeShortcutEnabled?: boolean
  openAIAssistantShortcutEnabled?: boolean
  executeQuery?: () => void
  onMount?: (editor: monacoEditor.IStandaloneCodeEditor, monaco: Monaco) => void
}

// [Joshen] This has overlap with components/interfaces/SQLEditor/MonacoEditor
// Can we try to de-dupe accordingly? Perhaps the SQL Editor could use this AIEditor
// We have a tendency to create multiple versions of the monaco editor like RLSCodeEditor
// so hoping to prevent that from snowballing
export const AIEditor = ({
  language = 'javascript',
  value,
  defaultValue = '',
  aiEndpoint,
  aiMetadata,
  initialPrompt,
  readOnly = false,
  autoFocus = false,
  className = '',
  options = {},
  onChange,
  onClose,
  closeShortcutEnabled = true,
  openAIAssistantShortcutEnabled = true,
  executeQuery,
  onMount,
}: AIEditorProps) => {
  const { toggleSidebar } = useSidebarManagerSnapshot()
  const editorRef = useRef<monacoEditor.IStandaloneCodeEditor | null>(null)
  const diffEditorRef = useRef<monacoEditor.IStandaloneDiffEditor | null>(null)
  const monacoRef = useRef<Monaco | null>(null)
  const closeActionDisposableRef = useRef<{ dispose: () => void } | null>(null)

  const isCommandMenuHotkeyEnabled = useIsShortcutEnabled(SHORTCUT_IDS.COMMAND_MENU_OPEN)
  const setCommandMenuOpen = useSetCommandMenuOpen()

  const executeQueryRef = useLatest(executeQuery)
  const commandMenuHotkeyEnabledRef = useLatest(isCommandMenuHotkeyEnabled)
  const setCommandMenuOpenRef = useLatest(setCommandMenuOpen)

  const [currentValue, setCurrentValue] = useState(value || defaultValue)
  const [isDiffMode, setIsDiffMode] = useState(false)
  const [isDiffEditorMounted, setIsDiffEditorMounted] = useState(false)
  const [diffValue, setDiffValue] = useState({ original: '', modified: '' })
  const [promptState, setPromptState] = useState({
    isOpen: Boolean(initialPrompt),
    selection: '',
    beforeSelection: '',
    afterSelection: '',
    startLineNumber: 0,
    endLineNumber: 0,
  })
  const [promptInput, setPromptInput] = useState(initialPrompt || '')

  const [isCompletionLoading, setIsCompletionLoading] = useState(false)

  const complete = useCallback(
    async (
      _prompt: string,
      options?: {
        headers?: Record<string, string>
        body?: { completionMetadata?: any }
      }
    ) => {
      try {
        if (!aiEndpoint) throw new Error('AI endpoint is not configured')
        setIsCompletionLoading(true)

        const response = await fetch(aiEndpoint, {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
            ...(options?.headers ?? {}),
          },
          body: JSON.stringify({
            ...(aiMetadata ?? {}),
            ...(options?.body ?? {}),
          }),
        })

        if (!response.ok) {
          const errorText = await response.text()
          throw new Error(errorText || 'Failed to generate completion')
        }

        const text: string = await response.json()

        const meta = options?.body?.completionMetadata ?? {}
        const beforeSelection: string = meta.textBeforeCursor ?? ''
        const afterSelection: string = meta.textAfterCursor ?? ''
        const selection: string = meta.selection ?? ''

        const original = beforeSelection + selection + afterSelection
        const modified = beforeSelection + text + afterSelection

        setDiffValue({ original, modified })
        setIsDiffMode(true)
      } catch (error: any) {
        toast.error(`Failed to generate: ${error?.message ?? 'Unknown error'}`)
      } finally {
        setIsCompletionLoading(false)
      }
    },
    [aiEndpoint, aiMetadata]
  )

  const handleReset = useCallback(() => {
    setIsDiffMode(false)
    setPromptState((prev) => ({ ...prev, isOpen: false }))
    setPromptInput('')
    editorRef.current?.focus()
  }, [])

  const handleAcceptDiff = useCallback(() => {
    if (diffValue.modified) {
      const newValue = diffValue.modified
      setCurrentValue(newValue)
      onChange?.(newValue)
      handleReset()
    }
  }, [diffValue.modified, onChange, handleReset])

  const handleRejectDiff = () => {
    handleReset()
  }

  const refreshCloseAction = useCallback(() => {
    closeActionDisposableRef.current?.dispose()
    closeActionDisposableRef.current = null

    const editor = editorRef.current
    const monaco = monacoRef.current

    if (!editor || !monaco || !onClose || !closeShortcutEnabled) return

    const action = editor.addAction({
      id: 'close-editor',
      label: 'Close editor',
      keybindings: [monaco.KeyMod.CtrlCmd + monaco.KeyCode.KeyE],
      contextMenuGroupId: 'operation',
      contextMenuOrder: 0,
      run: onClose,
    })

    closeActionDisposableRef.current = action ?? null
  }, [closeShortcutEnabled, onClose])

  const handleEditorOnMount: OnMount = (
    editor: monacoEditor.IStandaloneCodeEditor,
    monaco: Monaco
  ) => {
    editorRef.current = editor
    monacoRef.current = monaco
    onMount?.(editor, monaco)
    // Set prompt state to open if promptInput exists
    if (promptInput) {
      const model = editor.getModel()
      if (model) {
        const lineCount = model.getLineCount()
        setPromptState({
          isOpen: true,
          selection: model.getValue(),
          beforeSelection: '',
          afterSelection: '',
          startLineNumber: 1,
          endLineNumber: lineCount,
        })
      }
    }

    // [Joshen] Opting to ignore "Cannot find module" errors here as users are getting
    // confused with the error highlighting when importing external modules
    monaco.languages.typescript.typescriptDefaults.setDiagnosticsOptions({
      diagnosticCodesToIgnore: [2792],
    })

    if (language === 'javascript' || language === 'typescript') {
      // The Deno libs load as raw text — via the raw-loader rules in
      // next.config.ts (Next/turbopack) and the `rawTextLoader` plugin in
      // vite.config.ts (TanStack/Vite) — and are registered as Monaco extra
      // libs. The specifiers must stay string literals so both bundlers can
      // statically analyze and code-split them; the `as string` cast keeps
      // tsc from resolving the `.d.ts` files as declaration files (TS2846 /
      // "not a module") while erasing to a plain literal for the bundlers.
      const denoTypeLibs: Record<string, Promise<{ default: string }>> = {
        'edge-runtime': import('@/public/deno/edge-runtime.d.ts' as string),
        'lib.deno': import('@/public/deno/lib.deno.d.ts' as string),
      }
      for (const [lib, loading] of Object.entries(denoTypeLibs)) {
        loading
          .then((module) => {
            monaco.languages.typescript.typescriptDefaults.addExtraLib(module.default)
          })
          .catch((error) => {
            console.error(`Failed to load Deno ${lib} typings:`, error)
          })
      }
    }

    refreshCloseAction()

    // Add AI Assistant toggle keybinding (Cmd+I)
    if (openAIAssistantShortcutEnabled) {
      editor.addAction({
        id: 'toggle-ai-assistant',
        label: 'Toggle AI Assistant',
        keybindings: [monaco.KeyMod.CtrlCmd + monaco.KeyCode.KeyI],
        run: () => {
          toggleSidebar(SIDEBAR_KEYS.AI_ASSISTANT)
        },
      })
    }

    editor.addAction({
      id: 'generate-ai',
      label: 'Generate with AI',
      keybindings: [monaco.KeyMod.CtrlCmd | monaco.KeyMod.Shift | monaco.KeyCode.KeyK],
      run: () => {
        const selectionParts = getEditorSelectionParts(editor)
        if (!selectionParts) return
        setPromptState({ isOpen: true, ...selectionParts })
      },
    })

    // Monaco claims Cmd+K as a chord prefix, which swallows the global command
    // menu shortcut while the editor is focused. Intercept it here and open the
    // command menu directly so it works the same inside and outside the editor.
    editor.addCommand(monaco.KeyMod.CtrlCmd | monaco.KeyCode.KeyK, () => {
      if (commandMenuHotkeyEnabledRef.current) {
        setCommandMenuOpenRef.current(true)
      }
    })
  }

  const handlePrompt = async (
    prompt: string,
    context: {
      beforeSelection: string
      selection: string
      afterSelection: string
    }
  ) => {
    try {
      setPromptState((prev) => ({
        ...prev,
        selection: context.selection,
        beforeSelection: context.beforeSelection,
        afterSelection: context.afterSelection,
      }))

      const headerData = await constructHeaders()
      const authorizationHeader = headerData.get('Authorization')

      await complete(prompt, {
        ...(authorizationHeader ? { headers: { Authorization: authorizationHeader } } : undefined),
        body: {
          ...aiMetadata,
          completionMetadata: {
            textBeforeCursor: context.beforeSelection,
            textAfterCursor: context.afterSelection,
            language,
            prompt,
            selection: context.selection,
          },
        },
      })
    } catch (error) {
      setPromptState((prev) => ({ ...prev, isOpen: false }))
    }
  }

  useEffect(() => {
    setCurrentValue(value || defaultValue)
  }, [value, defaultValue])

  useEffect(() => {
    if (initialPrompt) {
      setPromptInput(initialPrompt)
      setPromptState({
        isOpen: Boolean(initialPrompt),
        selection: '',
        beforeSelection: '',
        afterSelection: '',
        startLineNumber: 0,
        endLineNumber: 0,
      })
    }
  }, [initialPrompt])

  useEffect(() => {
    if (!isDiffMode) {
      setIsDiffEditorMounted(false)
    }
  }, [isDiffMode])

  useEffect(() => {
    const handleKeyboard = (event: KeyboardEvent) => {
      if (event.key === 'Escape') {
        handleReset()
      } else if (event.key === 'Enter' && (event.metaKey || event.ctrlKey) && isDiffMode) {
        event.preventDefault()
        handleAcceptDiff()
      }
    }

    window.addEventListener('keydown', handleKeyboard)
    return () => window.removeEventListener('keydown', handleKeyboard)
  }, [isDiffMode, handleAcceptDiff, handleReset])

  return (
    <div className="flex-1 overflow-hidden flex flex-col h-full relative">
      {isDiffMode ? (
        <div className="w-full h-full">
          <DiffEditor
            language={language}
            original={diffValue.original}
            modified={diffValue.modified}
            onMount={(editor: monacoEditor.IStandaloneDiffEditor) => {
              diffEditorRef.current = editor
              setIsDiffEditorMounted(true)
            }}
          />
          {isDiffEditorMounted && (
            <ResizableAIWidget
              editor={diffEditorRef.current!}
              id="ask-ai-diff"
              value={promptInput}
              onChange={setPromptInput}
              onSubmit={(prompt: string) => {
                handlePrompt(prompt, {
                  beforeSelection: promptState.beforeSelection,
                  selection: promptState.selection || diffValue.modified,
                  afterSelection: promptState.afterSelection,
                })
              }}
              onAccept={handleAcceptDiff}
              onReject={handleRejectDiff}
              onCancel={handleReset}
              isDiffVisible={true}
              isLoading={isCompletionLoading}
              startLineNumber={Math.max(0, promptState.startLineNumber)}
              endLineNumber={promptState.endLineNumber}
            />
          )}
        </div>
      ) : (
        <div className="w-full h-full relative">
          <CodeEditor
            language={language}
            autofocus={autoFocus}
            className={className}
            isReadOnly={readOnly}
            options={options}
            value={currentValue}
            onInputChange={(value) => {
              const newValue = value || ''
              setCurrentValue(newValue)
              onChange?.(newValue)
            }}
            onMount={handleEditorOnMount}
            actions={{ runQuery: { enabled: true, callback: () => executeQueryRef.current?.() } }}
          />
          {promptState.isOpen && editorRef.current && (
            <ResizableAIWidget
              editor={editorRef.current}
              id="ask-ai"
              value={promptInput}
              onChange={setPromptInput}
              onSubmit={(prompt: string) => {
                handlePrompt(prompt, {
                  beforeSelection: promptState.beforeSelection,
                  selection: promptState.selection,
                  afterSelection: promptState.afterSelection,
                })
              }}
              onCancel={handleReset}
              isDiffVisible={false}
              isLoading={isCompletionLoading}
              startLineNumber={Math.max(0, promptState.startLineNumber)}
              endLineNumber={promptState.endLineNumber}
            />
          )}
          <AnimatePresence>
            {!promptState.isOpen && !currentValue && aiEndpoint && (
              <motion.p
                initial={{ y: 5, opacity: 0 }}
                animate={{ y: 0, opacity: 1 }}
                exit={{ y: 5, opacity: 0 }}
                className="text-foreground-lighter absolute bottom-4 left-4 z-10 font-mono text-xs flex items-center gap-1"
              >
                Hit{' '}
                <KeyboardShortcut
                  keys={['Meta', 'Shift', 'k']}
                  variant="inline"
                  className="text-xs text-foreground-lighter"
                />{' '}
                to edit with the Assistant
              </motion.p>
            )}
          </AnimatePresence>
        </div>
      )}
    </div>
  )
}