package executor

import (
	"reflect"
	"testing"

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

func testBatchFoldEnv(pkCols, schemaOrder []string, insertIncrements map[string]struct{}) batchFoldEnvironment {
	resolvePK := func(table *parser.TableName) ([]string, error) {
		return append([]string(nil), pkCols...), nil
	}
	resolveColumnOrder := func(table *parser.TableName) ([]string, error) {
		return append([]string(nil), schemaOrder...), nil
	}
	return batchFoldEnvironment{
		resolvePK:          resolvePK,
		resolveColumnOrder: resolveColumnOrder,
		lookupColumnTypes: func(table *parser.TableName) (map[string]string, error) {
			return nil, nil
		},
		shouldInsertIncrement: func(table *parser.TableName, column string) bool {
			if len(insertIncrements) == 0 {
				return false
			}
			_, ok := insertIncrements[normalizeBatchColumnName(column)]
			return ok
		},
	}
}

func requireFoldedInsertEntry(t *testing.T, entry BatchEntry) *parser.InsertStatement {
	t.Helper()
	stmt, ok := entry.Statement.(*parser.InsertStatement)
	if !ok {
		t.Fatalf("expected folded entry to be INSERT, got %T", entry.Statement)
	}
	if len(stmt.Values) != 1 {
		t.Fatalf("expected folded entry to contain 1 row, got %d", len(stmt.Values))
	}
	if len(stmt.Columns) != len(entry.Params) {
		t.Fatalf("expected params to match columns, cols=%d params=%d", len(stmt.Columns), len(entry.Params))
	}
	return stmt
}

func foldedRowByColumn(t *testing.T, entry BatchEntry) map[string]interface{} {
	t.Helper()
	stmt := requireFoldedInsertEntry(t, entry)
	row := make(map[string]interface{}, len(stmt.Columns))
	for i, col := range stmt.Columns {
		row[normalizeBatchColumnName(col)] = entry.Params[i]
	}
	return row
}

func TestFoldBatchEntriesSamePKInsertDifferentColumnsCombinesData(t *testing.T) {
	env := testBatchFoldEnv([]string{"id", "ck"}, []string{"id", "ck", "note", "content"}, nil)
	insertNote := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, note) VALUES (?, ?, ?)")
	insertContent := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, content) VALUES (?, ?, ?)")
	entries := []BatchEntry{
		{Statement: insertNote, Params: []interface{}{1, 2, "hello"}},
		{Statement: insertContent, Params: []interface{}{1, 2, "world"}},
	}

	folded, err := foldBatchEntriesWithEnv(entries, false, env)
	if err != nil {
		t.Fatalf("foldBatchEntriesWithEnv failed: %v", err)
	}
	if len(folded) != 1 {
		t.Fatalf("expected 1 folded entry, got %d", len(folded))
	}

	row := foldedRowByColumn(t, folded[0])
	if !reflect.DeepEqual(row, map[string]interface{}{
		"note":    "hello",
		"content": "world",
		"id":      1,
		"ck":      2,
	}) {
		t.Fatalf("unexpected folded row: %#v", row)
	}
	if len(folded[0].originals) != 2 {
		t.Fatalf("expected folded entry to keep 2 originals, got %d", len(folded[0].originals))
	}
}

func TestFoldBatchEntriesSamePKInsertThenUpdateCombinesData(t *testing.T) {
	env := testBatchFoldEnv([]string{"id", "ck"}, []string{"id", "ck", "note", "content"}, nil)
	insertNote := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, note) VALUES (?, ?, ?)")
	updateContent := mustParseStatement(t, "UPDATE ks.tbl SET content = ? WHERE id = ? AND ck = ?")
	entries := []BatchEntry{
		{Statement: insertNote, Params: []interface{}{1, 2, "hello"}},
		{Statement: updateContent, Params: []interface{}{"world", 1, 2}},
	}

	folded, err := foldBatchEntriesWithEnv(entries, true, env)
	if err != nil {
		t.Fatalf("foldBatchEntriesWithEnv failed: %v", err)
	}
	if len(folded) != 1 {
		t.Fatalf("expected 1 folded entry, got %d", len(folded))
	}
	if folded[0].mergeKind != batchMergeKindUpdate {
		t.Fatalf("expected mixed insert/update fold to keep update merge kind, got %d", folded[0].mergeKind)
	}

	row := foldedRowByColumn(t, folded[0])
	if !reflect.DeepEqual(row, map[string]interface{}{
		"note":    "hello",
		"content": "world",
		"id":      1,
		"ck":      2,
	}) {
		t.Fatalf("unexpected folded row: %#v", row)
	}
}

