llm-in-text/src/plugins/copilotPlugin.ts

import { Plugin, PluginKey, Selection } from '@milkdown/prose/state'
import { $prose, $ctx, $markSchema } from '@milkdown/kit/utils'
import type { Ctx } from '@milkdown/kit/core'
import type { EditorView } from '@milkdown/prose/view'

const COPILOT_PLUGIN_KEY = new PluginKey('milkdown-copilot')
const DEBOUNCE_MS = 500

let enabled = true

interface CopilotState {
  from: number
  to: number
  suggestion: string
}

interface CopilotConfig {
  fetchSuggestion: (prefix: string, suffix: string, signal?: AbortSignal) => Promise<string>
  debounceMs?: number
}

const initialState: CopilotState = {
  from: 0,
  to: 0,
  suggestion: ''
}

export const copilotConfigCtx = $ctx<CopilotConfig, 'copilotConfig'>({
  fetchSuggestion: async () => '',
  debounceMs: DEBOUNCE_MS
}, 'copilotConfig')

export const copilotGhostMark = $markSchema('copilot_ghost', () => ({
  excludes: '_',
  inclusive: true,
  parseDOM: [{ tag: 'span[data-copilot-ghost]' }],
  toDOM: () => ['span', { 'data-copilot-ghost': '', class: 'copilot-ghost-text' }, 0]
}))

let debounceTimer: ReturnType<typeof setTimeout> | null = null
let abortController: AbortController | null = null
let currentCtx: Ctx | null = null

function clearGhostText(view: EditorView) {
  const state = COPILOT_PLUGIN_KEY.getState(view.state)
  if (state && state.suggestion && state.from < state.to) {
    const tr = view.state.tr
      .delete(state.from, state.to)
      .setMeta(COPILOT_PLUGIN_KEY, { ...initialState })
    view.dispatch(tr)
  }
}

function insertGhostText(view: EditorView, suggestion: string, from: number) {
  if (!currentCtx || !suggestion) return
  
  const schema = view.state.schema
  const markType = schema.marks.copilot_ghost
  
  if (!markType) {
    console.error('[Copilot] copilot_ghost mark not found in schema')
    return
  }
  
  const tr = view.state.tr
  tr.insertText(suggestion, from)
  const endPos = from + suggestion.length
  tr.addMark(from, endPos, markType.create())
  tr.setMeta(COPILOT_PLUGIN_KEY, { from, to: endPos, suggestion })
  view.dispatch(tr)
}

function doFetchSuggestion(view: EditorView, pos: number, prefix: string, suffix: string) {
  if (!currentCtx) return
  
  const config = currentCtx.get(copilotConfigCtx.key)
  
  if (abortController) {
    abortController.abort()
    abortController = null
  }
  
  abortController = new AbortController()
  
  config.fetchSuggestion(prefix, suffix, abortController.signal)
    .then(suggestion => {
      if (view.state.selection.from !== pos) return
      
      if (suggestion) {
        insertGhostText(view, suggestion, pos)
      }
    })
    .catch(e => {
      if (e.name !== 'AbortError') {
        console.error('[Copilot] Error:', e)
      }
    })
    .finally(() => {
      abortController = null
    })
}

function scheduleFetch(view: EditorView, pos: number, prefix: string, suffix: string) {
  if (!enabled) return
  
  if (debounceTimer) {
    clearTimeout(debounceTimer)
    debounceTimer = null
  }
  
  debounceTimer = setTimeout(() => {
    debounceTimer = null
    doFetchSuggestion(view, pos, prefix, suffix)
  }, DEBOUNCE_MS)
}

function acceptSuggestion(view: EditorView) {
  const state = COPILOT_PLUGIN_KEY.getState(view.state)
  if (!state?.suggestion || state.from >= state.to) return false
  
  const tr = view.state.tr.removeMark(state.from, state.to, view.state.schema.marks.copilot_ghost)
  const endPos = Math.min(state.to, tr.doc.content.size)
  tr.setSelection(Selection.near(tr.doc.resolve(endPos)))
  tr.setMeta(COPILOT_PLUGIN_KEY, { ...initialState })
  view.dispatch(tr)
  return true
}

function rejectSuggestion(view: EditorView) {
  const state = COPILOT_PLUGIN_KEY.getState(view.state)
  if (!state?.suggestion) return false
  
  clearGhostText(view)
  return true
}

export const copilotPlugin = $prose((ctx) => {
  currentCtx = ctx
  
  return new Plugin<CopilotState>({
    key: COPILOT_PLUGIN_KEY,
    state: {
      init: () => ({ ...initialState }),
      apply: (tr, value) => {
        const meta = tr.getMeta(COPILOT_PLUGIN_KEY)
        if (meta !== undefined) {
          return meta
        }
        
        if (tr.docChanged && value.suggestion) {
          return { ...initialState }
        }
        
        return value
      }
    },
    props: {
      handleKeyDown: (view, event) => {
        const state = COPILOT_PLUGIN_KEY.getState(view.state)
        
        if (event.key === 'Tab' && state?.suggestion) {
          event.preventDefault()
          return acceptSuggestion(view)
        }
        
        if (event.key === 'Escape' && state?.suggestion) {
          event.preventDefault()
          return rejectSuggestion(view)
        }
        
        if (state?.suggestion && event.key !== 'Shift' && event.key !== 'Control' && event.key !== 'Alt' && event.key !== 'Meta') {
          clearGhostText(view)
        }
        
        return false
      },
      handleClick: (view, pos) => {
        const state = COPILOT_PLUGIN_KEY.getState(view.state)
        if (!state?.suggestion) return false
        
        if (pos >= state.from && pos < state.to) {
          return acceptSuggestion(view)
        }
        
        clearGhostText(view)
        return false
      }
    },
    view: () => ({
      update: (view, prevState) => {
        if (view.state.doc.eq(prevState.doc) && view.state.selection.eq(prevState.selection)) {
          return
        }
        
        const state = COPILOT_PLUGIN_KEY.getState(view.state)
        if (state?.suggestion) {
          return
        }
        
        if (!view.state.doc.eq(prevState.doc)) {
          const { from, to } = view.state.selection
          if (from !== to) return
          
          const doc = view.state.doc
          const prefix = doc.textBetween(0, from)
          const suffix = doc.textBetween(to, doc.content.size)
          
          scheduleFetch(view, from, prefix, suffix)
        }
      }
    })
  })
})

export { COPILOT_PLUGIN_KEY }

export function isCopilotEnabled(): boolean {
  return enabled
}

export function setCopilotEnabled(value: boolean): void {
  enabled = value
  
  if (!value) {
    if (debounceTimer) {
      clearTimeout(debounceTimer)
      debounceTimer = null
    }
    if (abortController) {
      abortController.abort()
      abortController = null
    }
  }
}