package executor

import (
	"testing"

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

func TestCoalesceBatchInsertGroupsWithResolverKeepsSameHashBucketInSameGroupAcrossLimits(t *testing.T) {
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	entries := []BatchEntry{
		{Statement: insert, Params: []interface{}{1, "a"}},
		{Statement: insert, Params: []interface{}{2, "b"}},
		{Statement: insert, Params: []interface{}{11, "c"}},
	}

	groups := coalesceBatchInsertGroupsWithResolver(entries, 2, 256*1024,
		func(table *parser.TableName) ([]string, error) {
			return []string{"id"}, nil
		},
		func(entry BatchEntry, insert *parser.InsertStatement, params []interface{}) []string {
			id, _ := params[0].(int)
			if id == 1 || id == 11 {
				return []string{"bucket:1"}
			}
			return []string{"bucket:2"}
		},
	)
	if len(groups) != 1 {
		t.Fatalf("expected same HASH bucket rows to stay in one group, 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 rows, got %d", len(stmt.Values))
	}
}

func TestCoalesceBatchInsertGroupsWithResolverFallsBackWhenHashBucketUnavailable(t *testing.T) {
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	entries := []BatchEntry{
		{Statement: insert, Params: []interface{}{1, "a"}},
		{Statement: insert, Params: []interface{}{2, "b"}},
		{Statement: insert, Params: []interface{}{11, "c"}},
	}

	groups := coalesceBatchInsertGroupsWithResolver(entries, 2, 256*1024,
		func(table *parser.TableName) ([]string, error) {
			return []string{"id"}, nil
		},
		func(entry BatchEntry, insert *parser.InsertStatement, params []interface{}) []string {
			id, _ := params[0].(int)
			if id == 2 {
				return nil
			}
			return []string{"bucket:1"}
		},
	)
	if len(groups) != 2 {
		t.Fatalf("expected missing HASH bucket info to fall back to PK splitting, got %d groups", len(groups))
	}
}

func TestTryCoalesceHomogeneousUnloggedBatchGroupsChunksHashBucketsByRowLimit(t *testing.T) {
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	entries := []BatchEntry{
		{Statement: insert, Params: []interface{}{1, "a"}},
		{Statement: insert, Params: []interface{}{11, "b"}},
		{Statement: insert, Params: []interface{}{21, "c"}},
	}
	resolvePK := func(table *parser.TableName) ([]string, error) {
		return []string{"id"}, nil
	}
	resolveSplitKeys := func(entry BatchEntry, insert *parser.InsertStatement, params []interface{}) []string {
		return []string{"bucket:1"}
	}
	planning := testBatchPlanningEnvironment(resolvePK, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "v"}, nil
	}, resolveSplitKeys)

	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 strict row chunking to split HASH bucket group, got %d groups", len(groups))
	}
	for i, group := range groups {
		if group.groupingStrategy != batchGroupingStrategyPartitionGroup {
			t.Fatalf("expected group %d to use partition grouping, got %q", i, group.groupingStrategy)
		}
		if len(group.partitionHits) != 1 || group.partitionHits[0] != "bucket:1" {
			t.Fatalf("unexpected partition hits for group %d: %#v", i, group.partitionHits)
		}
	}
	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 TestTryCoalesceHomogeneousUnloggedBatchGroupsRegroupsInterleavedHashBuckets(t *testing.T) {
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	entries := []BatchEntry{
		{Statement: insert, Params: []interface{}{1, "a"}},
		{Statement: insert, Params: []interface{}{2, "b"}},
		{Statement: insert, Params: []interface{}{11, "c"}},
	}
	resolvePK := func(table *parser.TableName) ([]string, error) {
		return []string{"id"}, nil
	}
	resolveSplitKeys := func(entry BatchEntry, insert *parser.InsertStatement, params []interface{}) []string {
		idIndex := -1
		for i, col := range insert.Columns {
			if col == "id" {
				idIndex = i
				break
			}
		}
		if idIndex < 0 || idIndex >= len(params) {
			return []string{"bucket:unknown"}
		}
		var id int64
		switch v := params[idIndex].(type) {
		case int:
			id = int64(v)
		case int64:
			id = v
		}
		if id == 1 || id == 11 {
			return []string{"bucket:1"}
		}
		return []string{"bucket:2"}
	}
	planning := testBatchPlanningEnvironment(resolvePK, func(table *parser.TableName) ([]string, error) {
		return []string{"id", "v"}, nil
	}, resolveSplitKeys)

	groups, ok, err := tryCoalesceHomogeneousUnloggedBatchGroupsWithPlanningEnvironment(entries, 10, 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 interleaved HASH buckets to regroup into 2 groups, 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 regrouped bucket chunk sizes 2 and 1, got %d and %d", len(stmt1.Values), len(stmt2.Values))
	}
	if len(groups[0].partitionHits) != 1 || groups[0].partitionHits[0] != "bucket:1" {
		t.Fatalf("unexpected first group partition hits: %#v", groups[0].partitionHits)
	}
	if len(groups[1].partitionHits) != 1 || groups[1].partitionHits[0] != "bucket:2" {
		t.Fatalf("unexpected second group partition hits: %#v", groups[1].partitionHits)
	}
}

