package encryptionkey

import (
	"context"
	"fmt"
	"strconv"

	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"

	globalngv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/region"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/service/workflow/common"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	"github.com/tidbcloud/tidb-management-service/utilities/grpc/metadata"
	"github.com/tidbcloud/tidb-management-service/utilities/workflow"
)

type SwitchEncryptionRequest struct {
	ClusterId        uint64
	TargetCmek       *globalngv1.CustomerManagedEncryptionKey // nil if switching to default_key or disabling
	TargetDefaultKey bool                                     // true if switching to default_key
	TargetDisable    bool                                     // true if disabling encryption
	ValidateOnly     bool
}

func (s *service) SwitchClusterEncryptionKey(
	ctx context.Context,
	req *SwitchEncryptionRequest,
) bizerr.BizError {
	// 1. Load cluster
	cluster, err := s.clusterRepo.Get(ctx, req.ClusterId, false)
	if err != nil {
		return bizerr.WrapError(err, "failed to load cluster %d", req.ClusterId)
	}

	if cluster.GetState() != globalngv1.Cluster_AVAILABLE {
		return bizerr.New(bizerr.ErrCodeParams, "cluster must be in AVAILABLE state, current: %s", cluster.GetState().String())
	}

	// 2. Load current encryption key (if any)
	var currentKey *entity.EncryptionKey
	if cluster.EncryptionKeyId != 0 {
		currentKey, err = s.encryptionKeyRepo.Load(ctx, cluster.EncryptionKeyId, false)
		if err != nil {
			return bizerr.WrapError(err, "failed to load current encryption key")
		}
	}

	// 3. Determine transition type
	transitionType, modeErr := determineTransitionType(currentKey, req)
	if modeErr != nil {
		return modeErr
	}

	// 3.5. No-op: disabling encryption on an already-unencrypted cluster
	if transitionType == "no-op" {
		return nil
	}

	// 4. Validate transition is allowed
	if validErr := validateTransition(transitionType); validErr != nil {
		return validErr
	}

	// 5. Check binding not in Modifying state
	isModifying, err := s.bindingRepo.IsModifying(ctx, req.ClusterId)
	if err != nil {
		return bizerr.WrapError(err, "failed to check binding state")
	}
	if isModifying {
		return bizerr.New(bizerr.ErrCodeParams, "encryption change already in progress for cluster %d", req.ClusterId)
	}

	// 6. Resolve target encryption key
	var targetKeyId uint64
	if req.TargetCmek != nil {
		// Enrich CMEK with cluster context — the OpenAPI layer only forwards key/principal fields.
		if req.TargetCmek.OrganizationId == 0 {
			req.TargetCmek.OrganizationId = cluster.GetOrganizationId()
		}
		if req.TargetCmek.CloudProviderId == 0 {
			req.TargetCmek.CloudProviderId = cluster.GetCloudProviderId()
		}
		if req.TargetCmek.RegionId == 0 {
			req.TargetCmek.RegionId = cluster.GetRegionId()
		}
		if req.TargetCmek.ServicePlan == 0 {
			req.TargetCmek.ServicePlan = cluster.GetServicePlan()
		}

		// Verify IAM principal
		verifyResp, verifyErr := s.verifyCmekAccessIamPrincipal(ctx, &globalngv1.VerifyCmekAccessIamPrincipalRequest{
			Key:       req.TargetCmek,
			ClusterId: req.ClusterId,
		}, cluster.GetSharedPoolId(), req.ClusterId)
		if verifyErr != nil {
			return verifyErr
		}
		if !verifyResp.GetValid() {
			return bizerr.New(bizerr.ErrCodeParams, "CMEK verification failed: %s", verifyResp.GetFailReason())
		}

		// Dedup: check if an EncryptionKey with the same master_key already exists for this tenant+region.
		masterKey := resolveTargetMasterKey(req)
		existingKey, findErr := s.encryptionKeyRepo.FindByMasterKeyAndTenant(ctx, masterKey, cluster.GetOrganizationId(), cluster.RegionId)
		if findErr != nil && findErr.ErrCode() != bizerr.ErrCodeRecordNotFound {
			return bizerr.WrapError(findErr, "failed to check existing encryption key")
		}
		if existingKey != nil {
			targetKeyId = existingKey.Id
		} else {
			createdCmek, createErr := s.CreateCustomerManagedEncryptionKey(ctx,
				req.TargetCmek,
				s.resolveProviderIdentityId(ctx, req.TargetCmek),
				cluster.GetSharedPoolId(),
				cluster.GetDataplaneId(),
			)
			if createErr != nil {
				return createErr
			}
			targetKeyId = createdCmek.GetCustomerManagedEncryptionKeyId()
		}
	} else if req.TargetDefaultKey {
		identity, identityErr := s.regionService.PickRegionProviderIdentity(ctx, cluster.RegionId, cluster.GetOrganizationId())
		if identityErr != nil {
			return bizerr.WrapError(identityErr, "failed to pick region provider identity")
		}
		defaultKey, dkErr := s.CreateDefaultKeyOrGrantPermission(ctx, cluster.GetOrganizationId(), identity.Id, cluster.GetDataplaneId())
		if dkErr != nil {
			return dkErr
		}
		targetKeyId = defaultKey.Id
	}
	// For disable: targetKeyId remains 0

	// 7. Determine cse-ctl mode
	cseCtlMode := determineCseCtlMode(currentKey, req)

	// Finish all validation, fast return if necessary.
	if req.ValidateOnly {
		return nil
	}

	metadataJson, merr := metadata.MarshalRequestMetaFromCtx(ctx)
	if merr != nil {
		tilog.Warn(ctx, "failed to marshal request meta", zap.Error(merr))
	}

	binding := &entity.ClusterEncryptionBinding{
		ClusterId:       req.ClusterId,
		TenantId:        cluster.GetOrganizationId(),
		EncryptionKeyId: targetKeyId,
		State:           entity.ClusterEncryptionBindingStateModifying,
		Message:         fmt.Sprintf("starting %s", transitionType),
	}

	// 8. Upsert binding + launch workflow in the same transaction
	tErr := s.transaction.Transaction(ctx, func(ctx context.Context) error {
		if upsertErr := s.bindingRepo.Upsert(ctx, binding); upsertErr != nil {
			return bizerr.WrapError(upsertErr, "failed to upsert cluster encryption binding")
		}

		tilog.Info(ctx, "cluster encryption binding upserted",
			zap.Uint64("cluster_id", req.ClusterId),
			zap.String("transition_type", transitionType),
			zap.String("cse_ctl_mode", cseCtlMode),
		)

		workflowErr := s.workflow.RunDagWithBizId(
			ctx,
			common.UpdateClusterEncryption,
			strconv.FormatUint(cluster.GetOrganizationId(), 10),
			"0",
			strconv.FormatUint(req.ClusterId, 10),
			map[string]string{
				workflow.KeyTraceId:       workflow.GetTraceId(ctx),
				workflow.KeyRequestMeta:   string(metadataJson),
				common.KeyClusterId:       strconv.FormatUint(req.ClusterId, 10),
				common.KeyEncryptionKeyId: strconv.FormatUint(targetKeyId, 10),
				common.KeyEncryptionMode:  cseCtlMode,
				common.KeyTransitionType:  transitionType,
			},
		)
		if workflowErr != nil {
			return bizerr.WrapError(workflowErr, "failed to start encryption change workflow")
		}

		return nil
	})
	if tErr != nil {
		return bizerr.WrapErrorWithCode(tErr, bizerr.GeneralInternalErrorCode, "failed to switch cluster encryption key")
	}

	return nil
}

