gemini-cli/packages/cli/src/ui/hooks/useCommandCompletion.tsx

/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import { useCallback, useMemo, useEffect, useState } from 'react';
import type { Suggestion } from '../components/SuggestionsDisplay.js';
import type { CommandContext, SlashCommand } from '../commands/types.js';
import type { TextBuffer } from '../components/shared/text-buffer.js';
import { logicalPosToOffset } from '../components/shared/text-buffer.js';
import { isSlashCommand } from '../utils/commandUtils.js';
import { toCodePoints } from '../utils/textUtils.js';
import { useAtCompletion } from './useAtCompletion.js';
import { useSlashCompletion } from './useSlashCompletion.js';
import { useShellCompletion } from './useShellCompletion.js';
import type { PromptCompletion } from './usePromptCompletion.js';
import {
  usePromptCompletion,
  PROMPT_COMPLETION_MIN_LENGTH,
} from './usePromptCompletion.js';
import type { Config } from '@google/gemini-cli-core';
import { useCompletion } from './useCompletion.js';

export enum CompletionMode {
  IDLE = 'IDLE',
  AT = 'AT',
  SLASH = 'SLASH',
  PROMPT = 'PROMPT',
  SHELL = 'SHELL',
}

export interface UseCommandCompletionReturn {
  suggestions: Suggestion[];
  activeSuggestionIndex: number;
  visibleStartIndex: number;
  showSuggestions: boolean;
  isLoadingSuggestions: boolean;
  isPerfectMatch: boolean;
  forceShowShellSuggestions: boolean;
  setForceShowShellSuggestions: (value: boolean) => void;
  isShellSuggestionsVisible: boolean;
  setActiveSuggestionIndex: React.Dispatch<React.SetStateAction<number>>;
  resetCompletionState: () => void;
  navigateUp: () => void;
  navigateDown: () => void;
  handleAutocomplete: (indexToUse: number) => void;
  promptCompletion: PromptCompletion;
  getCommandFromSuggestion: (
    suggestion: Suggestion,
  ) => SlashCommand | undefined;
  slashCompletionRange: {
    completionStart: number;
    completionEnd: number;
    getCommandFromSuggestion: (
      suggestion: Suggestion,
    ) => SlashCommand | undefined;
    isArgumentCompletion: boolean;
    leafCommand: SlashCommand | null;
  };
  getCompletedText: (suggestion: Suggestion) => string | null;
  completionMode: CompletionMode;
}

export interface UseCommandCompletionOptions {
  buffer: TextBuffer;
  cwd: string;
  slashCommands: readonly SlashCommand[];
  commandContext: CommandContext;
  reverseSearchActive?: boolean;
  shellModeActive: boolean;
  config?: Config;
  active: boolean;
}

