package premium_dm

//nolint:gci
import (
	"context"
	"encoding/base64"
	"slices"

	"github.com/pingcap/tidb/pkg/util/filter"
	tfilter "github.com/pingcap/tidb/pkg/util/table-filter"
	router "github.com/pingcap/tidb/pkg/util/table-router"
	dmconfig "github.com/pingcap/tiflow/dm/config"
	dmdbconfig "github.com/pingcap/tiflow/dm/config/dbconfig"
	"github.com/pingcap/tiflow/dm/config/security"
	"github.com/tidbcloud/dataflow-service-ng/app/repo"
	"github.com/tidbcloud/dataflow-service-ng/app/repo/entity"
	apputils "github.com/tidbcloud/dataflow-service-ng/app/utils"
	"github.com/tidbcloud/dataflow-service-ng/idl/pbgen/proto/premium_dm/v1beta2"
	"github.com/tidbcloud/dataflow-service-ng/internal/config"
	bizerr_wrapper "github.com/tidbcloud/dataflow-service-ng/pkg/bizerr-wrapper"
	encryptor_wrapper "github.com/tidbcloud/dataflow-service-ng/pkg/encryptor-wrapper"
	"github.com/tidbcloud/pkgv2/bizerr"
	tidbmgmtservicev1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"google.golang.org/protobuf/encoding/protojson"
)

// premiumPrecheckTypeToDMCheckerTaskName maps persisted PremiumPrecheckItemType to tiflow/dmconfig checker
// task names (agent, SubTaskConfig). This is the only place that couples our type to DM strings.
var premiumPrecheckTypeToDMCheckerTaskName = map[entity.PremiumPrecheckItemType]string{
	entity.PremiumPrecheckItemTypeDumpPrivilege:        dmconfig.DumpPrivilegeChecking,
	entity.PremiumPrecheckItemTypeReplicationPrivilege: dmconfig.ReplicationPrivilegeChecking,
	entity.PremiumPrecheckItemTypeVersion:              dmconfig.VersionChecking,
	entity.PremiumPrecheckItemTypeServerID:             dmconfig.ServerIDChecking,
	entity.PremiumPrecheckItemTypeBinlogEnable:         dmconfig.BinlogEnableChecking,
	entity.PremiumPrecheckItemTypeBinlogFormat:         dmconfig.BinlogFormatChecking,
	entity.PremiumPrecheckItemTypeBinlogRowImage:       dmconfig.BinlogRowImageChecking,
	entity.PremiumPrecheckItemTypeTableSchema:          dmconfig.TableSchemaChecking,
	entity.PremiumPrecheckItemTypeBinlogDB:             dmconfig.BinlogDBChecking,
	entity.PremiumPrecheckItemTypeConnNumber:           dmconfig.ConnNumberChecking,
	entity.PremiumPrecheckItemTypeTargetDBPrivilege:    dmconfig.TargetDBPrivilegeChecking,
	entity.PremiumPrecheckItemTypeMetaPosition:         dmconfig.MetaPositionChecking,
	entity.PremiumPrecheckItemTypeLightningTableEmpty:  dmconfig.LightningTableEmptyChecking,
	entity.PremiumPrecheckItemTypePrimaryKey:           PrimaryKeyChecking,
}

// PrimaryKeyChecking is an additional checking item not defined by the service module's TiFlow version.
const PrimaryKeyChecking = "primary_key"

// defaultTiDBPort is the default port used for TiDB connections.
const defaultTiDBPort = 4000

var dmCheckerTaskNameToPremiumPrecheckType = func() map[string]entity.PremiumPrecheckItemType {
	m := make(map[string]entity.PremiumPrecheckItemType, len(premiumPrecheckTypeToDMCheckerTaskName))
	for k, v := range premiumPrecheckTypeToDMCheckerTaskName {
		m[v] = k
	}
	return m
}()

// DMCheckerTaskNameForType returns the tiflow checker task name for a persisted PremiumPrecheckItemType.
func DMCheckerTaskNameForType(typ entity.PremiumPrecheckItemType) (string, bool) {
	s, ok := premiumPrecheckTypeToDMCheckerTaskName[typ]
	return s, ok
}

