package sqlancer.common.oracle;

import java.io.IOException;

import java.sql.SQLException;

import java.util.ArrayList;

import java.util.List;

import java.util.Optional;

import sqlancer.IgnoreMeException;

import sqlancer.Randomly;

import sqlancer.SQLGlobalState;

import sqlancer.common.DBMSCommon;

import sqlancer.common.ast.newast.Expression;

import sqlancer.common.ast.newast.Join;

import sqlancer.common.ast.newast.Select;

import sqlancer.common.gen.CERTGenerator;

import sqlancer.common.query.ExpectedErrors;

import sqlancer.common.query.SQLQueryAdapter;

import sqlancer.common.query.SQLancerResultSet;

import sqlancer.common.schema.AbstractSchema;

import sqlancer.common.schema.AbstractTable;

import sqlancer.common.schema.AbstractTableColumn;

import sqlancer.common.schema.AbstractTables;

public class CERTOracle<Z extends Select<J, E, T, C>, J extends Join<E, T, C>, E extends Expression<C>, S extends AbstractSchema<?, T>, T extends AbstractTable<C, ?, ?>, C extends AbstractTableColumn<?, ?>, G extends SQLGlobalState<?, S>>

implements TestOracle<G> {

private final G state;

private final CheckedFunction<SQLancerResultSet, Optional<Long>> rowCountParser;

private final CheckedFunction<SQLancerResultSet, Optional<String>> queryPlanParser;

private CERTGenerator<Z, J, E, T, C> gen;

private final ExpectedErrors errors;

public CERTOracle(G state, CERTGenerator<Z, J, E, T, C> gen, ExpectedErrors expectedErrors,

CheckedFunction<SQLancerResultSet, Optional<Long>> rowCountParser,

CheckedFunction<SQLancerResultSet, Optional<String>> queryPlanParser) {

if (state == null || gen == null || expectedErrors == null) {

throw new IllegalArgumentException("Null variables used to initialize test oracle.");

}

this.state = state;

this.gen = gen;

this.errors = expectedErrors;

this.rowCountParser = rowCountParser;

this.queryPlanParser = queryPlanParser;

}

@Override

public void check() throws SQLException {

S schema = state.getSchema();

AbstractTables<T, C> targetTables = TestOracleUtils.getRandomTableNonEmptyTables(schema);

gen = gen.setTablesAndColumns(targetTables);

List<E> fetchColumns = gen.generateFetchColumns(false);

Z select = gen.generateSelect();

select.setFetchColumns(fetchColumns);

select.setJoinClauses(gen.getRandomJoinClauses());

select.setFromList(gen.getTableRefs());

if (Randomly.getBoolean()) {

select.setWhereClause(gen.generateBooleanExpression());

}

if (Randomly.getBoolean()) {

select.setGroupByClause(fetchColumns);

if (Randomly.getBoolean()) {

select.setHavingClause(gen.generateBooleanExpression());

}

}

List<String> queryPlan1Sequences = new ArrayList<>();

List<String> queryPlan2Sequences = new ArrayList<>();

String queryString1 = gen.generateExplainQuery(select);

long rowCount1 = getRow(state, queryString1, queryPlan1Sequences);

boolean increase = gen.mutate(select);

String queryString2 = gen.generateExplainQuery(select);

long rowCount2 = getRow(state, queryString2, queryPlan2Sequences);

if (DBMSCommon.editDistance(queryPlan1Sequences, queryPlan2Sequences) > 1) {

return;

}

// Check the results

if (increase && rowCount1 > rowCount2 || !increase && rowCount1 < rowCount2) {

throw new AssertionError("Inconsistent result for query: " + queryString1 + "; --" + rowCount1 + "\n"

+ queryString2 + "; --" + rowCount2);

}

}

private Long getRow(SQLGlobalState<?, ?> globalState, String explainQuery, List<String> queryPlanSequences)

throws AssertionError, SQLException {

Optional<Long> row = Optional.empty();

// Log the query

if (globalState.getOptions().logEachSelect()) {

globalState.getLogger().writeCurrent(explainQuery);

try {

globalState.getLogger().getCurrentFileWriter().flush();

} catch (IOException e) {

e.printStackTrace();

}

}

// Get the row count

SQLQueryAdapter q = new SQLQueryAdapter(explainQuery, errors);

try (SQLancerResultSet rs = q.executeAndGet(globalState)) {

if (rs != null) {

while (rs.next()) {

Optional<Long> rowCount = rowCountParser.apply(rs);

if (row.isEmpty() && rowCount.isPresent()) {

row = rowCount;

}

Optional<String> queryPlanSequence = queryPlanParser.apply(rs);

queryPlanSequence.ifPresent(qps -> queryPlanSequences.add(qps));

}

}

} catch (IgnoreMeException e) {

throw new IgnoreMeException();

} catch (Exception e) {

throw new AssertionError(q.getQueryString(), e);

}

return row.orElseThrow(IgnoreMeException::new);

}

@FunctionalInterface

public interface CheckedFunction<T, R> {

R apply(T t) throws SQLException;

}

}