package executor

import (
	"context"
	"reflect"
	"runtime"
	"strings"
	"testing"

	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql/parser"
	"github.com/everything-on-mysql/cassandra-on-mysql/pkg/cql/translator"
)

func mustParseStatement(t *testing.T, cql string) parser.Statement {
	t.Helper()
	stmt, err := parser.NewParser(cql).Parse()
	if err != nil {
		t.Fatalf("failed to parse %q: %v", cql, err)
	}
	if stmt == nil {
		t.Fatalf("parsed nil statement for %q", cql)
	}
	return stmt
}

func TestValidateBatchEntry(t *testing.T) {
	tests := []struct {
		name    string
		cql     string
		wantErr bool
	}{
		{
			name:    "insert ok",
			cql:     "INSERT INTO ks.tbl (id) VALUES (1)",
			wantErr: false,
		},
		{
			name:    "insert if not exists",
			cql:     "INSERT INTO ks.tbl (id) VALUES (1) IF NOT EXISTS",
			wantErr: true,
		},
		{
			name:    "update if exists",
			cql:     "UPDATE ks.tbl SET v = 1 WHERE id = 1 IF EXISTS",
			wantErr: true,
		},
		{
			name:    "delete if condition",
			cql:     "DELETE FROM ks.tbl WHERE id = 1 IF v = 2",
			wantErr: true,
		},
		{
			name:    "select not allowed",
			cql:     "SELECT * FROM ks.tbl",
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			stmt := mustParseStatement(t, tt.cql)
			err := validateBatchEntry(stmt)
			if (err != nil) != tt.wantErr {
				t.Fatalf("validateBatchEntry error = %v, wantErr %t", err, tt.wantErr)
			}
		})
	}
}

func TestCQLWriteRowsForBatchEntriesCountsOriginalRequests(t *testing.T) {
	entries := []BatchEntry{
		{Statement: mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (1, 'a'), (2, 'b')")},
		{Statement: mustParseStatement(t, "UPDATE ks.tbl SET v = 'c' WHERE id = 1")},
		{Statement: mustParseStatement(t, "DELETE FROM ks.tbl WHERE id = 2")},
	}

	if got := cqlWriteRowsForBatchEntries(entries); got != 4 {
		t.Fatalf("expected request write rows before merge/fold, got %d", got)
	}
}

func TestCQLWriteRowsForBatchEntriesCountsFoldedOriginals(t *testing.T) {
	folded := BatchEntry{
		Statement: mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (1, 'merged')"),
		originals: []BatchEntry{
			{Statement: mustParseStatement(t, "UPDATE ks.tbl SET v = 'a' WHERE id = 1")},
			{Statement: mustParseStatement(t, "UPDATE ks.tbl SET v = 'b' WHERE id = 1")},
		},
	}

	if got := cqlWriteRowsForBatchEntries(folded.originals); got != 2 {
		t.Fatalf("expected folded original requests to count before fold, got %d", got)
	}
	if got := cqlWriteRowsForBatchEntry(folded); got != 1 {
		t.Fatalf("expected folded execution entry itself to remain one row, got %d", got)
	}
}

func TestValidateBatchCounterUsage(t *testing.T) {
	counterUpdate := mustParseStatement(t, "UPDATE ks.counter_table SET cnt = cnt + 1 WHERE id = 1")
	counterDelete := mustParseStatement(t, "DELETE FROM ks.counter_table WHERE id = 1")
	standardUpdate := mustParseStatement(t, "UPDATE ks.standard_table SET name = 'x' WHERE id = 1")
	standardDelete := mustParseStatement(t, "DELETE FROM ks.standard_table WHERE id = 1")
	insert := mustParseStatement(t, "INSERT INTO ks.standard_table (id, name) VALUES (1, 'x')")

	columnTypes := map[string]map[string]string{
		"ks.counter_table":  {"id": "int", "cnt": "counter"},
		"ks.standard_table": {"id": "int", "name": "text"},
	}

	lookup := func(table *parser.TableName) (map[string]string, error) {
		if table == nil {
			return nil, nil
		}
		key := table.Schema + "." + table.Name
		return columnTypes[key], nil
	}

	cases := []struct {
		name          string
		isCounter     bool
		entries       []parser.Statement
		expectFailure bool
	}{
		{
			name:          "logged batch with counter update",
			isCounter:     false,
			entries:       []parser.Statement{counterUpdate},
			expectFailure: true,
		},
		{
			name:          "logged batch with counter delete",
			isCounter:     false,
			entries:       []parser.Statement{counterDelete},
			expectFailure: true,
		},
		{
			name:      "counter batch with counter update",
			isCounter: true,
			entries:   []parser.Statement{counterUpdate},
		},
		{
			name:      "counter batch with counter delete",
			isCounter: true,
			entries:   []parser.Statement{counterDelete},
		},
		{
			name:          "counter batch with non-counter update",
			isCounter:     true,
			entries:       []parser.Statement{standardUpdate},
			expectFailure: true,
		},
		{
			name:          "counter batch with non-counter delete",
			isCounter:     true,
			entries:       []parser.Statement{standardDelete},
			expectFailure: true,
		},
		{
			name:          "counter batch with mixed entries",
			isCounter:     true,
			entries:       []parser.Statement{counterUpdate, standardUpdate},
			expectFailure: true,
		},
		{
			name:          "counter batch with insert",
			isCounter:     true,
			entries:       []parser.Statement{insert},
			expectFailure: true,
		},
		{
			name:      "logged batch with non-counter update",
			isCounter: false,
			entries:   []parser.Statement{standardUpdate},
		},
		{
			name:      "logged batch with non-counter delete",
			isCounter: false,
			entries:   []parser.Statement{standardDelete},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			err := validateBatchCounterUsage(tc.entries, tc.isCounter, lookup)
			if tc.expectFailure && err == nil {
				t.Fatalf("expected error")
			}
			if !tc.expectFailure && err != nil {
				t.Fatalf("unexpected error: %v", err)
			}
		})
	}
}

