package sqlancer.tidb.gen;

import java.sql.SQLException;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

import sqlancer.Query;
import sqlancer.QueryAdapter;
import sqlancer.Randomly;
import sqlancer.tidb.TiDBExpressionGenerator;
import sqlancer.tidb.TiDBProvider.TiDBGlobalState;
import sqlancer.tidb.TiDBSchema.TiDBColumn;
import sqlancer.tidb.TiDBSchema.TiDBTable;

public class TiDBInsertGenerator {

	private final TiDBGlobalState globalState;
	private final Set<String> errors = new HashSet<>();
	private TiDBExpressionGenerator gen;

	public TiDBInsertGenerator(TiDBGlobalState globalState) {
		this.globalState = globalState;
		errors.add("Duplicate entry");
		errors.add("cannot be null");
		errors.add("doesn't have a default value");
		errors.add("Out of range value");
		errors.add("Incorrect tinyint value");
		errors.add("Data truncation");
		errors.add("Bad Number");
		errors.add("The value specified for generated column"); // TODO: do not insert data into generated columns
	}

	public static Query getQuery(TiDBGlobalState globalState) throws SQLException {
		return new TiDBInsertGenerator(globalState).get();
	}

	private Query get() {
		TiDBTable table = globalState.getSchema().getRandomTable(t -> !t.isView());
		gen = new TiDBExpressionGenerator(globalState).setColumns(table.getColumns());
		StringBuilder sb = new StringBuilder();
		sb.append(Randomly.fromOptions("INSERT", "REPLACE"));
		sb.append(" INTO ");
		sb.append(table.getName());
		if (Randomly.getBoolean()) {
			sb.append(" VALUES ");
			List<TiDBColumn> columns = table.getColumns();
			insertColumns(sb, columns);
		} else {
			List<TiDBColumn> columnSubset = table.getRandomNonEmptyColumnSubset();
			sb.append("(");
			sb.append(columnSubset.stream().map(c -> c.getName()).collect(Collectors.joining(", ")));
			sb.append(") VALUES ");
			insertColumns(sb, columnSubset);
		}
		return new QueryAdapter(sb.toString(), errors);
	}

	private void insertColumns(StringBuilder sb, List<TiDBColumn> columns) {
		for (int nrRows = 0; nrRows < Randomly.smallNumber() + 1; nrRows++) {
			if (nrRows != 0) {
				sb.append(", ");
			}
			sb.append("(");
			for (int nrColumn = 0; nrColumn < columns.size(); nrColumn++) {
				if (nrColumn != 0) {
					sb.append(", ");
				}
				insertValue(sb, columns.get(nrColumn));
			}
			sb.append(")");
		}
	}

	private void insertValue(StringBuilder sb, TiDBColumn tiDBColumn) {
		sb.append(gen.generateConstant()); // TODO: try to insert valid data
	}

}

@@ -421,16 +421,58 @@
 DBTable;
+%0Aimport sqlancer.tidb.visitor.TiDBVisitor;
 %0A%0Apublic
@@ -1460,61 +1460,422 @@
 ;%0A%09%09
-sb.append(Randomly.fromOptions(%22INSERT%22, %22REPLACE%22));
+boolean isInsert = Randomly.getBoolean();%0A%09%09if (isInsert) %7B%0A%09%09%09sb.append(%22INSERT%22);%0A%09%09%7D else %7B%0A%09%09%09sb.append(%22REPLACE%22);%0A%09%09%7D%0A%09%09if (Randomly.getBooleanWithRatherLowProbability()) %7B%0A%09%09%09sb.append(%22 %22);%0A%09%09%09sb.append(Randomly.fromOptions(%22LOW_PRIORITY%22, %22HIGH_PRIORITY%22, %22DELAYED%22));%0A%09%09%7D%0A%09%09if (isInsert && Randomly.getBoolean() && false /* https://github.com/pingcap/tidb/issues/16025 */) %7B%0A%09%09%09sb.append(%22 IGNORE %22);%0A%09%09%7D
 %0A%09%09s
@@ -2286,16 +2286,16 @@
 UES %22);%0A
-
 %09%09%09inser
@@ -2326,16 +2326,236 @@
 t);%0A%09%09%7D%0A
+%09%09if (isInsert && Randomly.getBoolean()) %7B%0A%09%09%09sb.append(%22 ON DUPLICATE KEY UPDATE %22);%0A%09%09%09sb.append(table.getRandomColumn().getName());%0A%09%09%09sb.append(%22=%22);%0A%09%09%09sb.append(TiDBVisitor.asString(gen.generateExpression()));%0A%09%09%7D%0A
 %09%09return

Enhance the INSERT statement