func TestParseBatchHashPartitionHintFromShowCreate(t *testing.T) {
	hint, ok := parseBatchHashPartitionHintFromShowCreate("CREATE TABLE `ks`.`tbl` (`id` int NOT NULL) PARTITION BY HASH (`id`) PARTITIONS 10")
	if !ok {
		t.Fatalf("expected HASH partition hint to parse")
	}
	if hint.column != "id" {
		t.Fatalf("expected HASH partition column id, got %q", hint.column)
	}
	if hint.partitions != 10 {
		t.Fatalf("expected 10 partitions, got %d", hint.partitions)
	}

	if _, ok := parseBatchHashPartitionHintFromShowCreate("CREATE TABLE `ks`.`tbl` (`id` int NOT NULL) PARTITION BY KEY(`id`) PARTITIONS 10"); ok {
		t.Fatalf("expected KEY partition clause not to parse as HASH")
	}
}

func TestBatchHashPartitionBucketKey(t *testing.T) {
	stmt := mustParseStatement(t, "INSERT INTO ks.tbl (id, v) VALUES (?, ?)")
	insert, ok := stmt.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected INSERT statement, got %T", stmt)
	}

	key, ok := batchHashPartitionBucketKey(insert, []interface{}{11, "x"}, batchHashPartitionHint{
		column:     "id",
		partitions: 10,
	})
	if !ok {
		t.Fatalf("expected HASH bucket key to resolve")
	}
	if key != "ks.tbl#hash:1" {
		t.Fatalf("expected bucket key ks.tbl#hash:1, got %q", key)
	}
}

func TestSingleHashPartitionKeyColumn(t *testing.T) {
	column, ok := singleHashPartitionKeyColumn([]mysqlbackend.CQLColumnInfo{
		{Name: "id", Kind: "partition_key", CQLType: "int", Position: 0},
		{Name: "ck", Kind: "clustering", CQLType: "int", Position: 0},
	})
	if !ok {
		t.Fatalf("expected single HASH partition key to be accepted")
	}
	if column != "id" {
		t.Fatalf("expected partition key column id, got %q", column)
	}

	if _, ok := singleHashPartitionKeyColumn([]mysqlbackend.CQLColumnInfo{
		{Name: "a", Kind: "partition_key", CQLType: "int", Position: 0},
		{Name: "b", Kind: "partition_key", CQLType: "int", Position: 1},
	}); ok {
		t.Fatalf("expected composite partition key to be rejected for HASH-only grouping")
	}

	if _, ok := singleHashPartitionKeyColumn([]mysqlbackend.CQLColumnInfo{
		{Name: "id", Kind: "partition_key", CQLType: "varint", Position: 0},
	}); ok {
		t.Fatalf("expected non-HASH-friendly partition key type to be rejected")
	}
}