func TestSplitBatchParams(t *testing.T) {
	stmt := mustParseStatement(t, "BEGIN BATCH INSERT INTO ks.tbl (id, v) VALUES (?, ?); UPDATE ks.tbl SET v = ? WHERE id = ?; APPLY BATCH")
	batchStmt, ok := stmt.(*parser.BatchStatement)
	if !ok {
		t.Fatalf("expected BatchStatement, got %T", stmt)
	}

	params := []interface{}{1, "a", "b", 2}
	entries, err := splitBatchParams(batchStmt, params)
	if err != nil {
		t.Fatalf("splitBatchParams error: %v", err)
	}
	if len(entries) != 2 {
		t.Fatalf("expected 2 entries, got %d", len(entries))
	}

	if !reflect.DeepEqual(entries[0].Params, []interface{}{1, "a"}) {
		t.Fatalf("unexpected params for entry 0: %#v", entries[0].Params)
	}
	if !reflect.DeepEqual(entries[1].Params, []interface{}{"b", 2}) {
		t.Fatalf("unexpected params for entry 1: %#v", entries[1].Params)
	}
}

func TestSplitBatchParamsWithUsingTimestamp(t *testing.T) {
	stmt := mustParseStatement(t, "BEGIN BATCH USING TIMESTAMP ? INSERT INTO ks.tbl (id, v) VALUES (?, ?); UPDATE ks.tbl SET v = ? WHERE id = ?; APPLY BATCH")
	batchStmt, ok := stmt.(*parser.BatchStatement)
	if !ok {
		t.Fatalf("expected BatchStatement, got %T", stmt)
	}

	params := []interface{}{int64(5), 1, "a", "b", 2}
	entries, err := splitBatchParams(batchStmt, params)
	if err != nil {
		t.Fatalf("splitBatchParams error: %v", err)
	}
	if len(entries) != 2 {
		t.Fatalf("expected 2 entries, got %d", len(entries))
	}

	if !reflect.DeepEqual(entries[0].Params, []interface{}{1, "a"}) {
		t.Fatalf("unexpected params for entry 0: %#v", entries[0].Params)
	}
	if !reflect.DeepEqual(entries[1].Params, []interface{}{"b", 2}) {
		t.Fatalf("unexpected params for entry 1: %#v", entries[1].Params)
	}
}

func TestBatchExecutionPlan(t *testing.T) {
	exec := &Executor{}
	mode, workers := exec.batchExecutionPlan("UNLOGGED")
	if mode != batchExecutionTransactional {
		t.Fatalf("expected transactional mode when feature is disabled")
	}
	if workers != 1 {
		t.Fatalf("expected worker count 1 when feature is disabled, got %d", workers)
	}

	exec.unloggedBatchParallel = true
	exec.unloggedBatchWorkers = 8
	mode, workers = exec.batchExecutionPlan("UNLOGGED")
	if mode != batchExecutionBestEffortParallel {
		t.Fatalf("expected best-effort parallel mode for unlogged batch")
	}
	if workers != 8 {
		t.Fatalf("expected worker count 8, got %d", workers)
	}

	mode, workers = exec.batchExecutionPlan("LOGGED")
	if mode != batchExecutionTransactional {
		t.Fatalf("expected logged batch to remain transactional")
	}
	if workers != 1 {
		t.Fatalf("expected worker count 1 for logged batch, got %d", workers)
	}

	exec.unloggedBatchWorkers = 0
	mode, workers = exec.batchExecutionPlan("UNLOGGED")
	if mode != batchExecutionBestEffortParallel {
		t.Fatalf("expected best-effort parallel mode for unlogged batch")
	}
	want := runtime.NumCPU() * 4
	if want < 1 {
		want = 1
	}
	if workers != want {
		t.Fatalf("expected default workers %d, got %d", want, workers)
	}
}

func TestBatchInsertMergeLimits(t *testing.T) {
	exec := &Executor{}
	rows, bytes := exec.batchInsertMergeLimits()
	if rows != defaultBatchInsertMergeMaxRows {
		t.Fatalf("expected default rows %d, got %d", defaultBatchInsertMergeMaxRows, rows)
	}
	if bytes != defaultBatchInsertMergeMaxBytes {
		t.Fatalf("expected default bytes %d, got %d", defaultBatchInsertMergeMaxBytes, bytes)
	}

	exec.batchInsertMergeRows = 512
	exec.batchInsertMergeBytes = 512 * 1024
	rows, bytes = exec.batchInsertMergeLimits()
	if rows != 512 {
		t.Fatalf("expected configured rows 512, got %d", rows)
	}
	if bytes != 512*1024 {
		t.Fatalf("expected configured bytes %d, got %d", 512*1024, bytes)
	}
}

func testBatchPlanningEnvironment(resolvePK batchUpdatePKResolver, resolveColumnOrder func(table *parser.TableName) ([]string, error), resolveSplitKeys batchSplitKeyResolver) batchPlanningEnvironment {
	return batchPlanningEnvironment{
		resolvePK:        resolvePK,
		resolveSplitKeys: resolveSplitKeys,
		foldEnv: batchFoldEnvironment{
			resolvePK:          resolvePK,
			resolveColumnOrder: resolveColumnOrder,
			shouldInsertIncrement: func(table *parser.TableName, column string) bool {
				return false
			},
		},
	}
}