// PremiumPrecheckItemTypeForDMCheckerTaskName resolves a tiflow checker task name to our persisted type.
func PremiumPrecheckItemTypeForDMCheckerTaskName(name string) (entity.PremiumPrecheckItemType, bool) {
	typ, ok := dmCheckerTaskNameToPremiumPrecheckType[name]
	return typ, ok
}

// ApplyPremiumSourcePrecheckConfig enriches a base SubTaskConfig with premium source-side settings.
func ApplyPremiumSourcePrecheckConfig(
	baseCfg *dmconfig.SubTaskConfig,
	premiumSource *entity.PremiumDMSource,
	taskMode entity.PremiumDMTaskMode,
	importMode entity.PremiumDMTaskImportMode,
) *dmconfig.SubTaskConfig {
	if baseCfg == nil || premiumSource == nil {
		return baseCfg
	}

	routeRules := buildPremiumRouteRules(premiumSource)
	for i := range routeRules {
		baseCfg.RouteRules = append(baseCfg.RouteRules, &routeRules[i])
	}

	baseCfg.BAList = buildPremiumBlockAllowRules(premiumSource)
	applyPremiumSourceSecurity(baseCfg, premiumSource)
	applyPremiumBinlogMeta(baseCfg, premiumSource)

	if taskMode == entity.PremiumDMTaskModeAll && importMode != entity.PremiumDMTaskImportModeLogical {
		baseCfg.LoaderConfig.ImportMode = dmconfig.LoadModePhysical // TODO: change to LoadModeImportInto
	}
	return baseCfg
}

func premiumModeToString(mode v1beta2.PremiumDMTaskMode) string {
	switch mode {
	case v1beta2.PremiumDMTaskMode_ALL:
		return "all"
	case v1beta2.PremiumDMTaskMode_INCREMENTAL:
		return "incremental"
	default:
		return "all"
	}
}

func BuildPremiumPrecheckBaseConfig(
	cluster *tidbmgmtservicev1.Cluster,
	req *v1beta2.CreatePremiumMigrationReq,
	premiumSource *entity.PremiumDMSource,
) (*dmconfig.SubTaskConfig, bizerr.BizError) {
	if cluster == nil || req == nil || premiumSource == nil {
		return nil, bizerr_wrapper.NewInternalErr("invalid input: cluster, req, or premiumSource is nil")
	}

	sourceHost := premiumSource.ConnProfile.Target
	sourcePort := int(premiumSource.ConnProfile.Port)
	sourceUser := premiumSource.ConnProfile.User

	subTaskCfg := &dmconfig.SubTaskConfig{
		Name:           "precheck-task",
		Mode:           premiumModeToString(req.Mode),
		SourceID:       "test-precheck",
		LoaderConfig:   dmconfig.DefaultLoaderConfig(),
		SyncerConfig:   dmconfig.DefaultSyncerConfig(),
		MydumperConfig: dmconfig.DefaultMydumperConfig(),
		From: dmdbconfig.DBConfig{
			Host: sourceHost,
			Port: sourcePort,
			User: sourceUser,
		},
		To: dmdbconfig.DBConfig{
			Port: defaultTiDBPort,
			User: req.Target.User,
		},
	}

	return subTaskCfg, nil
}

func decryptPersistedPassword(encryptor encryptor_wrapper.Encryptor, password string) (string, bizerr.BizError) {
	if password == "" {
		return "", nil
	}
	decrypted, err := encryptor.DecryptBase64(password)
	if err != nil {
		return "", err
	}
	return decrypted, nil
}

