package sqlancer.postgres.oracle.tlp;

import java.sql.SQLException;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

import sqlancer.Randomly;
import sqlancer.TernaryLogicPartitioningOracleBase;
import sqlancer.TestOracle;
import sqlancer.gen.ExpressionGenerator;
import sqlancer.postgres.PostgresGlobalState;
import sqlancer.postgres.PostgresSchema;
import sqlancer.postgres.PostgresSchema.PostgresTable;
import sqlancer.postgres.PostgresSchema.PostgresTables;
import sqlancer.postgres.ast.PostgresColumnValue;
import sqlancer.postgres.ast.PostgresExpression;
import sqlancer.postgres.ast.PostgresJoin;
import sqlancer.postgres.ast.PostgresSelect;
import sqlancer.postgres.ast.PostgresSelect.ForClause;
import sqlancer.postgres.ast.PostgresSelect.PostgresFromTable;
import sqlancer.postgres.gen.PostgresCommon;
import sqlancer.postgres.gen.PostgresExpressionGenerator;
import sqlancer.postgres.oracle.PostgresNoRECOracle;

public class PostgresTLPBase extends TernaryLogicPartitioningOracleBase<PostgresExpression, PostgresGlobalState>
        implements TestOracle {

    PostgresSchema s;
    PostgresTables targetTables;
    PostgresExpressionGenerator gen;
    PostgresSelect select;

    public PostgresTLPBase(PostgresGlobalState state) {
        super(state);
        PostgresCommon.addCommonExpressionErrors(errors);
        PostgresCommon.addCommonFetchErrors(errors);
    }

    @Override
    public void check() throws SQLException {
        // clear left-over query string from previous test
        state.getState().queryString = null;
        s = state.getSchema();
        targetTables = s.getRandomTableNonEmptyTables();
        gen = new PostgresExpressionGenerator(state).setColumns(targetTables.getColumns());
        initializeTernaryPredicateVariants();
        select = new PostgresSelect();
        select.setFetchColumns(generateFetchColumns());
        List<PostgresTable> tables = targetTables.getTables();
        List<PostgresJoin> joins = PostgresNoRECOracle.getJoinStatements(state, targetTables.getColumns(), tables);
        List<PostgresExpression> tableList = tables.stream().map(t -> new PostgresFromTable(t, Randomly.getBoolean()))
                .collect(Collectors.toList());
        // TODO joins
        select.setFromList(tableList);
        select.setWhereClause(null);
        select.setJoinClauses(joins);
        if (Randomly.getBoolean()) {
            select.setForClause(ForClause.getRandom());
        }
    }

    List<PostgresExpression> generateFetchColumns() {
        return Arrays.asList(new PostgresColumnValue(targetTables.getColumns().get(0), null));
    }

    @Override
    protected ExpressionGenerator<PostgresExpression> getGen() {
        return gen;
    }

}

@@ -492,16 +492,74 @@
 Tables;%0A
+import sqlancer.postgres.PostgresSchema.PostgresDataType;%0A
 import s
@@ -600,16 +600,16 @@
 nValue;%0A
-
 import s
@@ -851,24 +851,128 @@
 sFromTable;%0A
+import sqlancer.postgres.ast.PostgresSelect.PostgresCTE;%0Aimport sqlancer.postgres.ast.PostgresConstant;%0A
 import sqlan
@@ -2885,24 +2885,24 @@
  getGen() %7B%0A
-
         retu
@@ -2916,10 +2916,1580 @@
 %0A    %7D%0A%0A
+    public static PostgresCTE createCTE(PostgresGlobalState globalState, PostgresTables tables) %7B%0A        PostgresExpressionGenerator gen = new PostgresExpressionGenerator(globalState).setColumns(tables.getColumns());%0A        PostgresSelect selectCTE = new PostgresSelect();%0A        selectCTE.setFromList(tables.getTables().stream().map(t -%3E new PostgresFromTable(t, Randomly.getBoolean()))%0A                .collect(Collectors.toList()));%0A        if (Randomly.getBoolean()) %7B%0A            selectCTE.setWhereClause(gen.generateExpression(0, PostgresDataType.BOOLEAN));%0A        %7D%0A        if (Randomly.getBooleanWithRatherLowProbability()) %7B%0A            selectCTE.setGroupByExpressions(gen.generateExpressions(Randomly.smallNumber() + 1));%0A            if (Randomly.getBoolean()) %7B%0A                selectCTE.setHavingClause(gen.generateHavingClause());%0A            %7D%0A        %7D%0A        if (Randomly.getBooleanWithRatherLowProbability()) %7B%0A            selectCTE.setOrderByExpressions(gen.generateOrderBy());%0A        %7D%0A        if (Randomly.getBoolean()) %7B%0A            selectCTE.setLimitClause(PostgresConstant.createIntConstant(Randomly.getPositiveOrZeroNonCachedInteger()));%0A            if (Randomly.getBoolean()) %7B%0A                selectCTE.setOffsetClause(%0A                        PostgresConstant.createIntConstant(Randomly.getPositiveOrZeroNonCachedInteger()));%0A            %7D%0A        %7D%0A        if (Randomly.getBooleanWithRatherLowProbability()) %7B%0A            selectCTE.setForClause(ForClause.getRandom());%0A        %7D%0A        return new PostgresCTE(selectCTE, %22cte%22);%0A    %7D%0A%0A
 %7D%0A

create CTE function added to Postgres TLP Base