package importcfg

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

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/core-svc/idl/pbgen/proto/core"
	corepb "github.com/tidbcloud/core-svc/idl/pbgen/proto/core"
	"github.com/tidbcloud/dataflow-service/app/models/config_manager"
	config_manager_impl "github.com/tidbcloud/dataflow-service/app/models/config_manager/impl"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/dataflow"
	"github.com/tidbcloud/dataflow-service/internal/db/entity"
	"github.com/tidbcloud/dataflow-service/pkg/grpc/grpcli"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/pkg/pingcap/util/config"
	"k8s.io/utils/ptr"
)

func generateLightningConfig(cfgStr string) *config.GenericConfig {
	cfg := config.New(map[string]interface{}{})
	_ = cfg.UnmarshalTOML([]byte(cfgStr))
	return config.New(cloudv1.TomlConfig(*cfg).MP)
}

var normalConfig = `[checkpoint]
driver = "file"
dsn = "/etc/endpoint/cp.pb"
enable = true

[cron]
log-progress = "5m"
switch-mode = "5m"

[lightning]
check-requirements = false
file = ""
index-concurrency = 4
level = "info"
pprof-port = 8289
table-concurrency = 6

[mydumper]
data-source-dir = "fakeURL"
no-schema = false
strict-format = false
[mydumper.csv]
  backslash-escape = true
  delimiter = "\""
  header = true
  not-null = false
  null = "\\N"
  separator = ","

[post-restore]
analyze = true
checksum = true

[security]
ca-path = "/var/lib/lightning-tls/ca.crt"
cert-path = "/var/lib/lightning-tls/tls.crt"
key-path = "/var/lib/lightning-tls/tls.key"

[tidb]
build-stats-concurrency = 4
distsql-scan-concurrency = 15
host = "fakeName-tidb"
index-serial-scan-concurrency = 2
log-level = "error"
password = ""
pd-addr = "fakeName-pd:2379"
port = 4000
status-port = 10080
tls = "false"
user = ""

[tikv-importer]
backend = "local"
sorted-kv-dir = "/etc/endpoint"
`

func TestGetFilePatternsInConfig(t *testing.T) {
	type testCase struct {
		CaseName     string
		ConfigStr    string
		ExpectResult []map[string]string
	}
	for _, tc := range []testCase{
		{
			CaseName:     "normal template",
			ConfigStr:    normalConfig,
			ExpectResult: nil,
		},
		{
			CaseName: "single file pattern",
			ConfigStr: normalConfig + `
[[mydumper.files]]
pattern = "(?i)^(?:[^/]*/)*([^/.]+)\\.(.*?)(?:\\.([0-9_]+))?\\.csv$"
schema = "$1"
table = "$2"
type = "csv"
`,
			ExpectResult: []map[string]string{
				{
					"pattern": "(?i)^(?:[^/]*/)*([^/.]+)\\.(.*?)(?:\\.([0-9_]+))?\\.csv$",
					"schema":  "$1",
					"table":   "$2",
					"type":    "csv",
				},
			},
		},
		{
			CaseName: "multiple file patterns",
			ConfigStr: normalConfig + `
[[mydumper.files]]
pattern = "(?i)^(?:[^/]*/)*testdb\\.tbl01(?:\\.([0-9_]+))?\\.csv$"
schema = "testdb"
table = "tbl01"
type = "csv"

[[mydumper.files]]
pattern = "(?i)^(?:[^/]*/)*testdb\\.tbl02(?:\\.([0-9_]+))?\\.csv$"
schema = "testdb"
table = "tbl02"
type = "csv"

[[mydumper.files]]
pattern = "(?i)^(?:[^/]*/)*testdb\\.tbl03(?:\\.([0-9_]+))?\\.csv$"
schema = "testdb"
table = "tbl03"
type = "csv"
`,
			ExpectResult: []map[string]string{
				{
					"pattern": "(?i)^(?:[^/]*/)*testdb\\.tbl01(?:\\.([0-9_]+))?\\.csv$",
					"schema":  "testdb",
					"table":   "tbl01",
					"type":    "csv",
				},
				{
					"pattern": "(?i)^(?:[^/]*/)*testdb\\.tbl02(?:\\.([0-9_]+))?\\.csv$",
					"schema":  "testdb",
					"table":   "tbl02",
					"type":    "csv",
				},
				{
					"pattern": "(?i)^(?:[^/]*/)*testdb\\.tbl03(?:\\.([0-9_]+))?\\.csv$",
					"schema":  "testdb",
					"table":   "tbl03",
					"type":    "csv",
				},
			},
		},
	} {
		cfg := config.New(map[string]interface{}{})
		err := cfg.UnmarshalTOML([]byte(tc.ConfigStr))
		require.NoErrorf(t, err, "test case name: %s", tc.CaseName)
		filePatterns := getFilePatternsInConfig(cfg)
		require.Equalf(t, tc.ExpectResult, filePatterns, "test case name: %s", tc.CaseName)
	}
}

