notebook-ai-extension/src/components/opencode_inline_prompt.ts

/**
 * Inline prompt widget attached to a cell's DOM when the AI button is clicked.
 * Renders two <select> boxes (provider + model), a textarea, send/cancel
 * buttons, and a scrollable history area that shows the current session's
 * messages (user + assistant). The model select's default for the chosen
 * provider is `default[providerID]` from the /config/providers response.
 *
 * v3-final: no settings — picks come from the OpenCode Serve providers
 * cache; history comes from /session-messages; cell source is NOT replaced
 * (the AI reply is rendered as markdown inside this history area).
 */
import { CodeCell } from '@jupyterlab/cells';
import { Notification } from '@jupyterlab/apputils';
import { Widget } from '@lumino/widgets';
import { marked } from 'marked';

import type { OpenCodeMessage, OpenCodeProvidersResponse } from '../types';

export interface IOpenCodeInlinePromptOptions {
  onSubmit: (text: string, providerId?: string, modelId?: string) => void;
  onCancel: () => void;
  disabled: boolean;
  providers: OpenCodeProvidersResponse | null;
}

interface FlatProvider {
  id: string;
  name?: string;
  models: { [modelID: string]: unknown };
}

function flattenProviders(
  providers: OpenCodeProvidersResponse | null
): { providers: FlatProvider[]; defaultMap: { [pid: string]: string } } {
  const out: FlatProvider[] = [];
  if (!providers || !providers.providers) {
    return { providers: out, defaultMap: (providers && providers.default) || {} };
  }
  for (const p of providers.providers) {
    if (!p.models || typeof p.models !== 'object') {
      continue;
    }
    out.push({ id: p.id, name: p.name, models: p.models });
  }
  return { providers: out, defaultMap: providers.default || {} };
}

function modelKeys(models: { [modelID: string]: unknown }): string[] {
  return Object.keys(models);
}

function pickDefaultModelId(
  providerId: string,
  models: { [modelID: string]: unknown },
  defaultMap: { [pid: string]: string }
): string | undefined {
  const keys = modelKeys(models);
  if (keys.length === 0) {
    return undefined;
  }
  const fromDefault = defaultMap[providerId];
  if (fromDefault && Object.prototype.hasOwnProperty.call(models, fromDefault)) {
    return fromDefault;
  }
  return keys[0];
}

export class OpenCodeInlinePrompt extends Widget {
  private _cell: CodeCell;
  private _textarea: HTMLTextAreaElement;
  private _sendBtn: HTMLButtonElement;
  private _cancelBtn: HTMLButtonElement;
  private _providerSelect: HTMLSelectElement | null = null;
  private _modelSelect: HTMLSelectElement | null = null;
  private _providerModels: FlatProvider[];
  private _defaultMap: { [pid: string]: string };
  private _historyEl: HTMLDivElement;