func buildPremiumRouteRules(premiumSource *entity.PremiumDMSource) []router.TableRule {
	rules := make([]router.TableRule, 0)
	if len(premiumSource.MigrationRules) > 0 {
		schemaExists := make(map[string]struct{})
		for _, rule := range premiumSource.MigrationRules {
			rules = append(rules, router.TableRule{
				SchemaPattern: rule.Source.SchemaPattern,
				TablePattern:  rule.Source.TablePattern,
				TargetSchema:  rule.Target.Schema,
				TargetTable:   rule.Target.Table,
			})
			if _, ok := schemaExists[rule.Source.SchemaPattern]; !ok {
				rules = append(rules, router.TableRule{
					SchemaPattern: rule.Source.SchemaPattern,
					TargetSchema:  rule.Target.Schema,
				})
				schemaExists[rule.Source.SchemaPattern] = struct{}{}
			}
		}
		return rules
	}

	schemaExists := make(map[string]struct{})
	for _, schema := range premiumSource.BlockAllowRules.DoDBs {
		rules = append(rules, router.TableRule{
			SchemaPattern: schema,
			TablePattern:  "*",
			TargetSchema:  schema,
		})
		schemaExists[schema] = struct{}{}
	}
	for _, tbl := range premiumSource.BlockAllowRules.DoTables {
		rule := router.TableRule{
			SchemaPattern: tbl.Schema,
			TablePattern:  tbl.Table,
			TargetSchema:  tbl.Schema,
		}
		if tbl.Table != "*" && tbl.Table != "~.*" {
			rule.TargetTable = tbl.Table
		}
		rules = append(rules, rule)
		if _, ok := schemaExists[tbl.Schema]; !ok {
			rules = append(rules, router.TableRule{
				SchemaPattern: tbl.Schema,
				TargetSchema:  tbl.Schema,
			})
			schemaExists[tbl.Schema] = struct{}{}
		}
	}
	return rules
}

func buildPremiumBlockAllowRules(premiumSource *entity.PremiumDMSource) *tfilter.MySQLReplicationRules {
	baRules := &tfilter.MySQLReplicationRules{}
	if len(premiumSource.MigrationRules) > 0 {
		schemaMap := make(map[string]struct{})
		tableList := make([]*filter.Table, 0)
		for _, rule := range premiumSource.MigrationRules {
			schemaMap[rule.Source.SchemaPattern] = struct{}{}
			tableList = append(tableList, &filter.Table{
				Schema: rule.Source.SchemaPattern,
				Name:   rule.Source.TablePattern,
			})
		}
		for schema := range schemaMap {
			baRules.DoDBs = append(baRules.DoDBs, schema)
		}
		baRules.DoTables = tableList
		return baRules
	}

	baRules.DoDBs = premiumSource.BlockAllowRules.DoDBs
	for _, tbl := range premiumSource.BlockAllowRules.DoTables {
		baRules.DoTables = append(baRules.DoTables, &filter.Table{
			Schema: tbl.Schema,
			Name:   tbl.Table,
		})
	}
	return baRules
}

func applyPremiumSourceSecurity(cfg *dmconfig.SubTaskConfig, premiumSource *entity.PremiumDMSource) {
	if premiumSource.ConnProfile.Security == nil {
		return
	}
	sec := premiumSource.ConnProfile.Security
	if len(sec.SslCaContent) == 0 && len(sec.SslCertContent) == 0 && len(sec.SslKeyContent) == 0 {
		return
	}
	cfg.From.Security = &security.Security{
		SSLCABase64:   base64.StdEncoding.EncodeToString(sec.SslCaContent),
		SSLCertBase64: base64.StdEncoding.EncodeToString(sec.SslCertContent),
		SSLKeyBase64:  base64.StdEncoding.EncodeToString(sec.SslKeyContent),
	}
	if len(sec.CertAllowedCn) > 0 {
		cfg.From.Security.CertAllowedCN = slices.Clone(sec.CertAllowedCn)
	}
}

func applyPremiumBinlogMeta(cfg *dmconfig.SubTaskConfig, premiumSource *entity.PremiumDMSource) {
	if len(premiumSource.BinlogName) == 0 && len(premiumSource.BinlogGTID) == 0 {
		return
	}
	cfg.Meta = &dmconfig.Meta{}
	if len(premiumSource.BinlogName) > 0 {
		cfg.Meta.BinLogName = premiumSource.BinlogName
	}
	if premiumSource.BinlogPos > 0 {
		cfg.Meta.BinLogPos = uint32(premiumSource.BinlogPos)
	}
	if len(premiumSource.BinlogGTID) > 0 {
		cfg.EnableGTID = true
		cfg.Meta.BinLogGTID = premiumSource.BinlogGTID
	}
}

