package migrationprecheck

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"

	globaldedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/entity"
	mockclusterproxy "github.com/tidbcloud/tidb-management-service/utilities/clusterproxy/mock_clusterproxy"
)

func TestCompatibilityCheckCatalog(t *testing.T) {
	expected := []string{
		CheckDuplicateCheckConstraints,
		CheckMySQLDBCollationCompatibility,
		CheckInvalidForeignKey,
		CheckTiDBVersion,
		CheckTiDBGCLifeTime,
		CheckTablesWithoutValidIndex,
		CheckSequence,
		CheckNewCollationEnabled,
		CheckTTL,
		CheckTiFlashReplicas,
	}
	require.Len(t, checkDefinitions, len(expected))
	for _, key := range expected {
		definition, ok := checkDefinitions[key]
		require.True(t, ok, key)
		require.NotNil(t, definition.run)
	}
}

func TestFirstScalarDecodesBase64(t *testing.T) {
	raw := queryResult(t, map[string][]string{"value": {"72h0m0s"}})
	value, err := firstScalar(raw)
	require.NoError(t, err)
	require.Equal(t, "72h0m0s", value)
}

func TestEvaluateTiDBGCLifeTime(t *testing.T) {
	tests := []struct {
		name   string
		value  string
		status globaldedicatedv1.MigrationCompatibilityCheckStatus
		action string
	}{
		{
			name: "below minimum", value: "71h59m59s",
			status: globaldedicatedv1.MigrationCompatibilityCheckStatus_RISK_FOUND,
			action: "Set tidb_gc_life_time system variable to at least 72h before starting the migration backup.",
		},
		{
			name: "at minimum", value: "72h0m0s",
			status: globaldedicatedv1.MigrationCompatibilityCheckStatus_NO_RISK_FOUND,
			action: noActionRequired,
		},
		{
			name: "unparseable", value: "unknown",
			status: globaldedicatedv1.MigrationCompatibilityCheckStatus_NEEDS_INPUT,
			action: "Confirm tidb_gc_life_time manually and rerun the check.",
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			result, err := evaluateTiDBGCLifeTime(test.value)
			require.NoError(t, err)
			require.Equal(t, test.status, result.status)
			require.Equal(t, test.action, result.action)
		})
	}
}

func TestEvaluateTiDBVersion(t *testing.T) {
	tests := []struct {
		name       string
		value      string
		status     globaldedicatedv1.MigrationCompatibilityCheckStatus
		action     string
		wantAction bool
		wantErr    bool
	}{
		{
			name: "new architecture", value: "8.0.11-TiDB-v7.5.5",
			status: globaldedicatedv1.MigrationCompatibilityCheckStatus_NO_RISK_FOUND,
		},
		{
			name: "old architecture", value: "5.7.25-TiDB-v7.4.0",
			status: globaldedicatedv1.MigrationCompatibilityCheckStatus_NO_RISK_FOUND,
			action: "use the TiCDC old architecture", wantAction: true,
		},
		{
			name: "unknown version", value: "unknown",
			wantErr: true,
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			result, err := evaluateTiDBVersion(test.value)
			if test.wantErr {
				require.Error(t, err)
				return
			}
			require.NoError(t, err)
			require.Equal(t, test.status, result.status)
			if test.wantAction {
				require.Contains(t, result.action, test.action)
			} else {
				require.Equal(t, noActionRequired, result.action)
			}
		})
	}
}

func TestRunTiFlashReplicas(t *testing.T) {
	t.Run("returns node and replica table counts", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		proxy := mockclusterproxy.NewMockClusterProxy(ctrl)
		proxy.EXPECT().QuerySQL(gomock.Any(), tiFlashReplicaTableCountSQL).Return(
			queryResult(t, map[string][]string{"tables_with_tiflash_replicas": {"3"}}), nil,
		)
		firstTiFlashCount := uint32(2)
		secondTiFlashCount := uint32(1)
		tiKVCount := uint32(3)
		cluster := &entity.Cluster{Cluster: &globaldedicatedv1.Cluster{ClusterComponents: []*globaldedicatedv1.Cluster_ClusterComponent{
			{ComponentType: globaldedicatedv1.ComponentType_TIFLASH, Count: &firstTiFlashCount},
			{ComponentType: globaldedicatedv1.ComponentType_TIKV, Count: &tiKVCount},
			{ComponentType: globaldedicatedv1.ComponentType_TIFLASH, Count: &secondTiFlashCount},
		}}}

		result, err := runTiFlashReplicas(context.Background(), proxy, cluster)
		require.NoError(t, err)
		require.Equal(t, globaldedicatedv1.MigrationCompatibilityCheckStatus_NO_RISK_FOUND, result.status)
		require.Equal(t, "TiFlash nodes: 3; tables with TiFlash replicas: 3.", result.evidence)
		require.Equal(t, noActionRequired, result.action)
	})

	t.Run("rejects an invalid replica table count", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		proxy := mockclusterproxy.NewMockClusterProxy(ctrl)
		proxy.EXPECT().QuerySQL(gomock.Any(), tiFlashReplicaTableCountSQL).Return(
			queryResult(t, map[string][]string{"tables_with_tiflash_replicas": {"unknown"}}), nil,
		)
		cluster := &entity.Cluster{Cluster: &globaldedicatedv1.Cluster{}}

		_, err := runTiFlashReplicas(context.Background(), proxy, cluster)
		require.Error(t, err)
	})
}

