package executor

import (
	"testing"

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

func TestCollectBatchPrimaryKeyDebugInfoForMergedUpdates(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)
	}

	prefix, inputKeys, mergeAppliedKeys, mergedKeys := collectBatchPrimaryKeyDebugInfo(groups, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if prefix != "ks.tbl(id,ck)" {
		t.Fatalf("unexpected prefix: got=%q want=%q", prefix, "ks.tbl(id,ck)")
	}

	expectedInput := []string{
		"10,7",
		"11,7",
	}
	if len(inputKeys) != len(expectedInput) {
		t.Fatalf("expected %d input keys, got %d (%#v)", len(expectedInput), len(inputKeys), inputKeys)
	}
	for i := range expectedInput {
		if inputKeys[i] != expectedInput[i] {
			t.Fatalf("unexpected input key %d: got=%q want=%q", i, inputKeys[i], expectedInput[i])
		}
	}
	if len(mergeAppliedKeys) != 0 {
		t.Fatalf("expected no same-primary-key merge-applied keys, got %#v", mergeAppliedKeys)
	}

	if len(mergedKeys) != 1 {
		t.Fatalf("expected 1 merged key group, got %d", len(mergedKeys))
	}
	if mergedKeys[0].OriginalStatementCount != 2 {
		t.Fatalf("expected merged group original count 2, got %d", mergedKeys[0].OriginalStatementCount)
	}
	if len(mergedKeys[0].PrimaryKeys) != 2 {
		t.Fatalf("expected 2 merged primary keys, got %d", len(mergedKeys[0].PrimaryKeys))
	}
	for i := range expectedInput {
		if mergedKeys[0].PrimaryKeys[i] != expectedInput[i] {
			t.Fatalf("unexpected merged key %d: got=%q want=%q", i, mergedKeys[0].PrimaryKeys[i], expectedInput[i])
		}
	}
}

func TestCollectBatchPrimaryKeyDebugInfoForCollapsedSamePrimaryKeyUpdates(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)
	}

	prefix, inputKeys, mergeAppliedKeys, mergedKeys := collectBatchPrimaryKeyDebugInfo(groups, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	if prefix != "ks.tbl(id,ck)" {
		t.Fatalf("unexpected prefix: got=%q want=%q", prefix, "ks.tbl(id,ck)")
	}
	expectedInput := []string{"10,7", "10,7"}
	if len(inputKeys) != len(expectedInput) {
		t.Fatalf("expected %d input keys, got %d (%#v)", len(expectedInput), len(inputKeys), inputKeys)
	}
	for i := range expectedInput {
		if inputKeys[i] != expectedInput[i] {
			t.Fatalf("unexpected input key %d: got=%q want=%q", i, inputKeys[i], expectedInput[i])
		}
	}
	expectedMergeApplied := []string{"10,7"}
	if len(mergeAppliedKeys) != len(expectedMergeApplied) {
		t.Fatalf("expected %d merge-applied keys, got %d (%#v)", len(expectedMergeApplied), len(mergeAppliedKeys), mergeAppliedKeys)
	}
	for i := range expectedMergeApplied {
		if mergeAppliedKeys[i] != expectedMergeApplied[i] {
			t.Fatalf("unexpected merge-applied key %d: got=%q want=%q", i, mergeAppliedKeys[i], expectedMergeApplied[i])
		}
	}
	if len(mergedKeys) != 1 {
		t.Fatalf("expected 1 merged key group, got %d", len(mergedKeys))
	}
	if len(mergedKeys[0].PrimaryKeys) != 1 || mergedKeys[0].PrimaryKeys[0] != "10,7" {
		t.Fatalf("unexpected merged keys %#v", mergedKeys[0].PrimaryKeys)
	}
}

func TestBatchEntryPrimaryKeyStringsForMergedInsert(t *testing.T) {
	insert1 := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, v) VALUES (?, ?, ?)")
	insert2 := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, v) VALUES (?, ?, ?)")
	groups := coalesceBatchInsertGroups([]BatchEntry{
		{Statement: insert1, Params: []interface{}{10, 7, "a"}},
		{Statement: insert2, Params: []interface{}{11, 7, "b"}},
	}, 64, 256*1024)

	if len(groups) != 1 {
		t.Fatalf("expected 1 merged insert group, got %d", len(groups))
	}

	keys := batchEntryPrimaryKeyStrings(groups[0].merged, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	})
	expected := []string{
		"ks.tbl(id=10, ck=7)",
		"ks.tbl(id=11, ck=7)",
	}
	if len(keys) != len(expected) {
		t.Fatalf("expected %d keys, got %d (%#v)", len(expected), len(keys), keys)
	}
	for i := range expected {
		if keys[i] != expected[i] {
			t.Fatalf("unexpected key %d: got=%q want=%q", i, keys[i], expected[i])
		}
	}
}