  constructor(
    cell: CodeCell,
    options: IOpenCodeInlinePromptOptions
  ) {
    super();
    this._cell = cell;
    this.addClass('opencode-inline-prompt');

    const flat = flattenProviders(options.providers);
    this._providerModels = flat.providers;
    this._defaultMap = flat.defaultMap;

    if (this._providerModels.length > 0) {
      this._providerSelect = document.createElement('select');
      this._providerSelect.className = 'opencode-provider-select';
      for (const p of this._providerModels) {
        const o = document.createElement('option');
        o.value = p.id;
        o.textContent = p.name ? `${p.id} (${p.name})` : p.id;
        this._providerSelect.appendChild(o);
      }
      this._providerSelect.value = this._providerModels[0].id;
      this._providerSelect.addEventListener('change', () => {
        this._rebuildModelSelect(this._providerSelect!.value);
      });
    }

    if (this._providerModels.length > 0) {
      this._modelSelect = document.createElement('select');
      this._modelSelect.className = 'opencode-model-select';
      this._rebuildModelSelect(this._providerSelect?.value);
    }

    this._textarea = document.createElement('textarea');
    this._textarea.placeholder = '向 AI 描述你想要的修改…';
    this._textarea.rows = 3;
    this._sendBtn = document.createElement('button');
    this._sendBtn.className = 'opencode-btn-send';
    this._sendBtn.textContent = '发送';
    this._sendBtn.disabled = options.disabled;
    this._sendBtn.addEventListener('click', () => {
      const text = this._textarea.value;
      if (!text.trim()) {
        return;
      }
      const providerId = this._providerSelect?.value;
      const modelId = this._modelSelect?.value || undefined;
      options.onSubmit(text, providerId, modelId);
    });
    // Chat-input convention: Enter submits, Shift+Enter / Ctrl+Enter
    // insert a newline (textarea default).
    this._textarea.addEventListener('keydown', (e) => {
      if (e.key === 'Enter' && !e.shiftKey && !e.ctrlKey) {
        e.preventDefault();
        this._sendBtn.click();
      }
    });
    this._cancelBtn = document.createElement('button');
    this._cancelBtn.className = 'opencode-btn-cancel';
    this._cancelBtn.textContent = '取消';
    this._cancelBtn.addEventListener('click', () => {
      options.onCancel();
    });

    const actions = document.createElement('div');
    actions.className = 'opencode-inline-actions';
    actions.appendChild(this._sendBtn);
    actions.appendChild(this._cancelBtn);

    if (this._providerSelect || this._modelSelect) {
      const row = document.createElement('div');
      row.className = 'opencode-prompt-providers';
      if (this._providerSelect) {
        const pLabel = document.createElement('label');
        const pSpan = document.createElement('span');
        pSpan.textContent = 'Provider:';
        pLabel.appendChild(pSpan);
        pLabel.appendChild(this._providerSelect);
        row.appendChild(pLabel);
      }
      if (this._modelSelect) {
        const mLabel = document.createElement('label');
        const mSpan = document.createElement('span');
        mSpan.textContent = 'Model:';
        mLabel.appendChild(mSpan);
        mLabel.appendChild(this._modelSelect);
        row.appendChild(mLabel);
      }
      this.node.appendChild(row);
    }

    // Scrollable history of the current session's messages. Empty until
    // setMessages() is called by the owning cell action.
    this._historyEl = document.createElement('div');
    this._historyEl.className = 'opencode-inline-history';
    this.node.appendChild(this._historyEl);

    this.node.appendChild(this._textarea);
    this.node.appendChild(actions);
  }

  private _rebuildModelSelect(providerId: string | undefined): void {
    if (!this._modelSelect) {
      return;
    }
    while (this._modelSelect.firstChild) {
      this._modelSelect.removeChild(this._modelSelect.firstChild);
    }
    if (!providerId) {
      this._modelSelect.disabled = true;
      return;
    }
    const provider = this._providerModels.find(p => p.id === providerId);
    if (!provider) {
      this._modelSelect.disabled = true;
      return;
    }
    const keys = modelKeys(provider.models);
    if (keys.length === 0) {
      this._modelSelect.disabled = true;
      return;
    }
    for (const k of keys) {
      const o = document.createElement('option');
      o.value = k;
      o.textContent = k;
      this._modelSelect.appendChild(o);
    }
    this._modelSelect.disabled = false;
    this._modelSelect.value =
      pickDefaultModelId(providerId, provider.models, this._defaultMap) ?? keys[0];
  }

  setDisabled(disabled: boolean): void {
    this._sendBtn.disabled = disabled;
    if (this._providerSelect) {
      this._providerSelect.disabled = disabled;
    }
    if (this._modelSelect) {
      this._modelSelect.disabled = disabled;
    }
  }

  /**
   * Clear the user input textarea (called by the cell action after a
   * successful response, so the user can type a follow-up message
   * without manually deleting the previous prompt).
   */
  clearInput(): void {
    this._textarea.value = '';
  }