func TestTryCoalesceHomogeneousUnloggedBatchGroupsGloballyFoldsInsertsAcrossLimits(t *testing.T) {
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	entries := []BatchEntry{
		{Statement: insert, Params: []interface{}{1, "a1"}},
		{Statement: insert, Params: []interface{}{2, "b1"}},
		{Statement: insert, Params: []interface{}{3, "c1"}},
		{Statement: insert, Params: []interface{}{1, "a2"}},
		{Statement: insert, Params: []interface{}{2, "b2"}},
		{Statement: insert, Params: []interface{}{3, "c2"}},
	}
	resolvePK := func(table *parser.TableName) ([]string, error) {
		return []string{"id"}, nil
	}
	planning := testBatchPlanningEnvironment(resolvePK, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "v"}, nil
	}, nil)

	groups, ok, err := tryCoalesceHomogeneousUnloggedBatchGroupsWithPlanningEnvironment(entries, 2, 256*1024, false, planning)
	if err != nil {
		t.Fatalf("tryCoalesceHomogeneousUnloggedBatchGroupsWithPlanningEnvironment failed: %v", err)
	}
	if !ok {
		t.Fatalf("expected homogeneous insert batch to use global fold planner")
	}
	if len(groups) != 2 {
		t.Fatalf("expected 2 groups after strict row chunking, got %d", len(groups))
	}
	stmt1, ok := groups[0].merged.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected first merged statement to be INSERT, got %T", groups[0].merged.Statement)
	}
	stmt2, ok := groups[1].merged.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected second merged statement to be INSERT, got %T", groups[1].merged.Statement)
	}
	if len(stmt1.Values) != 2 || len(stmt2.Values) != 1 {
		t.Fatalf("expected chunk sizes 2 and 1, got %d and %d", len(stmt1.Values), len(stmt2.Values))
	}
	_, _, _, mergedKeys := collectBatchPrimaryKeyDebugInfo(groups, resolvePK)
	if len(mergedKeys) != 2 {
		t.Fatalf("expected 2 merged key groups, got %d", len(mergedKeys))
	}
	if len(mergedKeys[0].PrimaryKeys) != 2 || len(mergedKeys[1].PrimaryKeys) != 1 {
		t.Fatalf("expected unique PK distribution 2 and 1, got %d and %d", len(mergedKeys[0].PrimaryKeys), len(mergedKeys[1].PrimaryKeys))
	}
	if mergedKeys[0].GroupBreakReason != batchGroupBreakReasonLimitRows {
		t.Fatalf("expected first group break reason %q, got %q", batchGroupBreakReasonLimitRows, mergedKeys[0].GroupBreakReason)
	}
	if mergedKeys[1].GroupBreakReason != batchGroupBreakReasonEndOfEntries {
		t.Fatalf("expected second group break reason %q, got %q", batchGroupBreakReasonEndOfEntries, mergedKeys[1].GroupBreakReason)
	}
	if !strings.HasPrefix(mergedKeys[0].BatchTemplateKey, "insert:ks.tbl|cols=") {
		t.Fatalf("unexpected batch template key: got=%q", mergedKeys[0].BatchTemplateKey)
	}
	if mergedKeys[0].BatchTemplateKey != mergedKeys[1].BatchTemplateKey {
		t.Fatalf("expected both groups to share the same template key, got %q and %q", mergedKeys[0].BatchTemplateKey, mergedKeys[1].BatchTemplateKey)
	}
}

func TestTryCoalesceHomogeneousUnloggedBatchGroupsGloballyFoldsUpdatesAcrossLimits(t *testing.T) {
	update := mustParseStatement(t, "UPDATE ks.tbl SET v = ? WHERE id = ?")
	entries := []BatchEntry{
		{Statement: update, Params: []interface{}{"a1", 1}},
		{Statement: update, Params: []interface{}{"b1", 2}},
		{Statement: update, Params: []interface{}{"c1", 3}},
		{Statement: update, Params: []interface{}{"a2", 1}},
		{Statement: update, Params: []interface{}{"b2", 2}},
		{Statement: update, Params: []interface{}{"c2", 3}},
	}
	resolvePK := func(table *parser.TableName) ([]string, error) {
		return []string{"id"}, nil
	}
	planning := testBatchPlanningEnvironment(resolvePK, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "v"}, nil
	}, nil)

	groups, ok, err := tryCoalesceHomogeneousUnloggedBatchGroupsWithPlanningEnvironment(entries, 2, 256*1024, true, planning)
	if err != nil {
		t.Fatalf("tryCoalesceHomogeneousUnloggedBatchGroupsWithPlanningEnvironment failed: %v", err)
	}
	if !ok {
		t.Fatalf("expected homogeneous update batch to use global fold planner")
	}
	if len(groups) != 2 {
		t.Fatalf("expected 2 groups after strict row chunking, got %d", len(groups))
	}
	stmt1, ok := groups[0].merged.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected first merged statement to be INSERT, got %T", groups[0].merged.Statement)
	}
	stmt2, ok := groups[1].merged.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected second merged statement to be INSERT, got %T", groups[1].merged.Statement)
	}
	if len(stmt1.Values) != 2 || len(stmt2.Values) != 1 {
		t.Fatalf("expected chunk sizes 2 and 1, got %d and %d", len(stmt1.Values), len(stmt2.Values))
	}
}

func TestTryCoalesceHomogeneousUnloggedBatchGroupsBucketsInterleavedTemplatesBeforeChunking(t *testing.T) {
	deltaUpdate := mustParseStatement(t, "UPDATE ks.tbl SET unread_count = unread_count + ?, mention_count = mention_count + ?, note = ? WHERE id = ?")
	plainUpdate := mustParseStatement(t, "UPDATE ks.tbl SET unread_count = ?, mention_count = ?, note = ? WHERE id = ?")
	resolvePK := func(table *parser.TableName) ([]string, error) {
		return []string{"id"}, nil
	}
	planning := testBatchPlanningEnvironment(resolvePK, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "unread_count", "mention_count", "note"}, nil
	}, nil)

	groups, ok, err := tryCoalesceHomogeneousUnloggedBatchGroupsWithPlanningEnvironment([]BatchEntry{
		{Statement: deltaUpdate, Params: []interface{}{int64(1), int64(2), "delta-1", int64(1)}},
		{Statement: plainUpdate, Params: []interface{}{int64(10), int64(20), "plain-2", int64(2)}},
		{Statement: deltaUpdate, Params: []interface{}{int64(3), int64(4), "delta-3", int64(3)}},
	}, 256, 256*1024, true, planning)
	if err != nil {
		t.Fatalf("tryCoalesceHomogeneousUnloggedBatchGroupsWithPlanningEnvironment failed: %v", err)
	}
	if !ok {
		t.Fatalf("expected homogeneous update batch to use global fold planner")
	}
	if len(groups) != 2 {
		t.Fatalf("expected interleaved templates to regroup into 2 groups, got %d", len(groups))
	}

	_, _, _, mergedKeys := collectBatchPrimaryKeyDebugInfo(groups, resolvePK)
	if len(mergedKeys) != 2 {
		t.Fatalf("expected 2 merged key groups, got %d", len(mergedKeys))
	}
	if !strings.Contains(mergedKeys[0].BatchTemplateKey, "increments=mention_count,unread_count") {
		t.Fatalf("expected first group to keep increment template, got %q", mergedKeys[0].BatchTemplateKey)
	}
	if strings.Contains(mergedKeys[1].BatchTemplateKey, "increments=") {
		t.Fatalf("expected second group to be plain template, got %q", mergedKeys[1].BatchTemplateKey)
	}
	if len(mergedKeys[0].PrimaryKeys) != 2 || len(mergedKeys[1].PrimaryKeys) != 1 {
		t.Fatalf("expected regrouped PK distribution 2 and 1, got %d and %d", len(mergedKeys[0].PrimaryKeys), len(mergedKeys[1].PrimaryKeys))
	}
}