func TestFoldBatchEntriesPlainThenDeltaBecomesSet(t *testing.T) {
	env := testBatchFoldEnv([]string{"id", "ck"}, []string{"id", "ck", "unread_count"}, nil)
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, unread_count) VALUES (?, ?, ?)")
	update := mustParseStatement(t, "UPDATE ks.tbl SET unread_count = unread_count + ? WHERE id = ? AND ck = ?")
	entries := []BatchEntry{
		{Statement: insert, Params: []interface{}{1, 2, 5}},
		{Statement: update, Params: []interface{}{3, 1, 2}},
	}

	folded, err := foldBatchEntriesWithEnv(entries, true, env)
	if err != nil {
		t.Fatalf("foldBatchEntriesWithEnv failed: %v", err)
	}
	if len(folded) != 1 {
		t.Fatalf("expected 1 folded entry, got %d", len(folded))
	}

	stmt := requireFoldedInsertEntry(t, folded[0])
	if len(stmt.OnDuplicateIncrementColumns) != 0 {
		t.Fatalf("expected plain+delta fold to materialize as absolute SET, got increment hints %#v", stmt.OnDuplicateIncrementColumns)
	}
	row := foldedRowByColumn(t, folded[0])
	if unread := row["unread_count"]; unread != int64(8) && unread != 8 {
		t.Fatalf("expected unread_count=8 after fold, got %#v", row["unread_count"])
	}
}

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

	folded, err := foldBatchEntriesWithEnv(entries, true, env)
	if err != nil {
		t.Fatalf("foldBatchEntriesWithEnv failed: %v", err)
	}
	if len(folded) != 1 {
		t.Fatalf("expected 1 folded entry, got %d", len(folded))
	}

	stmt := requireFoldedInsertEntry(t, folded[0])
	if _, ok := stmt.OnDuplicateIncrementColumns["unread_count"]; !ok {
		t.Fatalf("expected unread_count to remain delta after fold")
	}
	row := foldedRowByColumn(t, folded[0])
	if unread := row["unread_count"]; unread != int64(8) && unread != 8 {
		t.Fatalf("expected unread_count delta=8 after fold, got %#v", row["unread_count"])
	}
}

func TestFoldBatchEntriesUnsetSkippedAndNullPreserved(t *testing.T) {
	env := testBatchFoldEnv([]string{"id"}, []string{"id", "name", "idx", "note"}, nil)
	updateMixed := mustParseStatement(t, "UPDATE ks.tbl SET name = ?, idx = ? WHERE id = ?")
	updateNote := mustParseStatement(t, "UPDATE ks.tbl SET note = ? WHERE id = ?")
	entries := []BatchEntry{
		{Statement: updateMixed, Params: []interface{}{UnsetValue, 2, 1}},
		{Statement: updateNote, Params: []interface{}{nil, 1}},
	}

	folded, err := foldBatchEntriesWithEnv(entries, true, env)
	if err != nil {
		t.Fatalf("foldBatchEntriesWithEnv failed: %v", err)
	}
	if len(folded) != 1 {
		t.Fatalf("expected 1 folded entry, got %d", len(folded))
	}

	row := foldedRowByColumn(t, folded[0])
	if _, ok := row["name"]; ok {
		t.Fatalf("expected UNSET name to be omitted, got %#v", row)
	}
	if row["idx"] != 2 {
		t.Fatalf("expected idx=2, got %#v", row["idx"])
	}
	if value, ok := row["note"]; !ok || value != nil {
		t.Fatalf("expected note=NULL to be preserved, got %#v", row["note"])
	}
}

func TestFoldBatchEntriesConfiguredInsertIncrementStaysSeparate(t *testing.T) {
	env := testBatchFoldEnv([]string{"id"}, []string{"id", "unread_count"}, map[string]struct{}{"unread_count": {}})
	insert1 := mustParseStatement(t, "INSERT INTO ks.tbl (id, unread_count) VALUES (?, ?)")
	insert2 := mustParseStatement(t, "INSERT INTO ks.tbl (id, unread_count) VALUES (?, ?)")
	entries := []BatchEntry{
		{Statement: insert1, Params: []interface{}{1, 2}},
		{Statement: insert2, Params: []interface{}{1, 3}},
	}

	folded, err := foldBatchEntriesWithEnv(entries, false, env)
	if err != nil {
		t.Fatalf("foldBatchEntriesWithEnv failed: %v", err)
	}
	if len(folded) != 2 {
		t.Fatalf("expected configured increment inserts to stay separate, got %d entries", len(folded))
	}
	for i := range folded {
		if folded[i].folded {
			t.Fatalf("expected entry %d to remain raw when insert duplicate increment is configured", i)
		}
	}
}