func TestRunTTLCheck(t *testing.T) {
	tests := []struct {
		name           string
		globalValue    string
		candidateNames [][]string
		createTables   []string
		status         globaldedicatedv1.MigrationCompatibilityCheckStatus
		evidence       string
		action         string
	}{
		{
			name:           "global off with regular and partitioned TTL tables enabled",
			globalValue:    "0",
			candidateNames: [][]string{{"app", "ttl_orders"}, {"app", "ttl_partitioned"}},
			createTables: []string{
				"CREATE TABLE `ttl_orders` (`id` BIGINT PRIMARY KEY, `created_at` DATETIME) TTL=`created_at` + INTERVAL 1 DAY TTL_ENABLE='ON'",
				"CREATE TABLE `ttl_partitioned` (`id` BIGINT, `created_at` DATETIME, PRIMARY KEY (`id`, `created_at`)) TTL=`created_at` + INTERVAL 1 DAY TTL_ENABLE = 'ON' PARTITION BY RANGE COLUMNS (`created_at`) (PARTITION `p0` VALUES LESS THAN (MAXVALUE))",
			},
			status:   globaldedicatedv1.MigrationCompatibilityCheckStatus_NO_RISK_FOUND,
			evidence: "The global variable tidb_ttl_job_enable is OFF, and table-level TTL is enabled on the following tables: app.ttl_orders, app.ttl_partitioned.",
			action:   noActionRequired,
		},
		{
			name:           "global off without table-level TTL enabled",
			globalValue:    "OFF",
			candidateNames: [][]string{{"app", "ttl_disabled"}},
			createTables: []string{
				"CREATE TABLE `ttl_disabled` (`id` BIGINT PRIMARY KEY, `created_at` DATETIME) TTL=`created_at` + INTERVAL 1 DAY TTL_ENABLE='OFF'",
			},
			status:   globaldedicatedv1.MigrationCompatibilityCheckStatus_NO_RISK_FOUND,
			evidence: "The global variable tidb_ttl_job_enable is OFF,  and table-level TTL is not enabled on any table.",
			action:   noActionRequired,
		},
		{
			name:           "global on with table-level TTL enabled",
			globalValue:    "1",
			candidateNames: [][]string{{"app", "ttl_orders"}},
			createTables: []string{
				"CREATE TABLE `ttl_orders` (`id` BIGINT PRIMARY KEY, `created_at` DATETIME) TTL=`created_at` + INTERVAL 1 DAY TTL_ENABLE='ON'",
			},
			status:   globaldedicatedv1.MigrationCompatibilityCheckStatus_RISK_FOUND,
			evidence: "The global variable tidb_ttl_job_enable is ON, and table-level TTL is enabled on the following tables: app.ttl_orders.",
			action:   ttlRiskAction,
		},
		{
			name:        "global on without TTL tables",
			globalValue: "ON",
			status:      globaldedicatedv1.MigrationCompatibilityCheckStatus_NO_RISK_FOUND,
			evidence:    "The global variable tidb_ttl_job_enable is ON, but table-level TTL is not enabled on any table.",
			action:      noActionRequired,
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			ctrl := gomock.NewController(t)
			proxy := mockclusterproxy.NewMockClusterProxy(ctrl)
			calls := []any{
				proxy.EXPECT().QuerySQL(gomock.Any(), ttlJobEnabledSQL).Return(
					queryResult(t, map[string][]string{"tidb_ttl_job_enable": {test.globalValue}}), nil,
				),
				proxy.EXPECT().QuerySQL(gomock.Any(), ttlCandidateTablesSQL).Return(
					queryResult(t, map[string][]string{
						"table_schema": column(test.candidateNames, 0),
						"table_name":   column(test.candidateNames, 1),
					}), nil,
				),
			}
			for i, name := range test.candidateNames {
				calls = append(calls, proxy.EXPECT().QuerySQL(
					gomock.Any(), fmt.Sprintf("SHOW CREATE TABLE `%s`.`%s`", name[0], name[1]),
				).Return(queryResult(t, map[string][]string{
					"Table":        {name[1]},
					"Create Table": {test.createTables[i]},
				}), nil))
			}
			gomock.InOrder(calls...)

			result, err := runTTLCheck(context.Background(), proxy)
			require.NoError(t, err)
			require.Equal(t, test.status, result.status)
			require.Equal(t, test.evidence, result.evidence)
			require.Equal(t, test.action, result.action)
		})
	}
}

