package cmek

import (
	"context"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"strconv"
	"strings"

	regionaldedicatedv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/regional-server/v1"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globaldedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	"github.com/tidbcloud/tidb-management-service/utilities/featuregate"
)

func (s *customerManagedEncryptionKeyServiceImpl) GetCmekAccessIamPrincipal(
	ctx context.Context,
	req *globaldedicatedv1.GetCmekAccessIamPrincipalRequest,
) (*entity.CmekAccessIamPrincipal, bizerr.BizError) {
	ctx, span := Tracer.Start(ctx, "Domain.GetCmekAccessIamPrincipal")
	defer span.End()

	// check provider and region is valid
	err := s.checkCloudProviderAndRegion(ctx, req.CloudProviderId, req.RegionId, req.OrganizationId)
	if err != nil {
		return nil, err
	}

	// check cmek enable
	project, err := s.accountCli.GetProject(ctx, req.ProjectId)
	if err != nil {
		return nil, err
	}

	// if cmek not enabled, return empty
	if !project.AwsCmekEnabled ||
		!s.featGateCli.CheckFeatureEnabled(
			featuregate.AwsCmek,
			req.OrganizationId,
			req.ProjectId,
			req.GetCloudProviderId().String(),
			req.RegionId.String(),
		) {
		return nil, nil //nolint:nilnil  // nil, nil is a valid return for this method when the feature is not enabled
	}

	// get region
	region, err := s.regionRepo.SelectRegionById(ctx, req.RegionId, req.OrganizationId)
	if err != nil {
		return nil, err
	}

	if region.State != globalcommonv1.Region_ACTIVE {
		return nil, bizerr.New(bizerr.ErrCodePermission, "region is not active")
	}

	if region.CloudProviderId != req.CloudProviderId {
		return nil, bizerr.New(bizerr.ErrCodeParams, "cloud provider id not match")
	}

	// get access iam principal
	account, err := s.regionServerCli.GetAccount(
		ctx,
		region,
		&regionaldedicatedv1.GetAccountRequest{ProjectId: strconv.FormatUint(req.ProjectId, 10)},
	)

	// not found, need to assign account
	if err != nil && err.ErrCode() == bizerr.ErrCodeRecordNotFound {
		account, err = s.regionServerCli.AssignAccount(
			ctx,
			region,
			&regionaldedicatedv1.AssignAccountRequest{
				ProjectId:   strconv.FormatUint(req.ProjectId, 10),
				CmekEnabled: true,
			},
		)

		if err != nil {
			return nil, err
		}
	} else if err != nil {
		return nil, err
	}

	// double check account not nil
	if account == nil {
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "failed to get iam principal")
	}

	principal := &entity.CmekAccessIamPrincipal{}
	switch req.CloudProviderId {
	case globalcommonv1.CloudProvider_AWS:
		principal.AwsPrincipalAccountId = account.GetAccountId()
	case globalcommonv1.CloudProvider_AZURE:
		accountApp := account.GetAzureAppRegistration()
		principal.AzureAppRegistration = entity.AzureAppRegistration{
			AppRegistrationId:   accountApp.GetAppRegistrationId(),
			AppRegistrationName: accountApp.GetAppRegistrationName(),
			AppRegistrationUmi:  accountApp.GetAppRegistrationUmi(),
		}

		suffix, err := s.ensureCmekKeySuffix(ctx, req.ProjectId)
		if err != nil {
			return nil, err
		}
		principal.AzureCmekKeySuffix = suffix
	default:
		return nil, bizerr.New(bizerr.ErrCodeParams, "cloud provider not support")
	}
	return principal, nil
}

// ensureCmekKeySuffix ensures the project has a cmek suffix, if not, create one
func (s *customerManagedEncryptionKeyServiceImpl) ensureCmekKeySuffix(
	ctx context.Context,
	projectId uint64,
) (string, bizerr.BizError) {
	suffixEn, err := s.cmekRepo.GetProjectCmekKeySuffix(ctx, projectId)

	if err != nil && err.ErrCode() == bizerr.ErrCodeRecordNotFound {
		suffixEn, err = generateCmekKeySuffix(projectId)
		if err != nil {
			return "", err
		}

		err = s.cmekRepo.SaveProjectCmekKeySuffix(ctx, suffixEn)
		if err != nil {
			return "", err
		}
	}

	if suffixEn == nil {
		return "", bizerr.New(bizerr.ErrCodeUnknown, "failed to get or create cmek suffix")
	}

	return suffixEn.Suffix, nil
}

// generateCmekKeySuffix generates a cmek suffix for the project
// The suffix is generated by HMAC-SHA256(projectId, salt) and taking the first 16 bytes
// The salt is a random 32 bytes string, encoded in base64 URL encoding
// The suffix is encoded in base64 URL encoding
func generateCmekKeySuffix(
	projectId uint64,
) (suffix *entity.ProjectCmekKeySuffix, err bizerr.BizError) {
	if projectId == 0 {
		return nil, bizerr.New(bizerr.ErrCodeParams, "projectId is required")
	}

	replacer := strings.NewReplacer("+", "A", "/", "B", "=", "C")

	saltBytes := make([]byte, 32)
	if _, err := io.ReadFull(rand.Reader, saltBytes); err != nil {
		return nil, bizerr.New(bizerr.GeneralInternalErrorCode, "failed to generate cmek salt: %s", err.Error())
	}

	salt := replacer.Replace(base64.StdEncoding.EncodeToString(saltBytes))

	projectIdBytes := fmt.Appendf(nil, "%d", projectId)

	h := hmac.New(sha256.New, saltBytes)
	_, writeErr := h.Write(projectIdBytes)
	if writeErr != nil {
		return nil, bizerr.WrapErrorWithCode(
			writeErr,
			bizerr.GeneralInternalErrorCode,
			"failed to write projectId to hmac",
		)
	}

	fullHash := h.Sum(nil)

	suffixStr := replacer.Replace(base64.StdEncoding.EncodeToString(fullHash))[:16]

	return &entity.ProjectCmekKeySuffix{
		ProjectId: projectId,
		Salt:      salt,
		Suffix:    suffixStr,
	}, nil
}

// verifyCmekKeySuffix verifies the provided keyVaultKey matches the stored suffix
func verifyCmekKeySuffix(
	projectId uint64,
	keyVaultKey string,
	suffixEn *entity.ProjectCmekKeySuffix,
) (bool, bizerr.BizError) {
	if suffixEn == nil {
		return false, bizerr.New(bizerr.ErrCodeParams, "cmek suffix not found")
	}

	providedSuffix, extractErr := extractSuffixFromKey(keyVaultKey)
	if extractErr != nil {
		return false, bizerr.WrapErrorWithCode(
			extractErr,
			bizerr.ErrCodeParams,
			"failed to extract cmek suffix from keyVaultKey",
		)
	}

	return projectId == suffixEn.ProjectId && string(providedSuffix) == suffixEn.Suffix, nil
}

func extractSuffixFromKey(keyVaultKey string) ([]byte, error) {
	keyBytes := []byte(keyVaultKey)

	if len(keyBytes) <= 16 {
		return nil, errors.New("keyVaultKey is too short to contain a valid suffix")
	}

	suffixBytes := keyBytes[len(keyBytes)-16:]
	return suffixBytes, nil
}