package encryptionkey

import (
	"context"
	"strconv"
	"strings"

	"github.com/aws/aws-sdk-go-v2/aws/arn"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globalngv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/adapter/regionalserver"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	encryptionkeyrepo "github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/encryptionkey"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/region"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
)

func (s *service) GetCmekAccessIamPrincipal(ctx context.Context,
	request *globalngv1.GetCmekAccessIamPrincipalRequest) (*globalngv1.CmekAccessIamPrincipal, bizerr.BizError) {
	if request.GetServicePlan() != globalcommonv1.ServicePlan_Premium {
		return nil, bizerr.New(bizerr.ErrCodeParams, "only support Premium on AWS or Aliyun")
	}
	provider := request.GetCloudProviderId()
	if provider != globalcommonv1.CloudProvider_AWS && provider != globalcommonv1.CloudProvider_ALICLOUD {
		return nil, bizerr.New(bizerr.ErrCodeParams, "only support Premium on AWS or Aliyun")
	}

	cloudRegion, err := s.regionRepo.SelectRegion(ctx, region.Filter{
		CloudProvider: provider,
		RegionId:      request.GetRegionId(),
		ServicePlan:   request.GetServicePlan(),
		TenantId:      request.GetOrganizationId(),
	})

	if err != nil {
		return nil, bizerr.WrapError(err, "select region failed")
	}
	if cloudRegion == nil {
		return nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "region not found")
	}

	identity, err := s.regionService.PickRegionProviderIdentity(ctx, cloudRegion.RegionID, request.GetOrganizationId())
	if err != nil {
		return nil, bizerr.WrapError(err, "pick region provider identity failed")
	}

	resp := &globalngv1.CmekAccessIamPrincipal{
		OrganizationId:  request.GetOrganizationId(),
		CloudProviderId: provider,
		RegionId:        request.GetRegionId(),
		ServicePlan:     request.GetServicePlan(),
	}
	switch provider {
	case globalcommonv1.CloudProvider_AWS:
		resp.Principal = &globalngv1.CmekAccessIamPrincipal_Aws{
			Aws: &globalngv1.CmekAccessIamPrincipal_AwsCmekPrincipal{
				AccountId:  identity.GetIdentity(),
				ExternalId: strconv.FormatUint(request.GetOrganizationId(), 10),
			},
		}
	case globalcommonv1.CloudProvider_ALICLOUD:
		resp.Principal = &globalngv1.CmekAccessIamPrincipal_Aliyun{
			Aliyun: &globalngv1.CmekAccessIamPrincipal_AliyunCmekPrincipal{
				AccountId: aliyunAccountID(identity.GetIdentity()),
			},
		}
	}
	return resp, nil
}

func (s *service) VerifyCmekAccessIamPrincipal(ctx context.Context,
	request *globalngv1.VerifyCmekAccessIamPrincipalRequest,
) (*globalngv1.VerifyCmekAccessIamPrincipalResponse, bizerr.BizError) {
	return s.verifyCmekAccessIamPrincipal(ctx, request, "", request.GetClusterId())
}

type kmsCheckTarget struct {
	sharedPoolID string
}

func (s *service) resolveCheckKmsTarget(
	ctx context.Context,
	key *globalngv1.CustomerManagedEncryptionKey,
	sharedPoolId string,
	clusterId uint64,
) (kmsCheckTarget, bizerr.BizError) {
	if sharedPoolId != "" {
		return kmsCheckTarget{
			sharedPoolID: sharedPoolId,
		}, nil
	}
	if clusterId != 0 {
		cluster, err := s.clusterRepo.Get(ctx, clusterId, false)
		if err != nil {
			return kmsCheckTarget{}, bizerr.WrapError(err, "failed to load cluster %d", clusterId)
		}
		if cluster.GetSharedPoolId() == "" {
			return kmsCheckTarget{}, bizerr.New(bizerr.ErrCodeParams, "cluster %d does not have shared pool id", clusterId)
		}
		return kmsCheckTarget{
			sharedPoolID: cluster.GetSharedPoolId(),
		}, nil
	}

	if s.resourcePoolResolver == nil {
		return kmsCheckTarget{}, bizerr.New(bizerr.ErrCodeParams, "resource pool resolver is required for kms check")
	}
	resourcePool, err := s.resourcePoolResolver.GetOrCreateDefaultResourcePool(ctx, entity.ResourcePool{
		ResourcePool: &globalngv1.ResourcePool{
			ProviderId:  key.GetCloudProviderId(),
			RegionId:    key.GetRegionId(),
			OwnerType:   globalngv1.ResourcePool_ORGANIZATION,
			OwnerId:     key.GetOrganizationId(),
			ServicePlan: key.GetServicePlan(),
		},
	}, 0)
	if err != nil {
		return kmsCheckTarget{}, err
	}
	if resourcePool.GetId() == 0 {
		return kmsCheckTarget{}, bizerr.New(bizerr.ErrCodeParams, "default shared pool id is required for kms check")
	}
	return kmsCheckTarget{
		sharedPoolID: strconv.FormatUint(resourcePool.GetId(), 10),
	}, nil
}