func determineTransitionType(currentKey *entity.EncryptionKey, req *SwitchEncryptionRequest) (string, bizerr.BizError) {
	if currentKey == nil || currentKey.Id == 0 {
		// No encryption → enabling or already disabled
		if req.TargetCmek != nil {
			return "enable-cmek", nil
		}
		if req.TargetDefaultKey {
			return "enable-default-key", nil
		}
		if req.TargetDisable {
			return "no-op", nil
		}
		return "", bizerr.New(bizerr.ErrCodeParams, "must specify target_cmek, target_default_key, or target_disable")
	}
	// Already encrypted → switching or disabling
	if req.TargetDisable {
		return "disable-cmek", nil
	}

	// Same-key no-op: target matches current encryption — nothing to do.
	if req.TargetDefaultKey && !currentKey.CustomerManaged {
		return "no-op", nil
	}
	if req.TargetCmek != nil && currentKey.CustomerManaged {
		if targetMasterKey := resolveTargetMasterKey(req); targetMasterKey != "" && currentKey.MasterKey == targetMasterKey {
			return "no-op", nil
		}
	}

	if req.TargetCmek != nil {
		return "switch-encryption-key", nil
	}
	if req.TargetDefaultKey {
		return "switch-encryption-key", nil
	}
	return "", bizerr.New(bizerr.ErrCodeParams, "must specify target_cmek, target_default_key, or target_disable")
}