func UnmarshalPrecheckRequest(requestBytes []byte) (*v1beta2.CreatePremiumMigrationReq, error) {
	if len(requestBytes) == 0 {
		return nil, bizerr.New(bizerr.ErrCodeParams, "empty request bytes")
	}
	req := &v1beta2.CreatePremiumMigrationReq{}
	if err := protojson.Unmarshal(requestBytes, req); err != nil {
		return nil, bizerr.New(bizerr.ErrCodeParams, "failed to unmarshal precheck request: %v", err)
	}
	return req, nil
}

// BuildEndpointMapForCluster queries all active BizSinkEndpoints for the given
// tenant/sharedPool and returns a map keyed by endpoint ID. Callers use this map
// to look up endpoint details (e.g. EndpointAddress) when converting PRIVATELINK
// sources to entities.
func BuildEndpointMapForCluster(
	ctx context.Context,
	bizSinkEndpointRepo repo.BizSinkEndpointRepo,
	cluster *apputils.Cluster,
) (map[uint64]*entity.BizSinkEndpoint, error) {
	endpoints, bizErr := bizSinkEndpointRepo.List(ctx,
		func(qs entity.BizSinkEndpointQuerySet) entity.BizSinkEndpointQuerySet {
			return qs.TenantIDEq(entity.ID(cluster.GetTenantID())).SharedPoolIDEq(cluster.GetSharedPoolId()).
				StatusNotIn(entity.BizPLEndpointStatusDeleted)
		})
	if bizErr != nil {
		return nil, bizErr
	}
	m := make(map[uint64]*entity.BizSinkEndpoint, len(endpoints))
	for i := range endpoints {
		ep := &endpoints[i]
		m[uint64(ep.ID)] = ep
	}
	return m, nil
}

func ConvertRequestSourcesToEntities(
	tidbID string,
	tenantID string,
	sources []*v1beta2.PremiumDMSource,
	endpointMap map[uint64]*entity.BizSinkEndpoint,
) []entity.PremiumDMSource {
	if len(sources) == 0 {
		return nil
	}
	result := make([]entity.PremiumDMSource, 0, len(sources))
	for _, s := range sources {
		if s == nil {
			continue
		}
		result = append(result, entity.PremiumDMSource{
			ID:             entity.GenerateDMIDWithPrefix(entity.PremiumDMSourcePrefix),
			Stage:          entity.PremiumDMSourceStageCreating,
			TiDBID:         tidbID,
			TenantID:       tenantID,
			SourceType:     sourceTypeToEntity(s.GetSourceType()),
			ConnProfile:    connProfileToEntity(s.GetConnProfile(), endpointMap),
			MigrationRules: migrationRulesToEntity(s.GetPremiumMigrationRules()),
			BinlogName:     s.GetBinlogName(),
			BinlogPos:      s.GetBinlogPosition(),
			BinlogGTID:     s.GetBinlogGtid(),
		})
	}
	return result
}

func sourceTypeToEntity(src v1beta2.PremiumDMSourceType) entity.PremiumDBType {
	switch src {
	case v1beta2.PremiumDMSourceType_MYSQL:
		return entity.PremiumDBTypeMySQL
	case v1beta2.PremiumDMSourceType_AZURE_FLEXIBLE_MYSQL:
		return entity.PremiumDBTypeAzureFlexibleMySQL
	case v1beta2.PremiumDMSourceType_ALICLOUD_RDS_MYSQL:
		return entity.PremiumDBTypeAliCloudRDSMySQL
	default:
		return ""
	}
}