func (s *service) verifyCmekAccessIamPrincipal(
	ctx context.Context,
	request *globalngv1.VerifyCmekAccessIamPrincipalRequest,
	sharedPoolId string,
	clusterId uint64,
) (*globalngv1.VerifyCmekAccessIamPrincipalResponse, bizerr.BizError) {
	key := request.GetKey()
	provider := key.GetCloudProviderId()
	if key.GetServicePlan() != globalcommonv1.ServicePlan_Premium {
		return nil, bizerr.New(bizerr.ErrCodeParams, "only support Premium on AWS or Aliyun")
	}
	if provider != globalcommonv1.CloudProvider_AWS && provider != globalcommonv1.CloudProvider_ALICLOUD {
		return nil, bizerr.New(bizerr.ErrCodeParams, "only support Premium on AWS or Aliyun")
	}

	cloudRegion, err := s.regionRepo.SelectRegion(ctx, region.Filter{
		CloudProvider: provider,
		RegionId:      key.GetRegionId(),
		ServicePlan:   key.GetServicePlan(),
		TenantId:      key.GetOrganizationId(),
	})

	if err != nil {
		return nil, bizerr.WrapError(err, "select region failed")
	}
	if cloudRegion == nil {
		return nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "region not found")
	}

	identity, err := s.regionService.PickRegionProviderIdentity(ctx, cloudRegion.RegionID, key.GetOrganizationId())
	if err != nil {
		return nil, bizerr.WrapError(err, "pick region provider identity failed")
	}

	switch provider {
	case globalcommonv1.CloudProvider_AWS:
		requestAccountId := key.GetAwsPrincipal().GetAccountId()
		if requestAccountId == "" {
			// Fallback: request omitted principal, use region org identity
			tilog.Info(ctx, "aws principal not provided in request, falling back to region org identity",
				zap.String("identity", identity.GetIdentity()))
		} else if identity.GetIdentity() != requestAccountId {
			return nil, bizerr.New(bizerr.ErrCodeParams, "account not matched")
		}
		parsedArn, parseErr := arn.Parse(key.GetAwsKms().GetKmsKeyArn())
		if parseErr != nil {
			tilog.Error(ctx, "failed to parse arn", zap.Error(parseErr))
			return &globalngv1.VerifyCmekAccessIamPrincipalResponse{
				Valid:      false,
				FailReason: "not a valid kms key arn",
			}, nil
		}
		if parsedArn.Service != "kms" {
			tilog.Error(ctx, "not a valid kms key arn")
			return &globalngv1.VerifyCmekAccessIamPrincipalResponse{
				Valid:      false,
				FailReason: "not a valid kms key arn",
			}, nil
		}
		// Data-plane check: ask the region to actually exercise the key.
		// Runs last so trivially-invalid inputs never reach the network.
		checkTarget, err := s.resolveCheckKmsTarget(ctx, key, sharedPoolId, clusterId)
		if err != nil {
			return nil, err
		}
		if checkErr := s.nextgenCli.CheckKms(ctx, cloudRegion, &regionalserver.CheckKmsRequest{
			SharedPoolId: checkTarget.sharedPoolID,
			Provider:     "aws",
			KeyId:        parsedArn.String(),
			Region:       parsedArn.Region,
			EncryptionContext: map[string]string{
				"ExternalId": strconv.FormatUint(key.GetOrganizationId(), 10),
			},
		}); checkErr != nil {
			tilog.Warn(ctx, "data-plane kms check failed", zap.Error(checkErr))
			return nil, s.kmsVerifyFailureError(ctx, cloudRegion, checkTarget, checkErr, "data-plane kms check failed")
		}
	case globalcommonv1.CloudProvider_ALICLOUD:
		identityAccountID := aliyunAccountID(identity.GetIdentity())
		requestAccountID := key.GetAliyunPrincipal().GetAccountId()
		if requestAccountID == "" {
			// Fallback: request omitted principal, use region org identity
			tilog.Info(ctx, "alicloud principal not provided in request, falling back to region org identity",
				zap.String("identity", identityAccountID))
		} else if identityAccountID != aliyunAccountID(requestAccountID) {
			return nil, bizerr.New(bizerr.ErrCodeParams, "account not matched")
		}
		aliyunKms := key.GetAliyunKms()
		if aliyunKms.GetKmsKeyArn() == "" {
			return &globalngv1.VerifyCmekAccessIamPrincipalResponse{
				Valid:      false,
				FailReason: "kms_key_arn is required",
			}, nil
		}
		// Data-plane check: ask the region to actually encrypt/decrypt with the
		// key. Runs last so trivially-invalid inputs never reach the network.
		checkTarget, err := s.resolveCheckKmsTarget(ctx, key, sharedPoolId, clusterId)
		if err != nil {
			return nil, err
		}
		if checkErr := s.nextgenCli.CheckKms(ctx, cloudRegion, &regionalserver.CheckKmsRequest{
			SharedPoolId: checkTarget.sharedPoolID,
			Provider:     "alicloud",
			KeyId:        aliyunKms.GetKmsKeyArn(),
			Region:       cloudRegion.GetId().String(),
		}); checkErr != nil {
			tilog.Warn(ctx, "data-plane kms check failed", zap.Error(checkErr))
			return nil, s.kmsVerifyFailureError(ctx, cloudRegion, checkTarget, checkErr, "data-plane kms check failed")
		}
	}

	return &globalngv1.VerifyCmekAccessIamPrincipalResponse{
		Valid: true,
	}, nil
}