func column(rows [][]string, index int) []string {
	values := make([]string, 0, len(rows))
	for _, row := range rows {
		values = append(values, row[index])
	}
	return values
}

func TestRunDuplicateCheckConstraints(t *testing.T) {
	t.Run("TiDB before v7.4 skips information schema query", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		proxy := mockclusterproxy.NewMockClusterProxy(ctrl)
		proxy.EXPECT().QuerySQL(gomock.Any(), tidbVersionSQL).Return(
			queryResult(t, map[string][]string{"tidb_version": {"5.7.25-TiDB-v6.5.0"}}), nil,
		)

		result, err := runDuplicateCheckConstraints(context.Background(), proxy)
		require.NoError(t, err)
		require.Equal(t, globaldedicatedv1.MigrationCompatibilityCheckStatus_NO_RISK_FOUND, result.status)
		require.Equal(t, "The information_schema.CHECK_CONSTRAINTS table is added starting with v7.4.0.", result.evidence)
		require.Equal(t, noActionRequired, result.action)
	})

	t.Run("TiDB v7.4 or later checks duplicate constraints", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		proxy := mockclusterproxy.NewMockClusterProxy(ctrl)
		gomock.InOrder(
			proxy.EXPECT().QuerySQL(gomock.Any(), tidbVersionSQL).Return(
				queryResult(t, map[string][]string{"tidb_version": {"5.7.25-TiDB-v7.4.0"}}), nil,
			),
			proxy.EXPECT().QuerySQL(gomock.Any(), duplicateCheckConstraintsSQL).Return(queryResult(t, map[string][]string{
				"constraint_schema": {"app"},
				"constraint_name":   {"chk_id"},
				"total_count":       {"1"},
			}), nil),
		)

		result, err := runDuplicateCheckConstraints(context.Background(), proxy)
		require.NoError(t, err)
		require.Equal(t, globaldedicatedv1.MigrationCompatibilityCheckStatus_RISK_FOUND, result.status)
		require.Contains(t, result.evidence, "database app, constraint name chk_id")
	})
}