func validateTransition(transitionType string) bizerr.BizError {
	switch transitionType {
	// TODO: Support disable cmek in kernel layer. Planned but not started.
	case "disable-cmek":
		return bizerr.New(bizerr.ErrCodeParams, "cannot disable encryption: this operation is currently not allowed")
	default:
		return nil
	}
}

func determineCseCtlMode(currentKey *entity.EncryptionKey, _ *SwitchEncryptionRequest) string {
	// cse-ctl gen-meta only supports "enable-cmek" and "rotate-data-key" modes.
	// "enable-cmek" handles both first-time enablement and re-keying existing encryption.
	return "enable-cmek"
}

func resolveTargetMasterKey(req *SwitchEncryptionRequest) string {
	if req.TargetCmek != nil {
		if aws := req.TargetCmek.GetAwsKms(); aws != nil {
			return aws.GetKmsKeyArn()
		}
		if aliyun := req.TargetCmek.GetAliyunKms(); aliyun != nil {
			return aliyun.GetKmsKeyArn()
		}
	}
	return "default_key"
}

func (s *service) resolveProviderIdentityId(ctx context.Context, cmek *globalngv1.CustomerManagedEncryptionKey) uint64 {
	cloudRegion, rErr := s.regionRepo.SelectRegion(ctx, region.Filter{
		CloudProvider: cmek.GetCloudProviderId(),
		RegionId:      cmek.GetRegionId(),
		ServicePlan:   cmek.GetServicePlan(),
		TenantId:      cmek.GetOrganizationId(),
	})
	if rErr != nil || cloudRegion == nil {
		return 0
	}
	identity, err := s.regionService.PickRegionProviderIdentity(ctx, cloudRegion.RegionID, cmek.GetOrganizationId())
	if err != nil {
		return 0
	}
	return identity.Id
}

func (s *service) GetEncryptionBinding(ctx context.Context, clusterId uint64) (*entity.ClusterEncryptionBinding, bizerr.BizError) {
	return s.bindingRepo.LoadByClusterId(ctx, clusterId)
}

func (s *service) BuildCurrentEncryption(ctx context.Context, encryptionKeyId uint64) (*globalngv1.Cluster_DualLayerDataEncryption, bizerr.BizError) {
	key, err := s.encryptionKeyRepo.Load(ctx, encryptionKeyId, false)
	if err != nil {
		return nil, bizerr.WrapError(err, "failed to load encryption key %d", encryptionKeyId)
	}
	if !key.CustomerManaged {
		return &globalngv1.Cluster_DualLayerDataEncryption{
			Key: &globalngv1.Cluster_DualLayerDataEncryption_DefaultKey{},
		}, nil
	}
	providerIdentity, err := s.regionProviderIdentityRepo.Load(ctx, key.ProviderIdentityId)
	if err != nil {
		return nil, bizerr.WrapError(err, "failed to load provider identity for key %d", encryptionKeyId)
	}
	cmek, err := s.CvtEncryptionKeyToCmek(ctx, key, providerIdentity)
	if err != nil {
		return nil, err
	}
	return &globalngv1.Cluster_DualLayerDataEncryption{
		Key: &globalngv1.Cluster_DualLayerDataEncryption_Cmek{Cmek: cmek},
	}, nil
}