func TestCollectBatchPrimaryKeyDebugInfoIncludesPartitionGroupingPlan(t *testing.T) {
	insert1 := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, v) VALUES (?, ?, ?)")
	insert2 := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, v) VALUES (?, ?, ?)")
	resolvePK := func(table *parser.TableName) ([]string, error) {
		return []string{"id", "ck"}, nil
	}
	resolveSplitKeys := func(entry BatchEntry, insert *parser.InsertStatement, params []interface{}) []string {
		switch params[0] {
		case 10:
			return []string{"ks.tbl#hash:0"}
		case 11:
			return []string{"ks.tbl#hash:1"}
		default:
			return nil
		}
	}

	groups := coalesceBatchInsertGroupsWithResolver([]BatchEntry{
		{Statement: insert1, Params: []interface{}{10, 7, "a"}},
		{Statement: insert2, Params: []interface{}{11, 7, "b"}},
	}, 64, 256*1024, resolvePK, resolveSplitKeys)
	if len(groups) != 1 {
		t.Fatalf("expected 1 merged insert group, got %d", len(groups))
	}
	if groups[0].groupingStrategy != batchGroupingStrategyPartitionGroup {
		t.Fatalf("expected grouping strategy %q, got %q", batchGroupingStrategyPartitionGroup, groups[0].groupingStrategy)
	}
	expectedPartitions := []string{"ks.tbl#hash:0", "ks.tbl#hash:1"}
	if len(groups[0].partitionHits) != len(expectedPartitions) {
		t.Fatalf("expected %d partition hits, got %d (%#v)", len(expectedPartitions), len(groups[0].partitionHits), groups[0].partitionHits)
	}
	for i := range expectedPartitions {
		if groups[0].partitionHits[i] != expectedPartitions[i] {
			t.Fatalf("unexpected partition hit %d: got=%q want=%q", i, groups[0].partitionHits[i], expectedPartitions[i])
		}
	}

	_, _, _, mergedKeys := collectBatchPrimaryKeyDebugInfo(groups, resolvePK)
	if len(mergedKeys) != 1 {
		t.Fatalf("expected 1 merged key group, got %d", len(mergedKeys))
	}
	if mergedKeys[0].GroupingStrategy != string(batchGroupingStrategyPartitionGroup) {
		t.Fatalf("expected merged key grouping strategy %q, got %q", batchGroupingStrategyPartitionGroup, mergedKeys[0].GroupingStrategy)
	}
	if mergedKeys[0].BatchTemplateKey != "insert:ks.tbl|cols=id,ck,v" {
		t.Fatalf("unexpected batch template key: got=%q", mergedKeys[0].BatchTemplateKey)
	}
	if mergedKeys[0].GroupBreakReason != batchGroupBreakReasonEndOfEntries {
		t.Fatalf("expected group break reason %q, got %q", batchGroupBreakReasonEndOfEntries, mergedKeys[0].GroupBreakReason)
	}
	if len(mergedKeys[0].PartitionHits) != len(expectedPartitions) {
		t.Fatalf("expected %d merged partition hits, got %d (%#v)", len(expectedPartitions), len(mergedKeys[0].PartitionHits), mergedKeys[0].PartitionHits)
	}
	for i := range expectedPartitions {
		if mergedKeys[0].PartitionHits[i] != expectedPartitions[i] {
			t.Fatalf("unexpected merged partition hit %d: got=%q want=%q", i, mergedKeys[0].PartitionHits[i], expectedPartitions[i])
		}
	}
}

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

	groups := coalesceBatchInsertGroupsWithResolver([]BatchEntry{
		{Statement: insert, Params: []interface{}{1, "a1"}},
		{Statement: insert, Params: []interface{}{1, "a2"}},
		{Statement: insert, Params: []interface{}{2, "b1"}},
	}, 1, 256*1024, resolvePK)
	if len(groups) != 2 {
		t.Fatalf("expected 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 mergedKeys[0].GroupBreakReason != batchGroupBreakReasonDuplicateFoldSingletonUnabsorbed {
		t.Fatalf("expected first group break reason %q, got %q", batchGroupBreakReasonDuplicateFoldSingletonUnabsorbed, mergedKeys[0].GroupBreakReason)
	}
	if mergedKeys[0].BatchTemplateKey != "insert:ks.tbl|cols=id,v" {
		t.Fatalf("unexpected first batch template key: got=%q", mergedKeys[0].BatchTemplateKey)
	}
}