func TestLargeCompatibilityEvidenceIsSummarized(t *testing.T) {
	t.Run("duplicate check constraints", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		proxy := mockclusterproxy.NewMockClusterProxy(ctrl)
		schemas := make([]string, maximumEvidenceItems)
		constraints := make([]string, maximumEvidenceItems)
		for i := range maximumEvidenceItems {
			schemas[i] = "app"
			constraints[i] = fmt.Sprintf("chk_%02d", i)
		}
		gomock.InOrder(
			proxy.EXPECT().QuerySQL(gomock.Any(), tidbVersionSQL).Return(
				queryResult(t, map[string][]string{"tidb_version": {"5.7.25-TiDB-v7.4.0"}}), nil,
			),
			proxy.EXPECT().QuerySQL(gomock.Any(), duplicateCheckConstraintsSQL).Return(queryResult(t, map[string][]string{
				"constraint_schema": schemas,
				"constraint_name":   constraints,
				"total_count":       {"51"},
			}), nil),
		)

		result, err := runDuplicateCheckConstraints(context.Background(), proxy)
		require.NoError(t, err)
		require.Contains(t, result.evidence, "Showing the first 50 of 51 matches; 1 additional match is omitted.")
		require.LessOrEqual(t, len(result.evidence), maximumEvidenceBytes)
	})

	t.Run("invalid foreign keys", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		proxy := mockclusterproxy.NewMockClusterProxy(ctrl)
		constraints := make([]string, maximumEvidenceItems+1)
		for i := range constraints {
			constraints[i] = fmt.Sprintf(
				"CONSTRAINT `fk_%02d` FOREIGN KEY (`parent_id`) REFERENCES `parent` (`id`) /* FOREIGN KEY INVALID */",
				i,
			)
		}
		gomock.InOrder(
			proxy.EXPECT().QuerySQL(gomock.Any(), invalidForeignKeyCandidateSQL).Return(queryResult(t, map[string][]string{
				"table_schema": {"app"},
				"table_name":   {"child"},
			}), nil),
			proxy.EXPECT().QuerySQL(gomock.Any(), "SHOW CREATE TABLE `app`.`child`").Return(queryResult(t, map[string][]string{
				"Table":        {"child"},
				"Create Table": {"CREATE TABLE child (" + strings.Join(constraints, ",\n") + ")"},
			}), nil),
		)

		result, err := runInvalidForeignKeyCheck(context.Background(), proxy)
		require.NoError(t, err)
		require.Contains(t, result.evidence, "Showing the first 50 of 51 matches; 1 additional match is omitted.")
		require.NotContains(t, result.evidence, "fk_50")
		require.LessOrEqual(t, len(result.evidence), maximumEvidenceBytes)
	})

	t.Run("tables without valid indexes", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		proxy := mockclusterproxy.NewMockClusterProxy(ctrl)
		schemas := make([]string, maximumEvidenceItems)
		tables := make([]string, maximumEvidenceItems)
		for i := range maximumEvidenceItems {
			schemas[i] = "app"
			tables[i] = fmt.Sprintf("table_%02d", i)
		}
		proxy.EXPECT().QuerySQL(gomock.Any(), tablesWithoutValidIndexSQL).Return(queryResult(t, map[string][]string{
			"table_schema": schemas,
			"table_name":   tables,
			"total_count":  {"51"},
		}), nil)

		result, err := runTablesWithoutValidIndex(context.Background(), proxy)
		require.NoError(t, err)
		require.Contains(t, result.evidence, "Showing the first 50 of 51 matches; 1 additional match is omitted.")
		require.LessOrEqual(t, len(result.evidence), maximumEvidenceBytes)
	})

	t.Run("sequences", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		proxy := mockclusterproxy.NewMockClusterProxy(ctrl)
		schemas := make([]string, maximumEvidenceItems)
		sequences := make([]string, maximumEvidenceItems)
		for i := range maximumEvidenceItems {
			schemas[i] = "app"
			sequences[i] = fmt.Sprintf("seq_%02d", i)
		}
		proxy.EXPECT().QuerySQL(gomock.Any(), sequencesSQL).Return(queryResult(t, map[string][]string{
			"sequence_schema": schemas,
			"sequence_name":   sequences,
			"total_count":     {"51"},
		}), nil)

		result, err := runSequenceCheck(context.Background(), proxy)
		require.NoError(t, err)
		require.Contains(t, result.evidence, "Showing the first 50 of 51 matches; 1 additional match is omitted.")
		require.LessOrEqual(t, len(result.evidence), maximumEvidenceBytes)
	})

	t.Run("TTL tables", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		proxy := mockclusterproxy.NewMockClusterProxy(ctrl)
		candidateNames := make([][]string, maximumEvidenceItems+1)
		for i := range candidateNames {
			candidateNames[i] = []string{"app", fmt.Sprintf("ttl_%02d", i)}
		}
		calls := []any{
			proxy.EXPECT().QuerySQL(gomock.Any(), ttlJobEnabledSQL).Return(
				queryResult(t, map[string][]string{"tidb_ttl_job_enable": {"ON"}}), nil,
			),
			proxy.EXPECT().QuerySQL(gomock.Any(), ttlCandidateTablesSQL).Return(
				queryResult(t, map[string][]string{
					"table_schema": column(candidateNames, 0),
					"table_name":   column(candidateNames, 1),
				}), nil,
			),
		}
		for _, name := range candidateNames {
			calls = append(calls, proxy.EXPECT().QuerySQL(
				gomock.Any(), fmt.Sprintf("SHOW CREATE TABLE `%s`.`%s`", name[0], name[1]),
			).Return(queryResult(t, map[string][]string{
				"Table": {name[1]},
				"Create Table": {
					fmt.Sprintf(
						"CREATE TABLE `%s` (`id` BIGINT PRIMARY KEY, `created_at` DATETIME) TTL=`created_at` + INTERVAL 1 DAY TTL_ENABLE='ON'",
						name[1],
					),
				},
			}), nil))
		}
		gomock.InOrder(calls...)

		result, err := runTTLCheck(context.Background(), proxy)
		require.NoError(t, err)
		require.Equal(t, globaldedicatedv1.MigrationCompatibilityCheckStatus_RISK_FOUND, result.status)
		require.Contains(t, result.evidence, "Showing the first 50 of 51 matches; 1 additional match is omitted.")
		require.Contains(t, result.evidence, "app.ttl_00")
		require.Contains(t, result.evidence, "app.ttl_49")
		require.NotContains(t, result.evidence, "app.ttl_50")
		require.LessOrEqual(t, len(result.evidence), maximumEvidenceBytes)
	})
}

