/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package io.trino.cli;

import com.google.common.collect.ImmutableSet;
import org.jline.reader.Completer;
import org.jline.reader.impl.completer.StringsCompleter;

import java.util.Set;

import static java.util.Locale.ENGLISH;
import static java.util.stream.Collectors.toSet;

public final class Completion
{
    private static final Set<String> COMMANDS = ImmutableSet.of(
            "SELECT",
            "SHOW CATALOGS",
            "SHOW COLUMNS",
            "SHOW FUNCTIONS",
            "SHOW SCHEMAS",
            "SHOW SESSION",
            "SHOW TABLES",
            "CREATE TABLE",
            "DROP TABLE",
            "EXPLAIN",
            "DESCRIBE",
            "USE",
            "HELP",
            "QUIT");

    private Completion() {}

    public static Completer commandCompleter()
    {
        return new StringsCompleter(ImmutableSet.<String>builder()
                .addAll(COMMANDS)
                .addAll(COMMANDS.stream()
                        .map(value -> value.toLowerCase(ENGLISH))
                        .collect(toSet()))
                .build());
    }
}

@@ -589,46 +589,151 @@
 ort 
-com.google.common.collect.ImmutableSet
+org.jline.reader.Completer;%0Aimport org.jline.reader.impl.completer.AggregateCompleter;%0Aimport org.jline.reader.impl.completer.ArgumentCompleter
 ;%0Aim
@@ -746,32 +746,51 @@
 rg.jline.reader.
+impl.completer.Null
 Completer;%0Aimpor
@@ -859,18 +859,19 @@
 va.util.
-Se
+Lis
 t;%0A%0Aimpo
@@ -894,22 +894,21 @@
 til.
-Locale.ENGLISH
+Arrays.asList
 ;%0Aim
@@ -933,30 +933,32 @@
 til.
-stream.Collectors.toSe
+Collections.singletonLis
 t;%0A%0A
@@ -1005,73 +1005,86 @@
 ate 
-static final Set%3CString%3E COMMANDS = ImmutableSet.of(
+Completion() %7B%7D%0A%0A    public static Completer commandCompleter()
 %0A    
+%7B%0A
         
 %22SEL
@@ -1083,186 +1083,365 @@
     
-%22SELECT%22,%0A            %22SHOW CATALOGS%22,%0A            %22SHOW COLUMNS%22,%0A            %22SHOW FUNCTIONS%22,%0A            %22SHOW SCHEMAS%22,%0A            %22SHOW SESSION%22,%0A            %22SHOW TABLES%22
+return new AggregateCompleter(buildArgumentCompleter(%22CREATE%22, singletonList(%22TABLE%22)),%0A                buildArgumentCompleter(%22DESCRIBE%22),%0A                buildArgumentCompleter(%22DROP%22, singletonList(%22TABLE%22)),%0A                buildArgumentCompleter(%22EXPLAIN%22),%0A                buildArgumentCompleter(%22HELP%22),%0A                buildArgumentCompleter(%22QUIT%22)
 ,%0A  
@@ -1454,192 +1454,230 @@
     
+ 
-%22CREATE TABLE%22,%0A            %22DROP TABLE%22,%0A            %22EXPLAIN%22,%0A            %22DESCRIBE%22,%0A           
+   buildArgumentCompleter(%22SELECT%22),%0A                buildArgumentCompleter(%22SHOW%22, asList(%22CATALOGS%22, %22COLUMNS%22, %22FUNCTIONS%22, %22SCHEMAS%22,
  %22
-U
 SE
-%22,%0A            %22HELP%22,%0A            %22QUIT%22);%0A%0A    private Completion() %7B
+SSION%22, %22TABLES%22)),%0A                buildArgumentCompleter(%22USE%22));%0A    
 %7D%0A%0A    p
 ubli
@@ -1664,37 +1664,38 @@
 ));%0A    %7D%0A%0A    p
-ublic
+rivate
  static Complete
@@ -1700,264 +1700,367 @@
 ter 
-commandCompleter()%0A    %7B%0A        return new StringsCompleter(ImmutableSet.%3CString%3Ebuilder()%0A                .addAll(COMMANDS)%0A                .addAll(COMMANDS.stream()%0A                        .map(value -%3E value.toLowerCase(ENGLISH))%0A          
+buildArgumentCompleter(String command)%0A    %7B%0A        // NullCompleter is used to indicate the command is complete and the last word should not be repeated%0A        return new ArgumentCompleter(new StringsCompleter(command), NullCompleter.INSTANCE);%0A    %7D%0A%0A    private static Completer buildArgumentCompleter(String command, List%3CString%3E options)%0A
     
+%7B%0A
         
   .c
@@ -2059,53 +2059,121 @@
     
-  .collect(toSet()))%0A                .build()
+return new ArgumentCompleter(new StringsCompleter(command), new StringsCompleter(options), NullCompleter.INSTANCE
 );%0A 

Fix cli completion for compound commands