/*
 * 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.prestosql.testing.datatype;

import com.google.common.base.Joiner;
import io.prestosql.testing.sql.SqlExecutor;
import io.prestosql.testing.sql.TestTable;

import java.util.List;
import java.util.stream.Stream;

import static java.lang.String.format;
import static java.util.stream.Collectors.toList;
import static java.util.stream.IntStream.range;

public class CreateAndInsertDataSetup
        implements DataSetup
{
    private final SqlExecutor sqlExecutor;
    private final String tableNamePrefix;

    public CreateAndInsertDataSetup(SqlExecutor sqlExecutor, String tableNamePrefix)
    {
        this.sqlExecutor = sqlExecutor;
        this.tableNamePrefix = tableNamePrefix;
    }

    @Override
    public TestTable setupTestTable(List<ColumnSetup> inputs)
    {
        TestTable testTable = createTestTable(inputs);
        try {
            insertRows(testTable, inputs);
        }
        catch (Exception e) {
            closeQuietly(testTable);
            throw e;
        }
        return testTable;
    }

    // TODO replace with closeAllSuppress when https://github.com/airlift/airlift/pull/889 is available
    private static void closeQuietly(AutoCloseable autoCloseable)
    {
        //noinspection EmptyTryBlock
        try (AutoCloseable ignore = autoCloseable) {
            // suppress empty try-catch warning
        }
        catch (Exception ignored) {
        }
    }

    private void insertRows(TestTable testTable, List<ColumnSetup> inputs)
    {
        Stream<String> literals = inputs.stream()
                .map(ColumnSetup::getInputLiteral);
        String valueLiterals = Joiner.on(", ").join(literals.iterator());
        sqlExecutor.execute(format("INSERT INTO %s VALUES(%s)", testTable.getName(), valueLiterals));
    }

    private TestTable createTestTable(List<ColumnSetup> inputs)
    {
        return new TestTable(sqlExecutor, tableNamePrefix, "(" + columnDefinitions(inputs) + ")");
    }

    private String columnDefinitions(List<ColumnSetup> inputs)
    {
        List<String> columnTypeDefinitions = inputs.stream()
                .map(input -> input.getDeclaredType().orElseThrow(() -> new IllegalArgumentException("declared type not set")))
                .collect(toList());
        Stream<String> columnDefinitions = range(0, columnTypeDefinitions.size())
                .mapToObj(i -> format("col_%d %s", i, columnTypeDefinitions.get(i)));
        return Joiner.on(",\n").join(columnDefinitions.iterator());
    }
}

@@ -741,24 +741,59 @@
 .util.List;%0A
+import java.util.stream.IntStream;%0A
 import java.
@@ -898,62 +898,15 @@
 ors.
-toList;%0Aimport static java.util.stream.IntStream.ra
+joini
 ng
-e
 ;%0A%0Ap
@@ -2456,20 +2456,13 @@
 ix, 
-%22(%22 + column
+table
 Defi
@@ -2471,17 +2471,16 @@
 tion
-s
 (inputs)
  + %22
@@ -2479,14 +2479,8 @@
 uts)
- + %22)%22
 );%0A 
@@ -2504,22 +2504,21 @@
  String 
-column
+table
 Definiti
@@ -2519,17 +2519,16 @@
 finition
-s
 (List%3CCo
@@ -2564,188 +2564,399 @@
     
-List%3CString%3E columnTypeDefinitions = inputs.stream()%0A                .map(input -%3E input.getDeclaredType().orElseThrow(() -%3E new IllegalArgumentException(%22declared type not set%22)))
+if (inputs.stream().allMatch(input -%3E input.getDeclaredType().isPresent())) %7B%0A            // When all types are explicitly specified, use ordinary CREATE TABLE%0A            return IntStream.range(0, inputs.size())%0A                    .mapToObj(column -%3E format(%22col_%25d %25s%22, column, inputs.get(column).getDeclaredType().orElseThrow()))%0A                    .collect(joining(%22,%5Cn%22, %22(%5Cn%22, %22)%22));
 %0A   
@@ -2956,32 +2956,35 @@
 ));%0A        
+%7D%0A%0A
         
 .collect(toL
@@ -2975,27 +2975,48 @@
     
-.collect(toList
+return IntStream.range(0, inputs.size
 ())
-;
 %0A   
@@ -3024,81 +3024,158 @@
     
-Stream%3CString%3E columnDefinitions = range(0, columnTypeDefinitions.size
+        .mapToObj(column -%3E %7B%0A                    ColumnSetup input = inputs.get(column);%0A                    if (input.getDeclaredType().isEmpty
 ())
+ %7B
 %0A   
@@ -3191,22 +3191,22 @@
     
-.mapToObj(i -%3E
+        return
  for
@@ -3214,52 +3214,204 @@
 at(%22
+%25s AS 
 col_%25d
- %25s
 %22, i
-, columnTypeDefinitions.get(i))
+nput.getInputLiteral(), column);%0A                    %7D%0A%0A                    return format(%22CAST(%25s AS %25s) AS col_%25d%22, input.getInputLiteral(), input.getDeclaredType().get(), column
 );%0A 
@@ -3421,64 +3421,103 @@
     
-return Joiner.on(%22,%5Cn%22).join(columnDefinitions.iterator(
+        %7D)%0A                .collect(joining(%22,%5Cn%22, %22AS%5CnSELECT%5Cn%22, %22%5CnWHERE 'with no' = 'data'%22
 ));%0A

Support implicit types in CreateAndInsertDataSetup