func connProfileToEntity(
	pb *v1beta2.PremiumDMConnProfile,
	endpointMap map[uint64]*entity.BizSinkEndpoint,
) entity.PremiumConnProfile {
	if pb == nil {
		return entity.PremiumConnProfile{}
	}
	switch pb.ConnType {
	case v1beta2.PremiumDMConnType_PUBLIC:
		return entity.PremiumConnProfile{
			Type:     entity.PremiumConnTypePublic,
			Target:   pb.GetHost(),
			Port:     pb.GetPort(),
			User:     pb.GetUser(),
			Password: pb.GetPassword(),
			Security: securityToEntity(pb.GetSecurity()),
		}
	case v1beta2.PremiumDMConnType_PRIVATELINK:
		sinkID := pb.GetSinkEndpointId()
		target := ""
		if ep, ok := endpointMap[sinkID]; ok && ep != nil {
			target = ep.EndpointAddress
		}
		return entity.PremiumConnProfile{
			Type:           entity.PremiumConnTypePrivateLink,
			Target:         target,
			SinkEndpointID: sinkID,
			Port:           pb.GetPort(),
			User:           pb.GetUser(),
			Password:       pb.GetPassword(),
			Security:       securityToEntity(pb.GetSecurity()),
		}
	}
	return entity.PremiumConnProfile{}
}

func securityToEntity(pb *v1beta2.Security) *entity.PremiumSecurity {
	if pb == nil {
		return nil
	}
	return &entity.PremiumSecurity{
		CertAllowedCn:  pb.GetCertAllowedCn(),
		SslCaContent:   pb.GetSslCaContent(),
		SslCertContent: pb.GetSslCertContent(),
		SslKeyContent:  pb.GetSslKeyContent(),
	}
}

func migrationRulesToEntity(pb []*v1beta2.PremiumMigrationRule) entity.PremiumMigrationRules {
	if len(pb) == 0 {
		return nil
	}
	rules := make(entity.PremiumMigrationRules, 0, len(pb))
	for _, r := range pb {
		if r == nil {
			continue
		}
		rule := entity.PremiumMigrationRule{}
		if r.SourceTable != nil {
			rule.Source.SchemaPattern = r.SourceTable.GetSchemaPattern()
			rule.Source.TablePattern = r.SourceTable.GetTablePattern()
		}
		rule.Target = migrationRuleTargetToEntity(r.GetTargetTable())
		rules = append(rules, rule)
	}
	return rules
}

func migrationRuleTargetToEntity(pbTarget *v1beta2.PremiumMigrationRule_Table) entity.PremiumMigrationRuleTarget {
	if pbTarget == nil {
		return entity.PremiumMigrationRuleTarget{}
	}
	return entity.PremiumMigrationRuleTarget{Schema: pbTarget.GetSchema(), Table: pbTarget.GetTable()}
}

var (
	PremiumPrecheckFullModeTypes = []entity.PremiumPrecheckItemType{
		entity.PremiumPrecheckItemTypeDumpPrivilege,
		entity.PremiumPrecheckItemTypeVersion,
		entity.PremiumPrecheckItemTypeTableSchema,
		entity.PremiumPrecheckItemTypeConnNumber,
		entity.PremiumPrecheckItemTypeTargetDBPrivilege,
	}
	PremiumPrecheckIncrementalBaseTypes = []entity.PremiumPrecheckItemType{
		entity.PremiumPrecheckItemTypeReplicationPrivilege,
		entity.PremiumPrecheckItemTypeVersion,
		entity.PremiumPrecheckItemTypeServerID,
		entity.PremiumPrecheckItemTypeBinlogEnable,
		entity.PremiumPrecheckItemTypeBinlogFormat,
		entity.PremiumPrecheckItemTypeBinlogRowImage,
		entity.PremiumPrecheckItemTypeTableSchema,
		entity.PremiumPrecheckItemTypeBinlogDB,
		entity.PremiumPrecheckItemTypeTargetDBPrivilege,
	}
	PremiumPrecheckAllModeTypes = []entity.PremiumPrecheckItemType{
		entity.PremiumPrecheckItemTypeDumpPrivilege,
		entity.PremiumPrecheckItemTypeReplicationPrivilege,
		entity.PremiumPrecheckItemTypeVersion,
		entity.PremiumPrecheckItemTypeServerID,
		entity.PremiumPrecheckItemTypeBinlogEnable,
		entity.PremiumPrecheckItemTypeBinlogFormat,
		entity.PremiumPrecheckItemTypeBinlogRowImage,
		entity.PremiumPrecheckItemTypeTableSchema,
		entity.PremiumPrecheckItemTypeBinlogDB,
		entity.PremiumPrecheckItemTypeConnNumber,
		entity.PremiumPrecheckItemTypeTargetDBPrivilege,
	}
)