func TestSetTidbConfig(t *testing.T) {
	type testCase struct {
		caseName       string
		cluster        *corepb.Cluster
		expected       string
		expectedErrMsg string
	}

	normalCfg := `[tidb]
  host = "db-tidb"
  pd-addr = "db-pd:2379"
  port = %d
  user = "cloud_admin"

[tikv-importer]
  backend = "local"
`

	normalPortCfg := fmt.Sprintf(normalCfg, 4000)

	fakeToken := "fakeToken"
	for _, tc := range []testCase{
		{
			caseName:       "infra provider type is empty",
			cluster:        &corepb.Cluster{},
			expectedErrMsg: "not support infra provider type: ",
		},
		{
			caseName: "serverless token is nil",
			cluster: &corepb.Cluster{
				InfraProviderType: cloudv1.ProviderTypeAWSFreeTier,
			},
			expectedErrMsg: "serverless token is nil",
		},
		{
			caseName: "serverless normal config",
			cluster: &corepb.Cluster{
				InfraProviderType: cloudv1.ProviderTypeAWSFreeTier,
				ServerlessToken:   &fakeToken,
			},
			expected: `[tidb]
  host = "tidb-gateway-inner.tidb-admin.svc.cluster.local"
  pd-addr = "serverless-cluster-pd.tidb-serverless.svc.cluster.local:2379"
  port = 4000
  user = "fakeToken.cloud_admin"

[tikv-importer]
  backend = "tidb"
`,
		},
		{
			caseName: "eks normal config",
			cluster: &corepb.Cluster{
				InfraProviderType: cloudv1.ProviderTypeEKS,
			},
			expected: normalPortCfg,
		},
		{
			caseName: "gke normal config",
			cluster: &corepb.Cluster{
				InfraProviderType: cloudv1.ProviderTypeGKE,
			},
			expected: normalPortCfg,
		},
		{
			caseName: "aws(gardener) normal config",
			cluster: &corepb.Cluster{
				InfraProviderType: cloudv1.ProviderTypeAWS,
			},
			expected: normalPortCfg,
		},
		{
			caseName: "gcp(gardener) normal config",
			cluster: &corepb.Cluster{
				InfraProviderType: cloudv1.ProviderTypeGCP,
			},
			expected: normalPortCfg,
		},
		{
			caseName: "aws(gardener) custom config",
			cluster: &corepb.Cluster{
				InfraProviderType: cloudv1.ProviderTypeAWS,
				Connections: []*corepb.ClusterConnection{
					{
						Type: corepb.ClusterConnectionType_CLUSTER_CONNECTION_TYPE_PUBLIC,
						Port: 1024,
					},
				},
			},
			expected: fmt.Sprintf(normalCfg, 1024),
		},
		{
			caseName: "gcp(gardener) custom config",
			cluster: &corepb.Cluster{
				InfraProviderType: cloudv1.ProviderTypeGCP,
				Connections: []*corepb.ClusterConnection{
					{
						Type: corepb.ClusterConnectionType_CLUSTER_CONNECTION_TYPE_PUBLIC,
						Port: 1024,
					},
				},
			},
			expected: fmt.Sprintf(normalCfg, 1024),
		},
	} {
		cfg, err := setTidbConfig(generateLightningConfig(""), tc.cluster)
		if tc.expectedErrMsg != "" {
			require.Equalf(t, err.Error(), tc.expectedErrMsg, "test case name: %s", tc.caseName)
			continue
		}
		require.NoErrorf(t, err, "test case name: %s", tc.caseName)
		cfgBytes, err := cfg.MarshalTOML()
		require.NoErrorf(t, err, "test case name: %s", tc.caseName)
		require.Equalf(t, tc.expected, string(cfgBytes), "test case name: %s", tc.caseName)
	}
}

