package sqlancer.cockroachdb.gen;

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.cockroachdb.CockroachDBErrors;
import sqlancer.cockroachdb.CockroachDBProvider.CockroachDBGlobalState;
import sqlancer.cockroachdb.CockroachDBSchema.CockroachDBColumn;
import sqlancer.cockroachdb.CockroachDBSchema.CockroachDBTable;
import sqlancer.cockroachdb.CockroachDBVisitor;

public class CockroachDBInsertGenerator {

	public static Query insert(CockroachDBGlobalState globalState) {
		Set<String> errors = new HashSet<>();

		CockroachDBErrors.addExpressionErrors(errors); // e.g., caused by computed columns
		errors.add("violates not-null constraint");
		errors.add("violates unique constraint");
		errors.add("primary key column");
		errors.add("cannot write directly to computed column"); // TODO: do not select generated columns

		errors.add("failed to satisfy CHECK constraint");

		errors.add("violates foreign key constraint");
		errors.add("foreign key violation");
		errors.add("multi-part foreign key");
		StringBuilder sb = new StringBuilder();
		CockroachDBTable table = globalState.getSchema().getRandomTable(t -> !t.isView());
		if (Randomly.getBoolean()) {
			sb.append("INSERT INTO ");
		} else {
			sb.append("UPSERT INTO ");
			errors.add("UPSERT or INSERT...ON CONFLICT command cannot affect row a second time");
		}
		sb.append(table.getName());
		sb.append(" ");
		CockroachDBExpressionGenerator gen = new CockroachDBExpressionGenerator(globalState);
		if (Randomly.getBooleanWithSmallProbability()) {
			sb.append("DEFAULT VALUES");
		} else {
			List<CockroachDBColumn> columns = table.getRandomNonEmptyColumnSubset();
			sb.append("(");
			sb.append(columns.stream().map(c -> c.getName()).collect(Collectors.joining(", ")));
			sb.append(")");
			sb.append(" VALUES");
			for (int j = 0; j < Randomly.smallNumber() + 1; j++) {
				if (j != 0) {
					sb.append(", ");
				}
				sb.append("(");
				int i = 0;
				for (CockroachDBColumn c : columns) {
					if (i++ != 0) {
						sb.append(", ");
					}
					sb.append(CockroachDBVisitor.asString(gen.generateConstant(c.getColumnType())));
				}
				sb.append(")");
			}
		}
		if (Randomly.getBoolean() && false) {
			sb.append(" ON CONFLICT (");
			sb.append(table.getRandomNonEmptyColumnSubset().stream().map(c -> c.getName()).collect(Collectors.joining(", ")));
			sb.append(")");
			// WHERE clause not yet implemented, see https://github.com/cockroachdb/cockroach/issues/32557
			sb.append(" DO ");
			sb.append(" NOTHING ");
			errors.add("there is no unique or exclusion constraint matching the ON CONFLICT specification");
		}
		CockroachDBErrors.addTransactionErrors(errors);
		return new QueryAdapter(sb.toString(), errors);
	}

}

@@ -1279,28 +1279,43 @@
 sView());%0A%09%09
-if (
+boolean isUpsert = 
 Randomly.get
@@ -1323,16 +1323,33 @@
 oolean()
+;%0A%09%09if (!isUpsert
 ) %7B%0A%09%09%09s
@@ -2353,13 +2353,17 @@
  && 
-false
+!isUpsert
 ) %7B%0A
@@ -2658,31 +2658,627 @@
 %0A%09%09%09
-sb.append(%22 NOTHING %22);
+if (Randomly.getBoolean()) %7B%0A%09%09%09%09sb.append(%22 NOTHING %22);%0A%09%09%09%7D else %7B%0A%09%09%09%09// TODO: also support excluded. (see https://www.cockroachlabs.com/docs/stable/insert.html)%0A%09%09%09%09sb.append(%22 UPDATE SET %22);%0A%09%09%09%09List%3CCockroachDBColumn%3E columns = table.getRandomNonEmptyColumnSubset();%0A%09%09%09%09int i = 0;%0A%09%09%09%09for (CockroachDBColumn c : columns) %7B%0A%09%09%09%09%09if (i++ != 0) %7B%0A%09%09%09%09%09%09sb.append(%22, %22);%0A%09%09%09%09%09%7D%0A%09%09%09%09%09sb.append(c.getName());%0A%09%09%09%09%09sb.append(%22 = %22);%0A%09%09%09%09%09sb.append(CockroachDBVisitor.asString(gen.generateConstant(c.getColumnType())));%0A%09%09%09%09%7D%0A%09%09%09%09errors.add(%22UPSERT or INSERT...ON CONFLICT command cannot affect row a second time%22);%0A%09%09%09%7D
 %0A%09%09%09

Extend the ON CONFLICT clause in CockroachDB