export function useCommandCompletion({
  buffer,
  cwd,
  slashCommands,
  commandContext,
  reverseSearchActive = false,
  shellModeActive,
  config,
  active,
}: UseCommandCompletionOptions): UseCommandCompletionReturn {
  const [forceShowShellSuggestions, setForceShowShellSuggestions] =
    useState(false);

  const {
    suggestions,
    activeSuggestionIndex,
    visibleStartIndex,
    isLoadingSuggestions,
    isPerfectMatch,

    setSuggestions,
    setActiveSuggestionIndex,
    setIsLoadingSuggestions,
    setIsPerfectMatch,
    setVisibleStartIndex,

    resetCompletionState: baseResetCompletionState,
    navigateUp,
    navigateDown,
  } = useCompletion();

  const resetCompletionState = useCallback(() => {
    baseResetCompletionState();
    setForceShowShellSuggestions(false);
  }, [baseResetCompletionState]);

  const cursorRow = buffer.cursor[0];
  const cursorCol = buffer.cursor[1];

  const {
    completionMode,
    query: memoQuery,
    completionStart,
    completionEnd,
  } = useMemo(() => {
    const currentLine = buffer.lines[cursorRow] || '';
    const codePoints = toCodePoints(currentLine);

    if (shellModeActive) {
      return {
        completionMode:
          currentLine.trim().length === 0
            ? CompletionMode.IDLE
            : CompletionMode.SHELL,
        query: '',
        completionStart: -1,
        completionEnd: -1,
      };
    }

    // FIRST: Check for @ completion (scan backwards from cursor)
    // This must happen before slash command check so that `/cmd @file`
    // triggers file completion, not just slash command completion.
    for (let i = cursorCol - 1; i >= 0; i--) {
      const char = codePoints[i];

      if (char === ' ') {
        let backslashCount = 0;
        for (let j = i - 1; j >= 0 && codePoints[j] === '\\'; j--) {
          backslashCount++;
        }
        if (backslashCount % 2 === 0) {
          break;
        }
      } else if (char === '@') {
        let end = codePoints.length;
        for (let i = cursorCol; i < codePoints.length; i++) {
          if (codePoints[i] === ' ') {
            let backslashCount = 0;
            for (let j = i - 1; j >= 0 && codePoints[j] === '\\'; j--) {
              backslashCount++;
            }

            if (backslashCount % 2 === 0) {
              end = i;
              break;
            }
          }
        }
        const pathStart = i + 1;
        const partialPath = currentLine.substring(pathStart, end);
        return {
          completionMode: CompletionMode.AT,
          query: partialPath,
          completionStart: pathStart,
          completionEnd: end,
        };
      }
    }

    // THEN: Check for slash command (only if no @ completion is active)
    if (cursorRow === 0 && isSlashCommand(currentLine.trim())) {
      return {
        completionMode: CompletionMode.SLASH,
        query: currentLine,
        completionStart: 0,
        completionEnd: currentLine.length,
      };
    }

    // Check for prompt completion - only if enabled
    const trimmedText = buffer.text.trim();
    const isPromptCompletionEnabled = false;
    if (
      isPromptCompletionEnabled &&
      trimmedText.length >= PROMPT_COMPLETION_MIN_LENGTH &&
      !isSlashCommand(trimmedText) &&
      !trimmedText.includes('@')
    ) {
      return {
        completionMode: CompletionMode.PROMPT,
        query: trimmedText,
        completionStart: 0,
        completionEnd: trimmedText.length,
      };
    }

    return {
      completionMode: CompletionMode.IDLE,
      query: null,
      completionStart: -1,
      completionEnd: -1,
    };
  }, [cursorRow, cursorCol, buffer.lines, buffer.text, shellModeActive]);

  useAtCompletion({
    enabled: active && completionMode === CompletionMode.AT,
    pattern: memoQuery || '',
    config,
    cwd,
    setSuggestions,
    setIsLoadingSuggestions,
  });

  const slashCompletionRange = useSlashCompletion({
    enabled:
      active && completionMode === CompletionMode.SLASH && !shellModeActive,
    query: memoQuery,
    slashCommands,
    commandContext,
    setSuggestions,
    setIsLoadingSuggestions,
    setIsPerfectMatch,
  });

  const shellCompletionRange = useShellCompletion({
    enabled: active && completionMode === CompletionMode.SHELL,
    line: buffer.lines[cursorRow] || '',
    cursorCol,
    cwd,
    setSuggestions,
    setIsLoadingSuggestions,
  });

  const query =
    completionMode === CompletionMode.SHELL
      ? shellCompletionRange.query
      : memoQuery;

  const basePromptCompletion = usePromptCompletion({
    buffer,
  });

  const isShellSuggestionsVisible =
    completionMode !== CompletionMode.SHELL || forceShowShellSuggestions;

  const promptCompletion = useMemo(() => {
    if (
      completionMode === CompletionMode.SHELL &&
      suggestions.length === 1 &&
      query != null &&
      shellCompletionRange.completionStart === shellCompletionRange.activeStart
    ) {
      const suggestion = suggestions[0];
      const textToInsertBase = suggestion.value;

      if (
        textToInsertBase.startsWith(query) &&
        textToInsertBase.length > query.length
      ) {
        const currentLine = buffer.lines[cursorRow] || '';
        const start = shellCompletionRange.completionStart;
        const end = shellCompletionRange.completionEnd;

        let textToInsert = textToInsertBase;
        const charAfterCompletion = currentLine[end];
        if (
          charAfterCompletion !== ' ' &&
          !textToInsert.endsWith('/') &&
          !textToInsert.endsWith('\\')
        ) {
          textToInsert += ' ';
        }

        const newText =
          currentLine.substring(0, start) +
          textToInsert +
          currentLine.substring(end);

        return {
          text: newText,
          isActive: true,
          isLoading: false,
          accept: () => {
            buffer.replaceRangeByOffset(
              logicalPosToOffset(buffer.lines, cursorRow, start),
              logicalPosToOffset(buffer.lines, cursorRow, end),
              textToInsert,
            );
          },
          clear: () => {},
          markSelected: () => {},
        };
      }
    }
    return basePromptCompletion;
  }, [
    completionMode,
    suggestions,
    query,
    basePromptCompletion,
    buffer,
    cursorRow,
    shellCompletionRange,
  ]);

  useEffect(() => {
    setActiveSuggestionIndex(suggestions.length > 0 ? 0 : -1);
    setVisibleStartIndex(0);

    // Generic perfect match detection for non-slash modes or as a fallback
    if (completionMode !== CompletionMode.SLASH) {
      if (suggestions.length > 0) {
        const firstSuggestion = suggestions[0];
        setIsPerfectMatch(firstSuggestion.value === query);
      } else {
        setIsPerfectMatch(false);
      }
    }
  }, [
    suggestions,
    setActiveSuggestionIndex,
    setVisibleStartIndex,
    completionMode,
    query,
    setIsPerfectMatch,
  ]);

  useEffect(() => {
    if (
      !active ||
      completionMode === CompletionMode.IDLE ||
      reverseSearchActive
    ) {
      resetCompletionState();
    }
  }, [active, completionMode, reverseSearchActive, resetCompletionState]);

  const showSuggestions =
    active &&
    completionMode !== CompletionMode.IDLE &&
    !reverseSearchActive &&
    isShellSuggestionsVisible &&
    (isLoadingSuggestions || suggestions.length > 0);

  /**
   * Gets the completed text by replacing the completion range with the suggestion value.
   * This is the core string replacement logic used by both autocomplete and auto-execute.
   *
   * @param suggestion The suggestion to apply
   * @returns The completed text with the suggestion applied, or null if invalid
   */
  const getCompletedText = useCallback(
    (suggestion: Suggestion): string | null => {
      const currentLine = buffer.lines[cursorRow] || '';

      let start = completionStart;
      let end = completionEnd;
      if (completionMode === CompletionMode.SLASH) {
        start = slashCompletionRange.completionStart;
        end = slashCompletionRange.completionEnd;
      } else if (completionMode === CompletionMode.SHELL) {
        start = shellCompletionRange.completionStart;
        end = shellCompletionRange.completionEnd;
      }

      if (start === -1 || end === -1) {
        return null;
      }

      // Apply space padding for slash commands (needed for subcommands like "/chat list")
      let suggestionText = suggestion.value;
      if (completionMode === CompletionMode.SLASH) {
        // Add leading space if completing a subcommand (cursor is after parent command with no space)
        if (start === end && start > 1 && currentLine[start - 1] !== ' ') {
          suggestionText = ' ' + suggestionText;
        }
      }

      // Build the completed text with proper spacing
      return (
        currentLine.substring(0, start) +
        suggestionText +
        currentLine.substring(end)
      );
    },
    [
      cursorRow,
      buffer.lines,
      completionMode,
      completionStart,
      completionEnd,
      slashCompletionRange,
      shellCompletionRange,
    ],
  );

  const handleAutocomplete = useCallback(
    (indexToUse: number) => {
      if (indexToUse < 0 || indexToUse >= suggestions.length) {
        return;
      }
      const suggestion = suggestions[indexToUse];
      const completedText = getCompletedText(suggestion);

      if (completedText === null) {
        return;
      }

      let start = completionStart;
      let end = completionEnd;
      if (completionMode === CompletionMode.SLASH) {
        start = slashCompletionRange.completionStart;
        end = slashCompletionRange.completionEnd;
      } else if (completionMode === CompletionMode.SHELL) {
        start = shellCompletionRange.completionStart;
        end = shellCompletionRange.completionEnd;
      }

      // Add space padding for Tab completion (auto-execute gets padding from getCompletedText)
      let suggestionText = suggestion.value;
      if (completionMode === CompletionMode.SLASH) {
        if (
          start === end &&
          start > 1 &&
          (buffer.lines[cursorRow] || '')[start - 1] !== ' '
        ) {
          suggestionText = ' ' + suggestionText;
        }
      }

      const lineCodePoints = toCodePoints(buffer.lines[cursorRow] || '');
      const charAfterCompletion = lineCodePoints[end];
      if (
        charAfterCompletion !== ' ' &&
        !suggestionText.endsWith('/') &&
        !suggestionText.endsWith('\\')
      ) {
        suggestionText += ' ';
      }

      buffer.replaceRangeByOffset(
        logicalPosToOffset(buffer.lines, cursorRow, start),
        logicalPosToOffset(buffer.lines, cursorRow, end),
        suggestionText,
      );
    },
    [
      cursorRow,
      buffer,
      suggestions,
      completionMode,
      completionStart,
      completionEnd,
      slashCompletionRange,
      shellCompletionRange,
      getCompletedText,
    ],
  );

  return {
    suggestions,
    activeSuggestionIndex,
    visibleStartIndex,
    showSuggestions,
    isLoadingSuggestions,
    isPerfectMatch,
    forceShowShellSuggestions,
    setForceShowShellSuggestions,
    isShellSuggestionsVisible,
    setActiveSuggestionIndex,
    resetCompletionState,
    navigateUp,
    navigateDown,
    handleAutocomplete,
    promptCompletion,
    getCommandFromSuggestion: slashCompletionRange.getCommandFromSuggestion,
    slashCompletionRange,
    getCompletedText,
    completionMode,
  };
}