func TestSetFilePattern(t *testing.T) {
	type testCase struct {
		caseName string
		imp      entity.Import
		expected string
	}

	normalFilePattern := `[mydumper]

  [[mydumper.files]]
    pattern = "(?i)^(?:[^/]*/)*([^/.]+)-schema-create\\.sql$"
    schema = "$1"
    type = "schema-schema"

  [[mydumper.files]]
    pattern = "(?i)^(?:[^/]*/)*([^/.]+)\\.(.*?)-schema\\.sql$"
    schema = "$1"
    table = "$2"
    type = "table-schema"

  [[mydumper.files]]
    compression = "$3"
    pattern = "%s"
    schema = "$1"
    table = "$2"
    type = "%s"
`

	noCompressionFilePattern := `[mydumper]

  [[mydumper.files]]
    pattern = "(?i)^(?:[^/]*/)*([^/.]+)-schema-create\\.sql$"
    schema = "$1"
    type = "schema-schema"

  [[mydumper.files]]
    pattern = "(?i)^(?:[^/]*/)*([^/.]+)\\.(.*?)-schema\\.sql$"
    schema = "$1"
    table = "$2"
    type = "table-schema"

  [[mydumper.files]]
    pattern = "%s"
    schema = "$1"
    table = "$2"
    type = "parquet"
`
	escapeSlash := func(s string) string {
		return strings.ReplaceAll(s, "\\", "\\\\")
	}

	for _, tc := range []testCase{
		{
			caseName: "sql file pattern",
			imp: entity.Import{
				DataFormat: entity.ImportDataFormatSql,
			},
			expected: fmt.Sprintf(normalFilePattern, escapeSlash(DefaultSQLFilePattern), "sql"),
		},
		{
			caseName: "csv file pattern",
			imp: entity.Import{
				DataFormat: entity.ImportDataFormatCsv,
			},
			expected: fmt.Sprintf(normalFilePattern, escapeSlash(DefaultCSVFilePattern), "csv"),
		},
		{
			caseName: "parquet file pattern",
			imp: entity.Import{
				DataFormat: entity.ImportDataFormatParquet,
			},
			expected: fmt.Sprintf(noCompressionFilePattern, escapeSlash(DefaultRawParquetFilePattern)),
		},
		{
			caseName: "aurora file pattern",
			imp: entity.Import{
				DataFormat: entity.ImportDataFormatAuroraSnapshot,
			},
			expected: fmt.Sprintf(noCompressionFilePattern, escapeSlash(DefaultAuroraSnapshotFilePattern)),
		},
		{
			caseName: "sql file pattern + custom/no custom",
			imp: entity.Import{
				DataFormat: entity.ImportDataFormatSql,
				SourceURL:  "s3://bucket/prefix/",
				TargetTableInfos: entity.TargetTableInfos{
					Infos: []entity.TargetTableInfo{
						{
							Schema:     "test",
							Table:      "test",
							CustomFile: "s3://bucket/prefix/test*.sql",
						},
						{
							Schema: "test2",
							Table:  "test2",
						},
					},
				},
			},
			expected: `[mydumper]

  [[mydumper.files]]
    pattern = "^test.*\\.sql$"
    schema = "test"
    table = "test"
    type = "sql"

  [[mydumper.files]]
    compression = "$1"
    pattern = "(?i)^test2.test2(?:(?:\\.\\d+)|(?:\\.?(?:_\\d+){3}))?\\.sql(?:\\.(\\w+))?$"
    schema = "test2"
    table = "test2"
    type = "sql"
`,
		},
	} {
		cfg, err := setFilePattern(generateLightningConfig(""), tc.imp)
		require.NoErrorf(t, err, "test case name: %s", tc.caseName)
		cfgBytes, err := cfg.MarshalTOML()
		require.NoErrorf(t, err, "test case name: %s", tc.caseName)
		require.Equalf(t, tc.expected, string(cfgBytes), "test case name: %s", tc.caseName)
	}
}