func (s *service) GetCustomerManagedEncryptionKey(ctx context.Context,
	request *globalngv1.GetCustomerManagedEncryptionKeyRequest) (*globalngv1.CustomerManagedEncryptionKey, bizerr.BizError) {
	cmek, err := s.encryptionKeyRepo.Load(ctx, request.GetCustomerManagedEncryptionKeyId(), false)
	if err != nil {
		return nil, bizerr.WrapError(err, "get customer managed encryption key failed")
	}
	if !cmek.CustomerManaged {
		return nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "customer managed encryption key not found")
	}
	providerIdentity, err := s.regionProviderIdentityRepo.Load(ctx, cmek.ProviderIdentityId)
	if err != nil {
		return nil, bizerr.WrapError(err, "get customer managed encryption key failed")
	}
	return s.CvtEncryptionKeyToCmek(ctx, cmek, providerIdentity)
}

func (s *service) CvtEncryptionKeyToCmek(
	ctx context.Context,
	key *entity.EncryptionKey,
	providerIdentity *entity.RegionProviderIdentity,
) (*globalngv1.CustomerManagedEncryptionKey, bizerr.BizError) {
	if key == nil || providerIdentity == nil {
		tilog.Error(ctx, "encryption and provider identity cannot be nil")
		return nil, bizerr.New(bizerr.ErrCodeParams, "encryption and provider identity cannot be nil")
	}
	cmek := &globalngv1.CustomerManagedEncryptionKey{
		CustomerManagedEncryptionKeyId: key.Id,
		OrganizationId:                 key.TenantId,
		CloudProviderId:                key.ProviderName,
		RegionId:                       key.RegionName,
		ServicePlan:                    key.ServicePlan,
	}

	switch key.ProviderName {
	case globalcommonv1.CloudProvider_AWS:
		cmek.Key = &globalngv1.CustomerManagedEncryptionKey_AwsKms_{
			AwsKms: &globalngv1.CustomerManagedEncryptionKey_AwsKms{
				KmsKeyArn: key.MasterKey,
			},
		}
		cmek.Principal = &globalngv1.CustomerManagedEncryptionKey_AwsPrincipal{
			AwsPrincipal: &globalngv1.CmekAccessIamPrincipal_AwsCmekPrincipal{
				AccountId:  providerIdentity.Identity,
				ExternalId: strconv.FormatUint(key.TenantId, 10),
			},
		}
	case globalcommonv1.CloudProvider_ALICLOUD:
		cmek.Key = &globalngv1.CustomerManagedEncryptionKey_AliyunKms_{
			AliyunKms: &globalngv1.CustomerManagedEncryptionKey_AliyunKms{
				KmsKeyArn: key.MasterKey,
			},
		}
		cmek.Principal = &globalngv1.CustomerManagedEncryptionKey_AliyunPrincipal{
			AliyunPrincipal: &globalngv1.CmekAccessIamPrincipal_AliyunCmekPrincipal{
				AccountId: aliyunAccountID(providerIdentity.Identity),
			},
		}
	}
	return cmek, nil
}