// PremiumPrecheckItemTypesForMode returns the ordered precheck types for a task mode (same order as execution).
func PremiumPrecheckItemTypesForMode(
	taskMode entity.PremiumDMTaskMode,
	hasAnySourceBinlogCheckpoint bool,
	_ entity.PremiumDMTaskImportMode,
	needsPrimaryKeyCheck bool,
) []entity.PremiumPrecheckItemType {
	var types []entity.PremiumPrecheckItemType
	if needsPrimaryKeyCheck {
		types = append(types, entity.PremiumPrecheckItemTypePrimaryKey)
	}
	switch taskMode {
	case entity.PremiumDMTaskModeFull:
		types = append(types, PremiumPrecheckFullModeTypes...)
	case entity.PremiumDMTaskModeIncremental:
		types = append(types, PremiumPrecheckIncrementalBaseTypes...)
		if hasAnySourceBinlogCheckpoint {
			types = append(types, entity.PremiumPrecheckItemTypeMetaPosition)
		}
	case entity.PremiumDMTaskModeAll:
		types = append(types, PremiumPrecheckAllModeTypes...)
	}
	return types
}

// PremiumPrecheckItemTypesForCluster returns precheck types for a cluster after applying configured skip rules.
func PremiumPrecheckItemTypesForCluster(
	taskMode entity.PremiumDMTaskMode,
	hasAnySourceBinlogCheckpoint bool,
	importMode entity.PremiumDMTaskImportMode,
	needsPrimaryKeyCheck bool,
	skipConfigs config.PrecheckSkipConfigs,
	tenantID, projectID, clusterID uint64,
) []entity.PremiumPrecheckItemType {
	types := PremiumPrecheckItemTypesForMode(
		taskMode, hasAnySourceBinlogCheckpoint, importMode, needsPrimaryKeyCheck,
	)
	return applyPremiumPrecheckSkipConfigs(types, skipConfigs, tenantID, projectID, clusterID)
}

func applyPremiumPrecheckSkipConfigs(
	types []entity.PremiumPrecheckItemType,
	skipConfigs config.PrecheckSkipConfigs,
	tenantID, projectID, clusterID uint64,
) []entity.PremiumPrecheckItemType {
	if len(skipConfigs) == 0 {
		return types
	}

	skipped := skipConfigs.GetSkippedItems(tenantID, projectID, clusterID)
	if len(skipped) == 0 {
		return types
	}

	result := make([]entity.PremiumPrecheckItemType, 0, len(types))
	for _, typ := range types {
		taskName, ok := DMCheckerTaskNameForType(typ)
		if !ok {
			continue
		}
		if skipped[taskName] {
			continue
		}
		result = append(result, typ)
	}
	return result
}

func GetPremiumPrecheckExecutionItems(
	taskMode entity.PremiumDMTaskMode,
	hasAnySourceBinlogCheckpoint bool,
	importMode entity.PremiumDMTaskImportMode,
	needsPrimaryKeyCheck bool,
	precheckID string,
	skipConfigs config.PrecheckSkipConfigs,
	tenantID, projectID, clusterID uint64,
) []*entity.PremiumDMPrecheckItem {
	types := PremiumPrecheckItemTypesForCluster(
		taskMode,
		hasAnySourceBinlogCheckpoint,
		importMode,
		needsPrimaryKeyCheck,
		skipConfigs,
		tenantID,
		projectID,
		clusterID,
	)
	items := make([]*entity.PremiumDMPrecheckItem, 0, len(types))
	for _, typ := range types {
		taskName, ok := DMCheckerTaskNameForType(typ)
		if !ok {
			continue
		}
		items = append(items, &entity.PremiumDMPrecheckItem{
			ID:           generatePremiumPrecheckItemID(),
			Stage:        string(entity.PremiumDMPrecheckStagePending),
			Description:  taskName,
			Type:         typ,
			DMPrecheckID: precheckID,
		})
	}
	return items
}