func TestSetFilePatternWithTargetTable(t *testing.T) {
	type testCase struct {
		caseName string
		imp      entity.Import
		expected []map[string]string
	}

	var (
		schema    = "db"
		table     = "tbl"
		sourceUri = "s3://bucket/prefix/"
	)

	for _, tc := range []testCase{
		{
			caseName: "sql file pattern with custom file",
			imp: entity.Import{
				SourceURL:  sourceUri,
				DataFormat: entity.ImportDataFormatSql,
				TargetTableInfos: entity.TargetTableInfos{
					Infos: []entity.TargetTableInfo{
						{
							Schema:     schema,
							Table:      table,
							CustomFile: "s3://bucket/prefix/test*.sql",
						},
					},
				},
			},
			expected: []map[string]string{
				{
					"pattern": "^test.*\\.sql$",
					"schema":  schema,
					"table":   table,
					"type":    "sql",
				},
			},
		},
		{
			caseName: "sql file pattern",
			imp: entity.Import{
				SourceURL:  sourceUri,
				DataFormat: entity.ImportDataFormatSql,
				TargetTableInfos: entity.TargetTableInfos{
					Infos: []entity.TargetTableInfo{
						{
							Schema: schema,
							Table:  table,
						},
					},
				},
			},
			expected: []map[string]string{
				{
					"pattern":     buildSQLFilePatternByTargetTable(schema, table),
					"schema":      schema,
					"table":       table,
					"type":        "sql",
					"compression": "$1",
				},
			},
		},
		{
			caseName: "csv file pattern",
			imp: entity.Import{
				SourceURL:  sourceUri,
				DataFormat: entity.ImportDataFormatCsv,
				TargetTableInfos: entity.TargetTableInfos{
					Infos: []entity.TargetTableInfo{
						{
							Schema: schema,
							Table:  table,
						},
					},
				},
			},
			expected: []map[string]string{
				{
					"pattern":     buildCSVFilePatternByTargetTable(schema, table),
					"schema":      schema,
					"table":       table,
					"type":        "csv",
					"compression": "$1",
				},
			},
		},
		{
			caseName: "parquet file pattern",
			imp: entity.Import{
				SourceURL:  sourceUri,
				DataFormat: entity.ImportDataFormatParquet,
				TargetTableInfos: entity.TargetTableInfos{
					Infos: []entity.TargetTableInfo{
						{
							Schema: schema,
							Table:  table,
						},
					},
				},
			},
			expected: []map[string]string{
				{
					"pattern": buildRawParquetFilePatternByTargetTable(schema, table),
					"schema":  schema,
					"table":   table,
					"type":    "parquet",
				},
			},
		},
		{
			caseName: "aurora file pattern",
			imp: entity.Import{
				SourceURL:  sourceUri,
				DataFormat: entity.ImportDataFormatAuroraSnapshot,
				TargetTableInfos: entity.TargetTableInfos{
					Infos: []entity.TargetTableInfo{
						{
							Schema: schema,
							Table:  table,
						},
					},
				},
			},
			expected: []map[string]string{
				{
					"pattern": buildAuroraSnapshotFilePatternByTargetTable(schema, table),
					"schema":  schema,
					"table":   table,
					"type":    "parquet",
				},
			},
		},
	} {
		cfgMap, err := setFilePatternWithTargetTable(tc.imp)
		require.NoErrorf(t, err, "test case name: %s", tc.caseName)
		require.Equalf(t, tc.expected, cfgMap, "test case name: %s", tc.caseName)
	}
}