func TestTryCoalesceHomogeneousUnloggedBatchGroupsRejectsMixedTypes(t *testing.T) {
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	update := mustParseStatement(t, "UPDATE ks.tbl SET v = ? WHERE id = ?")
	resolvePK := func(table *parser.TableName) ([]string, error) {
		return []string{"id"}, nil
	}
	planning := testBatchPlanningEnvironment(resolvePK, nil, nil)

	_, ok, err := tryCoalesceHomogeneousUnloggedBatchGroupsWithPlanningEnvironment([]BatchEntry{
		{Statement: insert, Params: []interface{}{1, "a"}},
		{Statement: update, Params: []interface{}{"b", 1}},
	}, 4, 256*1024, true, planning)
	if err != nil {
		t.Fatalf("tryCoalesceHomogeneousUnloggedBatchGroupsWithPlanningEnvironment failed: %v", err)
	}
	if ok {
		t.Fatalf("expected mixed-type batch to fall back from global fold planner")
	}
}

func TestTryCoalesceHomogeneousUnloggedBatchGroupsReportsFallbackReasonForMixedTypes(t *testing.T) {
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	update := mustParseStatement(t, "UPDATE ks.tbl SET v = ? WHERE id = ?")
	resolvePK := func(table *parser.TableName) ([]string, error) {
		return []string{"id"}, nil
	}
	planning := testBatchPlanningEnvironment(resolvePK, nil, nil)

	_, decision, err := tryCoalesceHomogeneousUnloggedBatchGroupsWithDebug([]BatchEntry{
		{Statement: insert, Params: []interface{}{1, "a"}},
		{Statement: update, Params: []interface{}{"b", 1}},
	}, 4, 256*1024, true, planning)
	if err != nil {
		t.Fatalf("tryCoalesceHomogeneousUnloggedBatchGroupsWithDebug failed: %v", err)
	}
	if decision.Applied {
		t.Fatalf("expected mixed-type batch not to use global fold planner")
	}
	if decision.FallbackReason != batchGlobalFoldFallbackMixedStatementKinds {
		t.Fatalf("expected fallback reason %q, got %q", batchGlobalFoldFallbackMixedStatementKinds, decision.FallbackReason)
	}
}

func TestTryCoalesceHomogeneousUnloggedBatchGroupsRejectsBarrierStatements(t *testing.T) {
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?) USING TTL ?")
	resolvePK := func(table *parser.TableName) ([]string, error) {
		return []string{"id"}, nil
	}
	planning := testBatchPlanningEnvironment(resolvePK, nil, nil)

	_, ok, err := tryCoalesceHomogeneousUnloggedBatchGroupsWithPlanningEnvironment([]BatchEntry{
		{Statement: insert, Params: []interface{}{1, "a", int64(1)}},
	}, 4, 256*1024, false, planning)
	if err != nil {
		t.Fatalf("tryCoalesceHomogeneousUnloggedBatchGroupsWithPlanningEnvironment failed: %v", err)
	}
	if ok {
		t.Fatalf("expected barrier statement to fall back from global fold planner")
	}
}

func TestTryCoalesceHomogeneousUnloggedBatchGroupsReportsFallbackReasonForBarrierStatement(t *testing.T) {
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?) USING TTL ?")
	resolvePK := func(table *parser.TableName) ([]string, error) {
		return []string{"id"}, nil
	}
	planning := testBatchPlanningEnvironment(resolvePK, nil, nil)

	_, decision, err := tryCoalesceHomogeneousUnloggedBatchGroupsWithDebug([]BatchEntry{
		{Statement: insert, Params: []interface{}{1, "a", int64(1)}},
	}, 4, 256*1024, false, planning)
	if err != nil {
		t.Fatalf("tryCoalesceHomogeneousUnloggedBatchGroupsWithDebug failed: %v", err)
	}
	if decision.Applied {
		t.Fatalf("expected barrier statement not to use global fold planner")
	}
	if decision.FallbackReason != batchGlobalFoldFallbackBarrierStatement {
		t.Fatalf("expected fallback reason %q, got %q", batchGlobalFoldFallbackBarrierStatement, decision.FallbackReason)
	}
}

func TestCoalesceBatchInsertEntriesMergesCompatibleInserts(t *testing.T) {
	insert1 := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	insert2 := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")

	entries := []BatchEntry{
		{Statement: insert1, Params: []interface{}{1, "a"}},
		{Statement: insert2, Params: []interface{}{2, "b"}},
	}

	merged := coalesceBatchInsertEntries(entries, 32, 256*1024)
	if len(merged) != 1 {
		t.Fatalf("expected 1 merged entry, got %d", len(merged))
	}
	stmt, ok := merged[0].Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected merged entry to remain insert, got %T", merged[0].Statement)
	}
	if len(stmt.Values) != 2 {
		t.Fatalf("expected merged insert to have 2 rows, got %d", len(stmt.Values))
	}
	if !reflect.DeepEqual(merged[0].Params, []interface{}{1, "a", 2, "b"}) {
		t.Fatalf("unexpected merged params: %#v", merged[0].Params)
	}
}

func TestCoalesceBatchInsertEntriesDoesNotMergeDifferentTemplates(t *testing.T) {
	insert1 := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	insert2 := mustParseStatement(t, "INSERT INTO ks.tbl (id, v2) VALUES (?, ?)")
	update := mustParseStatement(t, "UPDATE ks.tbl SET v = ? WHERE id = ?")

	entries := []BatchEntry{
		{Statement: insert1, Params: []interface{}{1, "a"}},
		{Statement: insert2, Params: []interface{}{2, "b"}},
		{Statement: update, Params: []interface{}{"c", 1}},
	}

	merged := coalesceBatchInsertEntries(entries, 32, 256*1024)
	if len(merged) != 3 {
		t.Fatalf("expected 3 entries, got %d", len(merged))
	}
}