func TestSummarizedEvidenceByteLimit(t *testing.T) {
	evidence := summarizedEvidence(
		"Invalid foreign keys detected: ",
		[]string{strings.Repeat("x", maximumEvidenceBytes)},
		1,
	)

	require.LessOrEqual(t, len(evidence), maximumEvidenceBytes)
	require.Contains(t, evidence, "Showing the first 0 of 1 matches; 1 additional match is omitted.")
}

func TestRunMySQLDBCollationCompatibility(t *testing.T) {
	tests := []struct {
		name          string
		counts        [][2]string
		expectedState globaldedicatedv1.MigrationCompatibilityCheckStatus
		expectedText  string
	}{
		{
			name: "no collision", counts: [][2]string{{"3", "3"}, {"2", "2"}, {"1", "1"}},
			expectedState: globaldedicatedv1.MigrationCompatibilityCheckStatus_NO_RISK_FOUND,
			expectedText:  "No collation collision was detected.",
		},
		{
			name: "collision", counts: [][2]string{{"3", "2"}, {"2", "2"}, {"1", "1"}},
			expectedState: globaldedicatedv1.MigrationCompatibilityCheckStatus_RISK_FOUND,
			expectedText:  "mysql.db: total_count=3, grouped_count=2",
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			ctrl := gomock.NewController(t)
			proxy := mockclusterproxy.NewMockClusterProxy(ctrl)
			proxy.EXPECT().QuerySQL(gomock.Any(), mysqlDBSchemaSQL).
				Return(queryResult(t, map[string][]string{"value": {"1"}}), nil)
			for _, count := range test.counts {
				proxy.EXPECT().QuerySQL(gomock.Any(), gomock.Any()).Return(
					queryResult(t, map[string][]string{"total_count": {count[0]}, "grouped_count": {count[1]}}),
					nil,
				)
			}

			result, err := runMySQLDBCollationCompatibility(context.Background(), proxy)
			require.NoError(t, err)
			require.Equal(t, test.expectedState, result.status)
			require.Contains(t, result.evidence, test.expectedText)
		})
	}
}

func TestRunInvalidForeignKeyCheck(t *testing.T) {
	ctrl := gomock.NewController(t)
	proxy := mockclusterproxy.NewMockClusterProxy(ctrl)
	gomock.InOrder(
		proxy.EXPECT().QuerySQL(gomock.Any(), invalidForeignKeyCandidateSQL).Return(queryResult(t, map[string][]string{
			"table_schema": {"app"},
			"table_name":   {"orders"},
		}), nil),
		proxy.EXPECT().QuerySQL(gomock.Any(), "SHOW CREATE TABLE `app`.`orders`").Return(queryResult(t, map[string][]string{
			"Table": {"orders"},
			"Create Table": {
				"CREATE TABLE orders (CONSTRAINT `fk_customer` FOREIGN KEY (`customer_id`) REFERENCES `customers` (`id`) /* FOREIGN KEY INVALID */)",
			},
		}), nil),
	)

	result, err := runInvalidForeignKeyCheck(context.Background(), proxy)
	require.NoError(t, err)
	require.Equal(t, globaldedicatedv1.MigrationCompatibilityCheckStatus_RISK_FOUND, result.status)
	require.Contains(t, result.evidence, "app.orders.fk_customer")
}

func TestInvalidForeignKeyNames(t *testing.T) {
	createSQL := "CREATE TABLE `child` (\\n" +
		"CONSTRAINT `fk``one` FOREIGN KEY (`parent_id`) REFERENCES `parent` (`id`) /* FOREIGN KEY INVALID */,\\n" +
		"CONSTRAINT `fk_valid` FOREIGN KEY (`other_id`) REFERENCES `other` (`id`)\\n)"

	require.Equal(t, []string{"fk`one"}, invalidForeignKeyNames(createSQL))
}

func TestValidIdentifier(t *testing.T) {
	require.True(t, validIdentifier("orders_2026"))
	require.False(t, validIdentifier("orders`; DROP TABLE mysql.user"))
}

func queryResult(t *testing.T, columns map[string][]string) []byte {
	t.Helper()
	encoded := make(map[string][]string, len(columns))
	for name, values := range columns {
		encoded[name] = make([]string, 0, len(values))
		for _, value := range values {
			encoded[name] = append(encoded[name], base64.StdEncoding.EncodeToString([]byte(value)))
		}
	}
	raw, err := json.Marshal(encoded)
	require.NoError(t, err)
	return raw
}