func TestSetFilePatternWithNoTargetTable(t *testing.T) {
	type testCase struct {
		caseName string
		imp      entity.Import
		expected []map[string]string
	}

	schemaPatterns := []map[string]string{
		{
			"pattern": "(?i)^(?:[^/]*/)*([^/.]+)-schema-create\\.sql$",
			"schema":  "$1",
			"type":    "schema-schema",
		},
		{
			"pattern": "(?i)^(?:[^/]*/)*([^/.]+)\\.(.*?)-schema\\.sql$",
			"schema":  "$1",
			"table":   "$2",
			"type":    "table-schema",
		},
	}

	for _, tc := range []testCase{
		{
			caseName: "sql file pattern",
			imp: entity.Import{
				DataFormat: entity.ImportDataFormatSql,
			},
			expected: append(schemaPatterns, map[string]string{
				"pattern":     DefaultSQLFilePattern,
				"schema":      "$1",
				"table":       "$2",
				"type":        "sql",
				"compression": "$3",
			}),
		},
		{
			caseName: "csv file pattern",
			imp: entity.Import{
				DataFormat: entity.ImportDataFormatCsv,
			},
			expected: append(schemaPatterns, map[string]string{
				"pattern":     DefaultCSVFilePattern,
				"schema":      "$1",
				"table":       "$2",
				"type":        "csv",
				"compression": "$3",
			}),
		},
		{
			caseName: "parquet file pattern",
			imp: entity.Import{
				DataFormat: entity.ImportDataFormatParquet,
			},
			expected: append(schemaPatterns, map[string]string{
				"pattern": DefaultRawParquetFilePattern,
				"schema":  "$1",
				"table":   "$2",
				"type":    "parquet",
			}),
		},
		{
			caseName: "aurora file pattern",
			imp: entity.Import{
				DataFormat: entity.ImportDataFormatAuroraSnapshot,
			},
			expected: append(schemaPatterns, map[string]string{
				"pattern": DefaultAuroraSnapshotFilePattern,
				"schema":  "$1",
				"table":   "$2",
				"type":    "parquet",
			}),
		},
	} {
		cfgMap := setFilePatternWithNoTargetTable(tc.imp)
		require.Equalf(t, tc.expected, cfgMap, "test case name: %s", tc.caseName)
	}
}

func TestSetCsvCfg(t *testing.T) {
	type testCase struct {
		caseName string
		imp      entity.Import
		csvCfg   *dataflow.CustomCSVFormat
		expected string
	}

	normalCsvConfig := `[mydumper]
  [mydumper.csv]
    backslash-escape = false
    delimiter = "\""
    header = true
    header-schema-match = true
    not-null = false
    null = ""
    separator = ","
    trim-last-separator = false
`

	tr := true

	for _, tc := range []testCase{
		{
			caseName: "empty csv config",
			imp: entity.Import{
				DataFormat: entity.ImportDataFormatCsv,
			},
			expected: "",
		},
		{
			caseName: "non csv data format",
			imp: entity.Import{
				DataFormat: entity.ImportDataFormatSql,
			},
			expected: "",
		},
		{
			caseName: "csv config",
			imp: entity.Import{
				DataFormat: entity.ImportDataFormatCsv,
			},
			csvCfg: &dataflow.CustomCSVFormat{
				Separator:         ",",
				Header:            true,
				Delimiter:         "\"",
				HeaderSchemaMatch: &tr,
			},
			expected: normalCsvConfig,
		},
	} {
		if tc.csvCfg != nil {
			csvConfigBytes, err := json.Marshal(tc.csvCfg)
			require.NoErrorf(t, err, "test case name: %s", tc.caseName)
			tc.imp.CustomCSVFormat = string(csvConfigBytes)
		}

		cfg, err := setCsvCfg(generateLightningConfig(""), tc.imp)
		require.NoErrorf(t, err, "test case name: %s", tc.caseName)
		cfgBytes, err := cfg.MarshalTOML()
		require.NoErrorf(t, err, "test case name: %s", tc.caseName)
		require.Equalf(t, tc.expected, string(cfgBytes), "test case name: %s", tc.caseName)
	}
}