func TestCoalesceBatchInsertEntriesRespectsLimits(t *testing.T) {
	insert1 := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	insert2 := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")

	entries := []BatchEntry{
		{Statement: insert1, Params: []interface{}{1, "a"}},
		{Statement: insert2, Params: []interface{}{2, "b"}},
	}

	limitedRows := coalesceBatchInsertEntries(entries, 1, 256*1024)
	if len(limitedRows) != 2 {
		t.Fatalf("expected row limit to prevent merge, got %d entries", len(limitedRows))
	}

	limitedBytes := coalesceBatchInsertEntries(entries, 32, 1)
	if len(limitedBytes) != 2 {
		t.Fatalf("expected byte limit to prevent merge, got %d entries", len(limitedBytes))
	}
}

func TestCoalesceBatchInsertGroupsWithResolverKeepsSamePrimaryKeyInSameGroupAcrossLimits(t *testing.T) {
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, v) VALUES (?, ?, ?)")
	entries := []BatchEntry{
		{Statement: insert, Params: []interface{}{1, 1, "a1"}},
		{Statement: insert, Params: []interface{}{2, 1, "b1"}},
		{Statement: insert, Params: []interface{}{3, 1, "c1"}},
		{Statement: insert, Params: []interface{}{1, 1, "a2"}},
		{Statement: insert, Params: []interface{}{2, 1, "b2"}},
		{Statement: insert, Params: []interface{}{3, 1, "c2"}},
	}

	groups := coalesceBatchInsertGroupsWithResolver(entries, 4, 256*1024, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if len(groups) != 1 {
		t.Fatalf("expected same-primary-key inserts not to split across groups, got %d groups", len(groups))
	}
	stmt, ok := groups[0].merged.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected merged statement to remain insert, got %T", groups[0].merged.Statement)
	}
	if len(stmt.Values) != 3 {
		t.Fatalf("expected merged insert to contain 3 deduplicated rows, got %d", len(stmt.Values))
	}
}

func TestCoalesceBatchInsertGroupsWithResolverSplitsAtSafeBoundary(t *testing.T) {
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, v) VALUES (?, ?, ?)")
	entries := []BatchEntry{
		{Statement: insert, Params: []interface{}{1, 1, "a1"}},
		{Statement: insert, Params: []interface{}{1, 1, "a2"}},
		{Statement: insert, Params: []interface{}{2, 1, "b1"}},
		{Statement: insert, Params: []interface{}{2, 1, "b2"}},
		{Statement: insert, Params: []interface{}{3, 1, "c1"}},
		{Statement: insert, Params: []interface{}{3, 1, "c2"}},
		{Statement: insert, Params: []interface{}{4, 1, "d1"}},
		{Statement: insert, Params: []interface{}{4, 1, "d2"}},
		{Statement: insert, Params: []interface{}{5, 1, "e1"}},
		{Statement: insert, Params: []interface{}{5, 1, "e2"}},
	}

	groups := coalesceBatchInsertGroupsWithResolver(entries, 4, 256*1024, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if len(groups) != 2 {
		t.Fatalf("expected inserts to split at safe boundary, got %d groups", len(groups))
	}
	stmt1, ok := groups[0].merged.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected first merged statement to remain insert, got %T", groups[0].merged.Statement)
	}
	if len(stmt1.Values) != 4 {
		t.Fatalf("expected first merged insert to contain 4 rows, got %d", len(stmt1.Values))
	}
	stmt2, ok := groups[1].merged.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected second merged statement to remain insert, got %T", groups[1].merged.Statement)
	}
	if len(stmt2.Values) != 1 {
		t.Fatalf("expected second merged insert to contain 1 row, got %d", len(stmt2.Values))
	}
}

func TestAbsorbDuplicateFoldSingletonSegmentsMergesIntoNeighbor(t *testing.T) {
	makeCandidate := func(rows, bytes, originals int) batchMergeCandidate {
		return batchMergeCandidate{
			rows:      rows,
			bytes:     bytes,
			originals: make([]BatchEntry, originals),
		}
	}

	segments := []batchMergeSegment{
		{
			candidates:       []batchMergeCandidate{makeCandidate(3, 30, 3)},
			groupingStrategy: batchGroupingStrategyPrimaryKey,
		},
		{
			candidates:       []batchMergeCandidate{makeCandidate(1, 10, 2)},
			groupingStrategy: batchGroupingStrategyPrimaryKey,
		},
		{
			candidates:       []batchMergeCandidate{makeCandidate(2, 20, 2)},
			groupingStrategy: batchGroupingStrategyPrimaryKey,
		},
	}

	merged := absorbDuplicateFoldSingletonSegments(segments, 4, 100)
	if len(merged) != 2 {
		t.Fatalf("expected duplicate-fold singleton to be absorbed, got %d segments", len(merged))
	}
	rows0, bytes0 := batchMergeSegmentSize(merged[0])
	if rows0 != 4 || bytes0 != 40 {
		t.Fatalf("expected first merged segment to keep within limits, got rows=%d bytes=%d", rows0, bytes0)
	}
	if len(merged[0].candidates) != 2 {
		t.Fatalf("expected first merged segment to contain 2 candidates, got %d", len(merged[0].candidates))
	}
	rows1, bytes1 := batchMergeSegmentSize(merged[1])
	if rows1 != 2 || bytes1 != 20 {
		t.Fatalf("expected trailing segment to remain unchanged, got rows=%d bytes=%d", rows1, bytes1)
	}
}