func (s *service) CreateCustomerManagedEncryptionKey(ctx context.Context,
	cmek *globalngv1.CustomerManagedEncryptionKey,
	providerIdentityId uint64,
	sharedPoolId string,
	dataplaneId string,
) (*globalngv1.CustomerManagedEncryptionKey, bizerr.BizError) {
	if providerIdentityId == 0 {
		return nil, bizerr.New(bizerr.ErrCodeParams, "providerIdentity is required")
	}

	provider := cmek.GetCloudProviderId()
	if cmek.GetServicePlan() != globalcommonv1.ServicePlan_Premium {
		return nil, bizerr.New(bizerr.ErrCodeParams, "only support Premium on AWS or Aliyun")
	}
	if provider != globalcommonv1.CloudProvider_AWS && provider != globalcommonv1.CloudProvider_ALICLOUD {
		return nil, bizerr.New(bizerr.ErrCodeParams, "only support Premium on AWS or Aliyun")
	}
	cloudRegion, err := s.regionRepo.SelectRegion(ctx, region.Filter{
		CloudProvider: provider,
		RegionId:      cmek.GetRegionId(),
		ServicePlan:   cmek.GetServicePlan(),
		TenantId:      cmek.GetOrganizationId(),
	})

	if err != nil {
		return nil, bizerr.WrapError(err, "select region failed")
	}
	if cloudRegion == nil {
		return nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "region not found")
	}

	if err := s.ValidateCustomerManagedEncryptionKey(ctx, cmek, sharedPoolId); err != nil {
		return nil, err
	}
	tErr := s.transaction.Transaction(ctx, func(ctx context.Context) error {
		id := s.idGeneratorAdapter.MustNew()
		encryptionKey := &entity.EncryptionKey{
			Id:                 id,
			TenantId:           cmek.GetOrganizationId(),
			ProviderName:       provider,
			RegionName:         cmek.GetRegionId(),
			ServicePlan:        cmek.GetServicePlan(),
			ProviderIdentityId: providerIdentityId,
			CustomerManaged:    true,
			MasterKey:          resolveMasterKeyString(cmek),
			InitialDataplaneId: dataplaneId,
			ProviderId:         cloudRegion.ProviderID,
			RegionId:           cloudRegion.RegionID,
		}
		innerErr := s.encryptionKeyRepo.Create(ctx, encryptionKey)
		if innerErr != nil {
			return innerErr
		}
		cmek.CustomerManagedEncryptionKeyId = encryptionKey.Id
		return nil
	})

	if tErr != nil {
		return nil, bizerr.WrapError(tErr, "createDefaultKey customer managed encryption key failed")
	}
	return cmek, nil
}

func (s *service) ValidateCustomerManagedEncryptionKey(ctx context.Context,
	cmek *globalngv1.CustomerManagedEncryptionKey,
	sharedPoolId string,
) bizerr.BizError {
	verifyResult, verifyErr := s.verifyCmekAccessIamPrincipal(ctx, &globalngv1.VerifyCmekAccessIamPrincipalRequest{
		Key: cmek,
	}, sharedPoolId, 0)
	if verifyErr != nil {
		return verifyErr
	}
	if verifyResult == nil {
		return bizerr.New(bizerr.GeneralInternalErrorCode, "verify cmek failed: empty response")
	}
	if !verifyResult.Valid {
		return bizerr.New(bizerr.ErrCodeParams, "verify cmek failed, reason: %s", verifyResult.FailReason)
	}
	return nil
}