func TestSetFilter(t *testing.T) {
	testcases := []struct {
		targetInfos []entity.TargetTableInfo
		expected    string
	}{
		{
			targetInfos: []entity.TargetTableInfo{
				{Schema: "a", Table: "b"},
				{Schema: "c", Table: "d"},
			},
			expected: "[mydumper]\n  filter = [\"`a`.`b`\", \"`c`.`d`\"]\n",
		},
		{
			targetInfos: []entity.TargetTableInfo{},
			expected:    "",
		},
	}

	for _, tc := range testcases {
		cfg := setFilter(generateLightningConfig(""), tc.targetInfos)
		cfgBytes, err := cfg.MarshalTOML()
		require.NoError(t, err)
		require.Equal(t, tc.expected, string(cfgBytes))
	}
}

func TestGenerateImportConfig(t *testing.T) {
	type args struct {
		ctx          context.Context
		configCenter config_manager.ConfigCenterManager
		imp          entity.Import
		cluster      *core.Cluster
	}
	ctx := context.TODO()
	configCenterGrpcClient := grpcli.NewConfigCenterGrpcClient(
		ctx,
		"127.0.0.1:10409")
	configCenter := config_manager_impl.NewConfigCenterManager(configCenterGrpcClient)
	imp := entity.Import{
		Updatable: entity.Updatable{
			Base: entity.Base{
				ID: 9999,
			},
		},
		Name:           "import",
		ClusterID:      999,
		BucketRegion:   "us-east-1",
		AwsRoleArn:     "arn:aws:iam::XXXXXXXXXXXX:role/my-role",
		DataFormat:     entity.ImportDataFormatCsv,
		DataSourceType: entity.ImportDataSourceS3,
		TargetUsername: "cloud_admin",
		TargetPassword: "YWFhCg==",
		SourceURL:      "s3://some_bucket/some_path/",
		TargetTableInfos: entity.TargetTableInfos{
			Infos: []entity.TargetTableInfo{
				{Schema: "db", Table: "table"},
			},
		},
	}
	cluster := &core.Cluster{
		Id:                999,
		TenantId:          9,
		ProjectId:         99,
		CloudProviderName: "azure",
		InfraProviderType: "aks",
		RegionName:        "eastus",
		IsDevTier:         false,
		Version:           "v7.5.3",
		Attributes: &core.Cluster_Attributes{
			UseNewDedicatedArch: ptr.To(true),
		},
	}
	tests := []struct {
		name    string
		args    args
		wantErr bool
	}{
		{
			name: "OK",
			args: args{
				ctx:          context.TODO(),
				configCenter: configCenter,
				imp:          imp,
				cluster:      cluster,
			},
			wantErr: false,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, got1, err := GenerateImportConfig(tt.args.ctx, tt.args.configCenter, tt.args.imp, "", tt.args.cluster)
			if (err != nil) != tt.wantErr {
				t.Errorf("GenerateImportConfig() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			fmt.Println(got, got1)
		})
	}
}