func TestAbsorbDuplicateFoldSingletonSegmentsPreservesPartitionHits(t *testing.T) {
	makeCandidate := func(splitKey string, originals int) batchMergeCandidate {
		return batchMergeCandidate{
			rows:      1,
			bytes:     10,
			originals: make([]BatchEntry, originals),
			splitKeys: []string{splitKey},
		}
	}

	segments := []batchMergeSegment{
		{
			candidates:       []batchMergeCandidate{makeCandidate("ks.tbl#hash:0", 2)},
			groupingStrategy: batchGroupingStrategyPartitionGroup,
			partitionHits:    []string{"ks.tbl#hash:0"},
		},
		{
			candidates:       []batchMergeCandidate{makeCandidate("ks.tbl#hash:1", 1)},
			groupingStrategy: batchGroupingStrategyPartitionGroup,
			partitionHits:    []string{"ks.tbl#hash:1"},
		},
	}

	merged := absorbDuplicateFoldSingletonSegments(segments, 4, 100)
	if len(merged) != 1 {
		t.Fatalf("expected partition-group singleton to be absorbed, got %d segments", len(merged))
	}
	if len(merged[0].partitionHits) != 2 {
		t.Fatalf("expected merged partition hits to be preserved, got %#v", merged[0].partitionHits)
	}
	if merged[0].partitionHits[0] != "ks.tbl#hash:0" || merged[0].partitionHits[1] != "ks.tbl#hash:1" {
		t.Fatalf("unexpected merged partition hits: %#v", merged[0].partitionHits)
	}
}

func TestCoalesceBatchInsertEntriesSkipsUnsupportedInsertShapes(t *testing.T) {
	normal := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	ifNotExists := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?) IF NOT EXISTS")
	withTS := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?) USING TIMESTAMP ?")

	entries := []BatchEntry{
		{Statement: normal, Params: []interface{}{1, "a"}},
		{Statement: ifNotExists, Params: []interface{}{2, "b"}},
		{Statement: withTS, Params: []interface{}{3, "c", int64(1)}},
	}

	merged := coalesceBatchInsertEntries(entries, 32, 256*1024)
	if len(merged) != 3 {
		t.Fatalf("expected unsupported insert shapes to remain separate, got %d entries", len(merged))
	}
}

func TestBuildBestEffortBatchWorkItemsMergedEntriesCarryFallback(t *testing.T) {
	insert1 := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	insert2 := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")

	entries := []BatchEntry{
		{Statement: insert1, Params: []interface{}{1, "a"}},
		{Statement: insert2, Params: []interface{}{2, "b"}},
	}

	items := buildBestEffortBatchWorkItems(entries, 32, 256*1024)
	if len(items) != 1 {
		t.Fatalf("expected 1 work item, got %d", len(items))
	}
	if len(items[0].fallback) != 2 {
		t.Fatalf("expected merged work item to keep 2 fallback entries, got %d", len(items[0].fallback))
	}
	stmt, ok := items[0].entry.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected merged work item statement to be insert, got %T", items[0].entry.Statement)
	}
	if len(stmt.Values) != 2 {
		t.Fatalf("expected merged statement to contain 2 rows, got %d", len(stmt.Values))
	}
}

func TestBuildBestEffortBatchWorkItemsSingleEntriesDoNotCarryFallback(t *testing.T) {
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	update := mustParseStatement(t, "UPDATE ks.tbl SET v = ? WHERE id = ?")

	entries := []BatchEntry{
		{Statement: insert, Params: []interface{}{1, "a"}},
		{Statement: update, Params: []interface{}{"b", 1}},
	}

	items := buildBestEffortBatchWorkItems(entries, 32, 256*1024)
	if len(items) != 2 {
		t.Fatalf("expected 2 work items, got %d", len(items))
	}
	for i, item := range items {
		if len(item.fallback) != 0 {
			t.Fatalf("expected item %d not to require fallback, got %d entries", i, len(item.fallback))
		}
	}
}

func TestCoalesceBatchInsertEntriesPreparedLikeReuseKeepsDistinctBindMarkers(t *testing.T) {
	// Simulate prepared-batch reuse where multiple entries point to the same parsed statement.
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	entries := []BatchEntry{
		{Statement: insert, Params: []interface{}{1, "a"}},
		{Statement: insert, Params: []interface{}{nil, "bad"}},
	}

	merged := coalesceBatchInsertEntries(entries, 32, 256*1024)
	if len(merged) != 1 {
		t.Fatalf("expected 1 merged entry, got %d", len(merged))
	}
	stmt, ok := merged[0].Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected merged entry to remain insert, got %T", merged[0].Statement)
	}
	if len(stmt.Values) != 2 {
		t.Fatalf("expected 2 value rows, got %d", len(stmt.Values))
	}
	if got := countBindMarkersInStatement(stmt); got != 4 {
		t.Fatalf("expected 4 bind markers for 2x2 placeholders after merge, got %d", got)
	}
}

func TestCoalesceBatchGroupsWithUpdateMergeMergesEligibleUpdates(t *testing.T) {
	update1 := mustParseStatement(t, "UPDATE ks.tbl SET v = ?, n = ? WHERE id = ? AND ck = ?")
	update2 := mustParseStatement(t, "UPDATE ks.tbl SET v = ?, n = ? WHERE id = ? AND ck = ?")
	entries := []BatchEntry{
		{Statement: update1, Params: []interface{}{"a", 1, 10, 7}},
		{Statement: update2, Params: []interface{}{"b", 2, 11, 7}},
	}

	groups, err := coalesceBatchGroupsWithUpdateMerge(entries, 64, 256*1024, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if err != nil {
		t.Fatalf("coalesceBatchGroupsWithUpdateMerge failed: %v", err)
	}
	if len(groups) != 1 {
		t.Fatalf("expected 1 merged group, got %d", len(groups))
	}
	stmt, ok := groups[0].merged.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected merged statement to be INSERT, got %T", groups[0].merged.Statement)
	}
	if len(stmt.Values) != 2 {
		t.Fatalf("expected merged INSERT to contain 2 rows, got %d", len(stmt.Values))
	}
	if len(groups[0].originals) != 2 {
		t.Fatalf("expected merged group to keep 2 originals, got %d", len(groups[0].originals))
	}
}