// resolveMasterKeyString returns the master key string to store in the database.
// For both AWS and Aliyun, this is the KMS key ARN.
func resolveMasterKeyString(cmek *globalngv1.CustomerManagedEncryptionKey) string {
	switch cmek.GetCloudProviderId() {
	case globalcommonv1.CloudProvider_AWS:
		return cmek.GetAwsKms().GetKmsKeyArn()
	case globalcommonv1.CloudProvider_ALICLOUD:
		return cmek.GetAliyunKms().GetKmsKeyArn()
	}
	return ""
}

func aliyunAccountID(identity string) string {
	identity = strings.TrimSpace(identity)
	if identity == "" {
		return identity
	}
	if _, err := strconv.ParseUint(identity, 10, 64); err == nil {
		return identity
	}
	const ramPrefix = "acs:ram::"
	if after, ok := strings.CutPrefix(identity, ramPrefix); ok {
		accountID, _, found := strings.Cut(after, ":")
		if found && accountID != "" {
			return accountID
		}
	}
	return identity
}

func (s *service) GetCustomerManagedEncryptionKeyByClusterId(ctx context.Context,
	request *globalngv1.GetCustomerManagedEncryptionKeyByClusterIdRequest,
) (*globalngv1.CustomerManagedEncryptionKey, bizerr.BizError) {
	cluster, err := s.clusterRepo.Get(ctx, request.GetClusterId(), false)
	if err != nil {
		return nil, bizerr.WrapError(err, "failed to load cluster %d", request.GetClusterId())
	}
	if cluster.EncryptionKeyId == 0 {
		return nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "cluster %d has no encryption key", request.GetClusterId())
	}
	key, err := s.encryptionKeyRepo.Load(ctx, cluster.EncryptionKeyId, false)
	if err != nil {
		return nil, bizerr.WrapError(err, "failed to load encryption key %d", cluster.EncryptionKeyId)
	}
	if !key.CustomerManaged {
		return nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "cluster %d does not use a customer-managed encryption key", request.GetClusterId())
	}
	providerIdentity, err := s.regionProviderIdentityRepo.Load(ctx, key.ProviderIdentityId)
	if err != nil {
		return nil, bizerr.WrapError(err, "failed to load provider identity")
	}
	return s.CvtEncryptionKeyToCmek(ctx, key, providerIdentity)
}

func (s *service) ListCustomerManagedEncryptionKeys(ctx context.Context,
	request *globalngv1.ListCustomerManagedEncryptionKeysRequest,
) (*globalngv1.ListCustomerManagedEncryptionKeysResponse, bizerr.BizError) {
	customerManaged := true
	filter := &encryptionkeyrepo.Filter{
		TenantId:        request.GetOrganizationId(),
		Ids:             request.GetCustomerManagedEncryptionKeyIds(),
		CustomerManaged: &customerManaged,
	}
	if len(request.GetCloudProviderIds()) == 1 {
		filter.ProviderName = request.GetCloudProviderIds()[0]
	}
	if len(request.GetRegionIds()) == 1 {
		filter.RegionName = request.GetRegionIds()[0]
	}
	if len(request.GetServicePlans()) == 1 {
		filter.ServicePlan = request.GetServicePlans()[0]
	}
	keys, err := s.encryptionKeyRepo.ListEncryptionKeys(ctx, filter)
	if err != nil {
		return nil, bizerr.WrapError(err, "failed to list encryption keys")
	}

	var cmeks []*globalngv1.CustomerManagedEncryptionKey
	for _, key := range keys {
		providerIdentity, piErr := s.regionProviderIdentityRepo.Load(ctx, key.ProviderIdentityId)
		if piErr != nil {
			tilog.Warn(ctx, "failed to load provider identity for key",
				zap.Uint64("key_id", key.Id),
				zap.Error(piErr))
			continue
		}
		cmek, cvtErr := s.CvtEncryptionKeyToCmek(ctx, key, providerIdentity)
		if cvtErr != nil {
			tilog.Warn(ctx, "failed to convert encryption key",
				zap.Uint64("key_id", key.Id),
				zap.Error(cvtErr))
			continue
		}
		cmeks = append(cmeks, cmek)
	}
	return &globalngv1.ListCustomerManagedEncryptionKeysResponse{
		CustomerManagedEncryptionKeys: cmeks,
	}, nil
}