package sqlancer.duckdb;

import java.io.IOException;

import java.sql.Connection;

import java.sql.DriverManager;

import java.sql.SQLException;

import java.util.HashSet;

import java.util.Set;

import java.util.stream.Collectors;

import sqlancer.AbstractAction;

import sqlancer.CompositeTestOracle;

import sqlancer.GlobalState;

import sqlancer.IgnoreMeException;

import sqlancer.Main.QueryManager;

import sqlancer.Main.StateLogger;

import sqlancer.ProviderAdapter;

import sqlancer.Query;

import sqlancer.QueryAdapter;

import sqlancer.QueryProvider;

import sqlancer.Randomly;

import sqlancer.StatementExecutor;

import sqlancer.TestOracle;

import sqlancer.duckdb.DuckDBProvider.DuckDBGlobalState;

import sqlancer.duckdb.gen.DuckDBDeleteGenerator;

import sqlancer.duckdb.gen.DuckDBIndexGenerator;

import sqlancer.duckdb.gen.DuckDBInsertGenerator;

import sqlancer.duckdb.gen.DuckDBRandomQuerySynthesizer;

import sqlancer.duckdb.gen.DuckDBTableGenerator;

import sqlancer.duckdb.gen.DuckDBUpdateGenerator;

import sqlancer.duckdb.gen.DuckDBViewGenerator;

public class DuckDBProvider extends ProviderAdapter<DuckDBGlobalState, DuckDBOptions> {

public DuckDBProvider() {

super(DuckDBGlobalState.class, DuckDBOptions.class);

}

public enum Action implements AbstractAction<DuckDBGlobalState> {

INSERT(DuckDBInsertGenerator::getQuery), //

CREATE_INDEX(DuckDBIndexGenerator::getQuery), //

VACUUM((g) -> new QueryAdapter("VACUUM;")), //

ANALYZE((g) -> new QueryAdapter("ANALYZE;")), //

DELETE(DuckDBDeleteGenerator::generate), //

UPDATE(DuckDBUpdateGenerator::getQuery), //

CREATE_VIEW(DuckDBViewGenerator::generate), //

EXPLAIN((g) -> {

Set<String> errors = new HashSet<>();

DuckDBErrors.addExpressionErrors(errors);

DuckDBErrors.addGroupByErrors(errors);

return new QueryAdapter(

"EXPLAIN " + DuckDBToStringVisitor

.asString(DuckDBRandomQuerySynthesizer.generateSelect(g, Randomly.smallNumber() + 1)),

errors);

});

private final QueryProvider<DuckDBGlobalState> queryProvider;

Action(QueryProvider<DuckDBGlobalState> queryProvider) {

this.queryProvider = queryProvider;

}

@Override

public Query getQuery(DuckDBGlobalState state) throws SQLException {

return queryProvider.getQuery(state);

}

}

private static int mapActions(DuckDBGlobalState globalState, Action a) {

Randomly r = globalState.getRandomly();

switch (a) {

case INSERT:

return r.getInteger(0, globalState.getOptions().getMaxNumberInserts());

case CREATE_INDEX:

if (!globalState.getDmbsSpecificOptions().testIndexes) {

return 0;

}

// fall through

case UPDATE:

return r.getInteger(0, globalState.getDmbsSpecificOptions().maxNumUpdates + 1);

case VACUUM: // seems to be ignored

case ANALYZE: // seems to be ignored

case EXPLAIN:

return r.getInteger(0, 2);

case DELETE:

return r.getInteger(0, globalState.getDmbsSpecificOptions().maxNumDeletes + 1);

case CREATE_VIEW:

return r.getInteger(0, globalState.getDmbsSpecificOptions().maxNumViews + 1);

default:

throw new AssertionError(a);

}

}

public static class DuckDBGlobalState extends GlobalState<DuckDBOptions> {

private DuckDBSchema schema;

public void setSchema(DuckDBSchema schema) {

this.schema = schema;

}

public DuckDBSchema getSchema() {

return schema;

}

}

@Override

public void generateAndTestDatabase(DuckDBGlobalState globalState) throws SQLException {

StateLogger logger = globalState.getLogger();

QueryManager manager = globalState.getManager();

globalState.setSchema(DuckDBSchema.fromConnection(globalState.getConnection(), globalState.getDatabaseName()));

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

globalState.getLogger().writeCurrent(globalState.getState());

}

for (int i = 0; i < Randomly.fromOptions(1, 2); i++) {

boolean success = false;

do {

Query qt = new DuckDBTableGenerator().getQuery(globalState);

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

globalState.getLogger().writeCurrent(qt.getQueryString());

}

success = manager.execute(qt);

globalState.setSchema(

DuckDBSchema.fromConnection(globalState.getConnection(), globalState.getDatabaseName()));

} while (!success);

}

if (globalState.getSchema().getDatabaseTables().size() == 0) {

throw new IgnoreMeException(); // TODO

}

StatementExecutor<DuckDBGlobalState, Action> se = new StatementExecutor<>(globalState, Action.values(),

DuckDBProvider::mapActions, (q) -> {

if (q.couldAffectSchema()) {

globalState.setSchema(DuckDBSchema.fromConnection(globalState.getConnection(),

globalState.getDatabaseName()));

}

if (globalState.getSchema().getDatabaseTables().isEmpty()) {

throw new IgnoreMeException();

}

});

se.executeStatements();

manager.incrementCreateDatabase();

TestOracle oracle = new CompositeTestOracle(globalState.getDmbsSpecificOptions().oracle.stream().map(o -> {

try {

return o.create(globalState);

} catch (SQLException e1) {

throw new AssertionError(e1);

}

}).collect(Collectors.toList()));

for (int i = 0; i < globalState.getOptions().getNrQueries(); i++) {

try {

oracle.check();

manager.incrementSelectQueryCount();

} catch (IgnoreMeException e) {

}

}

try {

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

logger.getCurrentFileWriter().close();

logger.currentFileWriter = null;

}

} catch (IOException e) {

// TODO Auto-generated catch block

e.printStackTrace();

}

globalState.getConnection().close();

}

@Override

public Connection createDatabase(DuckDBGlobalState globalState) throws SQLException {

String url = "jdbc:duckdb:";

return DriverManager.getConnection(url, globalState.getOptions().getUserName(),

globalState.getOptions().getPassword());

}

@Override

public String getDBMSName() {

return "duckdb";

}

}