func TestCoalesceBatchGroupsWithUpdateMergeMergesSamePrimaryKeyWhenTemplateMatches(t *testing.T) {
	update1 := mustParseStatement(t, "UPDATE ks.tbl SET v = ?, n = ? WHERE id = ? AND ck = ?")
	update2 := mustParseStatement(t, "UPDATE ks.tbl SET v = ?, n = ? WHERE id = ? AND ck = ?")
	entries := []BatchEntry{
		{Statement: update1, Params: []interface{}{"a", 1, 10, 7}},
		{Statement: update2, Params: []interface{}{"b", 2, 10, 7}},
	}

	groups, err := coalesceBatchGroupsWithUpdateMerge(entries, 64, 256*1024, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if err != nil {
		t.Fatalf("coalesceBatchGroupsWithUpdateMerge failed: %v", err)
	}
	if len(groups) != 1 {
		t.Fatalf("expected 1 merged group, got %d", len(groups))
	}
	stmt, ok := groups[0].merged.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected merged statement to be INSERT, got %T", groups[0].merged.Statement)
	}
	if len(stmt.Values) != 1 {
		t.Fatalf("expected same-primary-key updates to collapse to 1 row, got %d", len(stmt.Values))
	}
	if len(groups[0].originals) != 2 {
		t.Fatalf("expected merged group to keep 2 originals, got %d", len(groups[0].originals))
	}
	if !reflect.DeepEqual(groups[0].merged.Params, []interface{}{"b", 2, 10, 7}) {
		t.Fatalf("expected last write to win for same-primary-key updates, got %#v", groups[0].merged.Params)
	}
}

func TestCoalesceBatchGroupsWithUpdateMergeKeepsSamePrimaryKeyInSameGroupAcrossLimits(t *testing.T) {
	update := mustParseStatement(t, "UPDATE ks.tbl SET v = ?, n = ? WHERE id = ? AND ck = ?")
	entries := []BatchEntry{
		{Statement: update, Params: []interface{}{"a1", 1, 1, 1}},
		{Statement: update, Params: []interface{}{"b1", 2, 2, 1}},
		{Statement: update, Params: []interface{}{"c1", 3, 3, 1}},
		{Statement: update, Params: []interface{}{"a2", 4, 1, 1}},
		{Statement: update, Params: []interface{}{"b2", 5, 2, 1}},
		{Statement: update, Params: []interface{}{"c2", 6, 3, 1}},
	}

	groups, err := coalesceBatchGroupsWithUpdateMerge(entries, 4, 256*1024, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if err != nil {
		t.Fatalf("coalesceBatchGroupsWithUpdateMerge failed: %v", err)
	}
	if len(groups) != 1 {
		t.Fatalf("expected same-primary-key updates not to split across groups, got %d groups", len(groups))
	}
	stmt, ok := groups[0].merged.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected merged statement to be INSERT, got %T", groups[0].merged.Statement)
	}
	if len(stmt.Values) != 3 {
		t.Fatalf("expected merged INSERT to contain 3 deduplicated rows, got %d", len(stmt.Values))
	}
}

func TestCoalesceBatchGroupsWithUpdateMergeCollapsesSamePrimaryKeyArithmeticValues(t *testing.T) {
	update1 := mustParseStatement(t, "UPDATE ks.tbl SET unread_count = unread_count + ?, mention_count = mention_count + ? WHERE id = ? AND ck = ?")
	update2 := mustParseStatement(t, "UPDATE ks.tbl SET unread_count = unread_count + ?, mention_count = mention_count + ? WHERE id = ? AND ck = ?")
	entries := []BatchEntry{
		{Statement: update1, Params: []interface{}{3, 2, 10, 7}},
		{Statement: update2, Params: []interface{}{5, 1, 10, 7}},
	}

	groups, err := coalesceBatchGroupsWithUpdateMerge(entries, 64, 256*1024, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if err != nil {
		t.Fatalf("coalesceBatchGroupsWithUpdateMerge failed: %v", err)
	}
	if len(groups) != 1 {
		t.Fatalf("expected 1 merged group, got %d", len(groups))
	}

	stmt, ok := groups[0].merged.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected merged statement to be INSERT, got %T", groups[0].merged.Statement)
	}
	if len(stmt.Values) != 1 {
		t.Fatalf("expected same-primary-key arithmetic updates to collapse to 1 row, got %d", len(stmt.Values))
	}
	if !reflect.DeepEqual(groups[0].merged.Params, []interface{}{int64(8), int64(3), 10, 7}) &&
		!reflect.DeepEqual(groups[0].merged.Params, []interface{}{8, 3, 10, 7}) {
		t.Fatalf("expected arithmetic values to accumulate, got %#v", groups[0].merged.Params)
	}
}

func TestCoalesceBatchGroupsWithUpdateMergeMergesArithmeticUpdatesWithIncrementHints(t *testing.T) {
	update1 := mustParseStatement(t, "UPDATE ks.tbl SET unread_count = unread_count + ?, mention_count = mention_count + ? WHERE id = ? AND ck = ?")
	update2 := mustParseStatement(t, "UPDATE ks.tbl SET unread_count = unread_count + ?, mention_count = mention_count + ? WHERE id = ? AND ck = ?")
	entries := []BatchEntry{
		{Statement: update1, Params: []interface{}{3, 2, 10, 7}},
		{Statement: update2, Params: []interface{}{5, 1, 11, 7}},
	}

	groups, err := coalesceBatchGroupsWithUpdateMerge(entries, 64, 256*1024, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if err != nil {
		t.Fatalf("coalesceBatchGroupsWithUpdateMerge failed: %v", err)
	}
	if len(groups) != 1 {
		t.Fatalf("expected 1 merged group, got %d", len(groups))
	}

	stmt, ok := groups[0].merged.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected merged statement to be INSERT, got %T", groups[0].merged.Statement)
	}
	if len(stmt.Values) != 2 {
		t.Fatalf("expected merged INSERT to contain 2 rows, got %d", len(stmt.Values))
	}
	if _, ok := stmt.Values[0][0].(*parser.BindMarker); !ok {
		t.Fatalf("expected arithmetic unread_count value to be rewritten as delta bind marker, got %T", stmt.Values[0][0])
	}
	if _, ok := stmt.Values[0][1].(*parser.BindMarker); !ok {
		t.Fatalf("expected arithmetic mention_count value to be rewritten as delta bind marker, got %T", stmt.Values[0][1])
	}
	if _, ok := stmt.OnDuplicateIncrementColumns["unread_count"]; !ok {
		t.Fatalf("expected unread_count to be marked as duplicate increment column")
	}
	if _, ok := stmt.OnDuplicateIncrementColumns["mention_count"]; !ok {
		t.Fatalf("expected mention_count to be marked as duplicate increment column")
	}
}