func TestFoldBatchEntriesUsesSchemaColumnOrder(t *testing.T) {
	env := testBatchFoldEnv([]string{"id", "ck"}, []string{"id", "ck", "content", "note"}, nil)
	insertNote := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, note) VALUES (?, ?, ?)")
	insertContent := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, content) VALUES (?, ?, ?)")
	entries := []BatchEntry{
		{Statement: insertNote, Params: []interface{}{1, 2, "hello"}},
		{Statement: insertContent, Params: []interface{}{1, 2, "world"}},
	}

	folded, err := foldBatchEntriesWithEnv(entries, false, env)
	if err != nil {
		t.Fatalf("foldBatchEntriesWithEnv failed: %v", err)
	}
	if len(folded) != 1 {
		t.Fatalf("expected 1 folded entry, got %d", len(folded))
	}

	stmt := requireFoldedInsertEntry(t, folded[0])
	expected := []string{"content", "note", "id", "ck"}
	if !reflect.DeepEqual(stmt.Columns, expected) {
		t.Fatalf("unexpected folded column order: got=%v want=%v", stmt.Columns, expected)
	}
}

func TestFoldBatchEntriesBarrierFlushesPendingState(t *testing.T) {
	env := testBatchFoldEnv([]string{"id"}, []string{"id", "name", "idx"}, nil)
	insertName := mustParseStatement(t, "INSERT INTO ks.tbl (id, name) VALUES (?, ?)")
	deleteStmt := mustParseStatement(t, "DELETE FROM ks.tbl WHERE id = ?")
	insertIdx := mustParseStatement(t, "INSERT INTO ks.tbl (id, idx) VALUES (?, ?)")
	entries := []BatchEntry{
		{Statement: insertName, Params: []interface{}{1, "alpha"}},
		{Statement: deleteStmt, Params: []interface{}{9}},
		{Statement: insertIdx, Params: []interface{}{1, 2}},
	}

	folded, err := foldBatchEntriesWithEnv(entries, false, env)
	if err != nil {
		t.Fatalf("foldBatchEntriesWithEnv failed: %v", err)
	}
	if len(folded) != 3 {
		t.Fatalf("expected barrier to flush into 3 entries, got %d", len(folded))
	}
	if !folded[0].folded {
		t.Fatalf("expected first entry to be folded before barrier")
	}
	if _, ok := folded[1].Statement.(*parser.DeleteStatement); !ok {
		t.Fatalf("expected middle entry to remain DELETE barrier, got %T", folded[1].Statement)
	}
	if !folded[2].folded {
		t.Fatalf("expected final entry after barrier to fold independently")
	}
}

func TestFoldBatchEntriesInsertPrimaryKeyLookupErrorFallsBackRaw(t *testing.T) {
	insert := mustParseStatement(t, "INSERT INTO ks.tbl (id, note) VALUES (?, ?)")
	env := batchFoldEnvironment{
		resolvePK: func(table *parser.TableName) ([]string, error) {
			return nil, assertErr("pk lookup failed")
		},
		resolveColumnOrder: func(table *parser.TableName) ([]string, error) {
			return []string{"id", "note"}, nil
		},
		lookupColumnTypes: func(table *parser.TableName) (map[string]string, error) {
			return nil, nil
		},
	}

	folded, err := foldBatchEntriesWithEnv([]BatchEntry{{Statement: insert, Params: []interface{}{1, "hello"}}}, false, env)
	if err != nil {
		t.Fatalf("foldBatchEntriesWithEnv failed: %v", err)
	}
	if len(folded) != 1 {
		t.Fatalf("expected 1 entry, got %d", len(folded))
	}
	if folded[0].folded {
		t.Fatalf("expected entry to fall back to raw execution on pk lookup error")
	}
}

func TestFoldBatchEntriesColumnOrderLookupErrorStillFolds(t *testing.T) {
	insertNote := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, note) VALUES (?, ?, ?)")
	insertContent := mustParseStatement(t, "INSERT INTO ks.tbl (id, ck, content) VALUES (?, ?, ?)")
	env := batchFoldEnvironment{
		resolvePK: func(table *parser.TableName) ([]string, error) {
			return []string{"id", "ck"}, nil
		},
		resolveColumnOrder: func(table *parser.TableName) ([]string, error) {
			return nil, assertErr("column order unavailable")
		},
		lookupColumnTypes: func(table *parser.TableName) (map[string]string, error) {
			return nil, nil
		},
	}

	folded, err := foldBatchEntriesWithEnv([]BatchEntry{
		{Statement: insertNote, Params: []interface{}{1, 2, "hello"}},
		{Statement: insertContent, Params: []interface{}{1, 2, "world"}},
	}, false, env)
	if err != nil {
		t.Fatalf("foldBatchEntriesWithEnv failed: %v", err)
	}
	if len(folded) != 1 {
		t.Fatalf("expected 1 folded entry, got %d", len(folded))
	}
	stmt := requireFoldedInsertEntry(t, folded[0])
	expected := []string{"note", "content", "id", "ck"}
	if !reflect.DeepEqual(stmt.Columns, expected) {
		t.Fatalf("unexpected fallback column order: got=%v want=%v", stmt.Columns, expected)
	}
}

type batchFoldTestError string

func (e batchFoldTestError) Error() string { return string(e) }

func assertErr(msg string) error {
	return batchFoldTestError(msg)
}