From 58dc43963d8d31edb467adc853cfeb46ba94a6f1 Mon Sep 17 00:00:00 2001
From: acheron <acheroncrypto@gmail.com>
Date: Sun, 19 May 2024 10:08:40 +0200
Subject: [PATCH] Add ability to temporarily change autocomplete handlers

---
 client/src/utils/pg/terminal/autocomplete.ts | 22 +++++++++++
 client/src/utils/pg/terminal/terminal.ts     | 40 ++++++++++++++++----
 2 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/client/src/utils/pg/terminal/autocomplete.ts b/client/src/utils/pg/terminal/autocomplete.ts
index 402baeb3..06b27c1b 100644
--- a/client/src/utils/pg/terminal/autocomplete.ts
+++ b/client/src/utils/pg/terminal/autocomplete.ts
@@ -8,6 +8,7 @@ type AutoCompleteHandler = (tokens: string[], index: number) => string[];
 
 /** Terminal autocomplete functionality */
 export class PgAutocomplete {
+  /** Autocomplete handlers */
   private _handlers: AutoCompleteHandler[];
 
   constructor(handlers: AutoCompleteHandler[]) {
@@ -23,6 +24,27 @@ export class PgAutocomplete {
     return this._handlers.length > 0;
   }
 
+  /**
+   * Temporarily set autocomplete handlers to the given handler.
+   *
+   * @param handler handler to set
+   * @param opts handler options:
+   * - `append`: whether to append the handler to the existing handlers
+   * @returns an object with `restore` callback to restore the handlers
+   */
+  temporarilySetHandlers(
+    handler: AutoCompleteHandler,
+    opts?: { append?: boolean }
+  ) {
+    const initialHandlers = this._handlers;
+    this._handlers = opts?.append ? [...initialHandlers, handler] : [handler];
+    return {
+      restore: () => {
+        this._handlers = initialHandlers;
+      },
+    };
+  }
+
   /**
    * Collect the autocomplete canditates from the given input.
    *
diff --git a/client/src/utils/pg/terminal/terminal.ts b/client/src/utils/pg/terminal/terminal.ts
index 0883d23b..76875562 100644
--- a/client/src/utils/pg/terminal/terminal.ts
+++ b/client/src/utils/pg/terminal/terminal.ts
@@ -326,6 +326,7 @@ export class PgTerm {
   private _fitAddon: FitAddon;
   private _tty: PgTty;
   private _shell: PgShell;
+  private _autocomplete: PgAutocomplete;
   private _isOpen: boolean;
 
   constructor(cmdManager: CommandManager, xtermOptions?: ITerminalOptions) {
@@ -342,12 +343,17 @@ export class PgTerm {
 
     // Create Shell and TTY
     const history = new PgHistory(20);
-    const autocomplete = new PgAutocomplete([
+    this._autocomplete = new PgAutocomplete([
       () => history.getEntries(),
       () => cmdManager.getNames(),
     ]);
-    this._tty = new PgTty(this._xterm, cmdManager, autocomplete);
-    this._shell = new PgShell(this._tty, cmdManager, autocomplete, history);
+    this._tty = new PgTty(this._xterm, cmdManager, this._autocomplete);
+    this._shell = new PgShell(
+      this._tty,
+      cmdManager,
+      this._autocomplete,
+      history
+    );
 
     // Add a custom resize handler that clears the prompt using the previous
     // configuration, updates the cached terminal size information and then
@@ -563,24 +569,42 @@ export class PgTerm {
     this.focus();
 
     let convertedMsg = msg;
+    let restore;
     if (opts?.default) {
-      convertedMsg += ` (default: ${opts.default})`;
+      const value = opts.default;
+      convertedMsg += ` (default: ${value})`;
+      restore = this._autocomplete.temporarilySetHandlers(() => [value], {
+        append: true,
+      }).restore;
     }
     if (opts?.choice) {
       // Show multi choice items
-      convertedMsg += opts.choice.items.reduce(
+      const items = opts.choice.items;
+      convertedMsg += items.reduce(
         (acc, cur, i) => acc + `\n[${i}] - ${cur}`,
         "\n"
       );
+      restore = this._autocomplete.temporarilySetHandlers(() =>
+        items.map((_, i) => i.toString())
+      ).restore;
     } else if (opts?.confirm) {
       convertedMsg += PgTerminal.secondaryText(` [yes/no]`);
+      restore = this._autocomplete.temporarilySetHandlers(() => [
+        "yes",
+        "no",
+      ]).restore;
     }
 
-    let userInput = await this._shell.waitForUserInput(convertedMsg);
-    if (!userInput && opts?.default) {
-      userInput = opts.default;
+    let userInput;
+    try {
+      userInput = await this._shell.waitForUserInput(convertedMsg);
+    } finally {
+      restore?.();
     }
 
+    // Set the input to the default if it exists on empty input
+    if (!userInput && opts?.default) userInput = opts.default;
+
     // Default validators
     if (opts && !opts.validator) {
       // Validate confirm