func TestCoalesceBatchGroupsWithUpdateMergeArithmeticBuildsMultiValuesUpsertSQL(t *testing.T) {
	update1 := mustParseStatement(t, "UPDATE ks.tbl SET unread_count = unread_count + ?, mention_count = mention_count + ? WHERE id = ? AND ck = ?")
	update2 := mustParseStatement(t, "UPDATE ks.tbl SET unread_count = unread_count + ?, mention_count = mention_count + ? WHERE id = ? AND ck = ?")
	entries := []BatchEntry{
		{Statement: update1, Params: []interface{}{3, 2, 10, 7}},
		{Statement: update2, Params: []interface{}{5, 1, 11, 7}},
	}

	groups, err := coalesceBatchGroupsWithUpdateMerge(entries, 64, 256*1024, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if err != nil {
		t.Fatalf("coalesceBatchGroupsWithUpdateMerge failed: %v", err)
	}
	if len(groups) != 1 {
		t.Fatalf("expected 1 merged group, got %d", len(groups))
	}

	insert, ok := groups[0].merged.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected merged statement to be INSERT, got %T", groups[0].merged.Statement)
	}

	exec := &Executor{translator: translator.NewTranslator()}
	plan, err := exec.buildInsertPlan(context.Background(), "", "", insert)
	if err != nil {
		t.Fatalf("buildInsertPlan failed: %v", err)
	}
	expectedSQL := "INSERT INTO `ks`.`tbl` (`unread_count`, `mention_count`, `id`, `ck`) VALUES (?, ?, ?, ?), (?, ?, ?, ?) ON DUPLICATE KEY UPDATE `unread_count` = COALESCE(`unread_count`, 0) + VALUES(`unread_count`), `mention_count` = COALESCE(`mention_count`, 0) + VALUES(`mention_count`), `id` = VALUES(`id`), `ck` = VALUES(`ck`)"
	if plan.sql != expectedSQL {
		t.Fatalf("unexpected merged SQL: %s", plan.sql)
	}
}

func TestCoalesceBatchGroupsWithUpdateMergeDoesNotMergeMixedArithmeticAndPlainAssignments(t *testing.T) {
	arithmetic := mustParseStatement(t, "UPDATE ks.tbl SET unread_count = unread_count + ? WHERE id = ? AND ck = ?")
	plain := mustParseStatement(t, "UPDATE ks.tbl SET unread_count = ? WHERE id = ? AND ck = ?")
	entries := []BatchEntry{
		{Statement: arithmetic, Params: []interface{}{3, 10, 7}},
		{Statement: plain, Params: []interface{}{5, 11, 7}},
	}

	groups, err := coalesceBatchGroupsWithUpdateMerge(entries, 64, 256*1024, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if err != nil {
		t.Fatalf("coalesceBatchGroupsWithUpdateMerge failed: %v", err)
	}
	if len(groups) != 2 {
		t.Fatalf("expected mixed arithmetic/plain updates not to merge, got %d groups", len(groups))
	}
	for i := range groups {
		if _, ok := groups[i].merged.Statement.(*parser.UpdateStatement); !ok {
			t.Fatalf("expected group %d to remain UPDATE, got %T", i, groups[i].merged.Statement)
		}
	}
}

func TestCoalesceBatchGroupsWithUpdateMergeDoesNotMergeSamePrimaryKeyDifferentTemplates(t *testing.T) {
	withNote := mustParseStatement(t, "UPDATE ks.tbl SET v = ?, n = ? WHERE id = ? AND ck = ?")
	withoutNote := mustParseStatement(t, "UPDATE ks.tbl SET v = ? WHERE id = ? AND ck = ?")
	entries := []BatchEntry{
		{Statement: withNote, Params: []interface{}{"a", 1, 10, 7}},
		{Statement: withoutNote, Params: []interface{}{"b", 10, 7}},
	}

	groups, err := coalesceBatchGroupsWithUpdateMerge(entries, 64, 256*1024, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if err != nil {
		t.Fatalf("coalesceBatchGroupsWithUpdateMerge failed: %v", err)
	}
	if len(groups) != 2 {
		t.Fatalf("expected same-primary-key updates with different templates not to merge, got %d groups", len(groups))
	}
	for i := range groups {
		if _, ok := groups[i].merged.Statement.(*parser.UpdateStatement); !ok {
			t.Fatalf("expected group %d to remain UPDATE, got %T", i, groups[i].merged.Statement)
		}
	}
}

func TestCoalesceBatchGroupsWithUpdateMergeKeepsSingleUpdateUnchanged(t *testing.T) {
	update := mustParseStatement(t, "UPDATE ks.tbl SET v = ? WHERE id = ? AND ck = ?")
	entries := []BatchEntry{
		{Statement: update, Params: []interface{}{"a", 10, 7}},
	}

	groups, err := coalesceBatchGroupsWithUpdateMerge(entries, 64, 256*1024, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if err != nil {
		t.Fatalf("coalesceBatchGroupsWithUpdateMerge failed: %v", err)
	}
	if len(groups) != 1 {
		t.Fatalf("expected 1 group, got %d", len(groups))
	}
	if _, ok := groups[0].merged.Statement.(*parser.UpdateStatement); !ok {
		t.Fatalf("expected single update group to remain UPDATE, got %T", groups[0].merged.Statement)
	}
}

func TestCoalesceBatchGroupsWithUpdateMergeRequiresFullPrimaryKeyEquality(t *testing.T) {
	update1 := mustParseStatement(t, "UPDATE ks.tbl SET v = ? WHERE id = ?")
	update2 := mustParseStatement(t, "UPDATE ks.tbl SET v = ? WHERE id = ?")
	entries := []BatchEntry{
		{Statement: update1, Params: []interface{}{"a", 10}},
		{Statement: update2, Params: []interface{}{"b", 11}},
	}

	groups, err := coalesceBatchGroupsWithUpdateMerge(entries, 64, 256*1024, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if err != nil {
		t.Fatalf("coalesceBatchGroupsWithUpdateMerge failed: %v", err)
	}
	if len(groups) != 2 {
		t.Fatalf("expected 2 non-merged groups, got %d", len(groups))
	}
	for i := range groups {
		if _, ok := groups[i].merged.Statement.(*parser.UpdateStatement); !ok {
			t.Fatalf("expected group %d to remain UPDATE, got %T", i, groups[i].merged.Statement)
		}
	}
}