  /**
   * Render the current session's messages into the scrollable history
   * area. User messages are plain text; assistant messages are rendered
   * as markdown via marked. Auto-scrolls to the bottom so the most
   * recent message is visible.
   */
  setMessages(messages: OpenCodeMessage[]): void {
    this._historyEl.textContent = '';
    for (const msg of messages) {
      const wrap = document.createElement('div');
      wrap.className = `opencode-msg opencode-msg-${msg.role}`;
      if (msg.role === 'user') {
        wrap.textContent = msg.content;
      } else {
        // Assistant: render the markdown, then for every fenced code
        // block (<pre>) attach a small toolbar with Copy / Insert /
        // Replace buttons that act on THAT block's code (not the whole
        // message). Messages without code fences get no toolbar.
        const rendered = document.createElement('div');
        rendered.className = 'opencode-msg-content';
        rendered.innerHTML = marked.parse(msg.content) as string;

        const codeBlocks = rendered.querySelectorAll('pre');
        codeBlocks.forEach(pre => {
          const code = pre.querySelector('code');
          if (!code) {
            return;
          }
          // Snapshot the code text before we mutate the DOM so the
          // buttons act on the exact code inside this block (no
          // surrounding ```fences```, no toolbar label text).
          const codeText = code.textContent ?? '';
          const wrap = document.createElement('div');
          wrap.className = 'opencode-code-block';
          const toolbar = document.createElement('div');
          toolbar.className = 'opencode-code-toolbar';
          const mkBtn = (
            label: string,
            title: string,
            onClick: () => void
          ) => {
            const b = document.createElement('button');
            b.type = 'button';
            b.className = 'opencode-code-btn';
            b.textContent = label;
            b.title = title;
            b.addEventListener('click', onClick);
            return b;
          };
          toolbar.appendChild(
            mkBtn('📋 复制', '复制代码', () => {
              void this._copyToClipboard(codeText);
            })
          );
          toolbar.appendChild(
            mkBtn('↪ 插入', '在光标位置插入', () => {
              this._insertAtCursor(codeText);
            })
          );
          toolbar.appendChild(
            mkBtn('⟳ 替换', '替换单元格内容', () => {
              this._replaceCell(codeText);
            })
          );
          wrap.appendChild(toolbar);
          // Move the original <pre> into the wrapper.
          pre.parentNode!.insertBefore(wrap, pre);
          wrap.appendChild(pre);
        });

        wrap.appendChild(rendered);
      }
      this._historyEl.appendChild(wrap);
    }
    this._historyEl.scrollTop = this._historyEl.scrollHeight;
  }

  /**
   * Copy the given text to the clipboard. Prefers navigator.clipboard
   * (available in secure contexts incl. JupyterLab). Falls back to a
   * hidden textarea + execCommand for older / non-secure contexts.
   */
  private async _copyToClipboard(text: string): Promise<void> {
    try {
      if (
        typeof navigator !== 'undefined' &&
        navigator.clipboard &&
        typeof navigator.clipboard.writeText === 'function'
      ) {
        await navigator.clipboard.writeText(text);
        Notification.info('已复制到剪贴板');
        return;
      }
      const ta = document.createElement('textarea');
      ta.value = text;
      ta.style.position = 'fixed';
      ta.style.opacity = '0';
      this.node.appendChild(ta);
      ta.select();
      document.execCommand('copy');
      ta.remove();
      Notification.info('已复制到剪贴板');
    } catch (e) {
      Notification.error('复制失败: ' + (e as Error).message);
    }
  }

  /**
   * Insert text at the current editor cursor position. If the cell has
   * no editor (e.g. a non-Code cell) or the cursor can't be read, falls
   * back to appending at the end of the cell.
   */
  private _insertAtCursor(text: string): void {
    const cell = this._cell;
    if (!cell || !cell.model || !cell.model.sharedModel) {
      return;
    }
    const sharedModel = cell.model.sharedModel;
    const source = sharedModel.getSource() as string;
    let offset = source.length;
    try {
      const editor = (cell as any).editor;
      const pos = editor && editor.getCursorPosition
        ? editor.getCursorPosition()
        : null;
      if (pos) {
        const lines = source.split('\n');
        const line = Math.max(0, Math.min(pos.line, lines.length - 1));
        const col = Math.max(
          0,
          Math.min(pos.column, (lines[line] ?? '').length)
        );
        offset = 0;
        for (let i = 0; i < line; i++) {
          offset += lines[i].length + 1;
        }
        offset += col;
      }
    } catch {
      // Fall back to append on any error reading the cursor.
      offset = source.length;
    }
    const newSource = source.slice(0, offset) + text + source.slice(offset);
    sharedModel.setSource(newSource);
    Notification.info('已插入到光标位置');
  }

  /** Replace the entire cell source with the given text. */
  private _replaceCell(text: string): void {
    const cell = this._cell;
    if (!cell || !cell.model || !cell.model.sharedModel) {
      return;
    }
    cell.model.sharedModel.setSource(text);
    Notification.info('已替换单元格内容');
  }
}