package server

import (
	"bytes"
	"context"
	"fmt"

	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	"github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/azkeys"
	"github.com/samber/lo"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"github.com/tidbcloud/tiinf/monitor/tilog"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apiv1 "github.com/tidbcloud/infra-provider/api/tidbcloud/cluster-server/v1"
	"github.com/tidbcloud/infra-provider/pkg/cloud/azure"
)

const (
	dummyAzureEncryptionKey = "dummyAzureEncryptionKey"
)

func AzureQueryLatestKeyvaultKeyVersion(ctx context.Context, azureCli azure.Client, azureCmek *apiv1.Cluster_SecurityConfigurations_AzureCmek) (string, error) {
	var err error
	cred, err := azidentity.NewWorkloadIdentityCredential(&azidentity.WorkloadIdentityCredentialOptions{
		ClientID:                   azureCmek.ApplicationId,
		AdditionallyAllowedTenants: []string{"*"},
	})
	if err != nil {
		return "", fmt.Errorf("create workload identity credential: %w", err)
	}

	azureCli = azureCli.WithCredential(cred)

	candidates, err := azureCli.QueryKeyvaultKeyVersions(ctx, azureCmek.KeyVaultName, azureCmek.KeyName)
	if err != nil {
		return "", fmt.Errorf("query key versions: %w", err)
	}
	keys := lo.Filter(candidates, func(key *azkeys.KeyProperties, _ int) bool {
		return lo.FromPtr(key.Attributes.Enabled)
	})
	if len(keys) == 0 {
		return "", fmt.Errorf("no enabled key versions found for key %s in vault %s", azureCmek.KeyName, azureCmek.KeyVaultName)
	}

	latestKey := lo.MaxBy(keys, func(a, b *azkeys.KeyProperties) bool {
		return a.Attributes.Created.After(*b.Attributes.Created)
	})
	_, _, keyVersion, err := cloudv1.ParseAzureKeyvaultKeyURI(string(lo.FromPtr(latestKey.KID)))
	return keyVersion, err
}

// azureValidateKeyvaultKeyPermission wrap and then unwrap a dummy key to validate whether the permission grant is proper
func azureValidateKeyvaultKeyPermission(ctx context.Context, azureCli azure.Client, azureCmek *cloudv1.AzureCMEKSpec) (valid bool, message string, _ error) {
	const encryptionAlgorithm = azkeys.EncryptionAlgorithmRSA15
	logger := tilog.WithContextS(ctx)

	cred, err := azidentity.NewWorkloadIdentityCredential(&azidentity.WorkloadIdentityCredentialOptions{
		ClientID:                   azureCmek.ApplicationID,
		AdditionallyAllowedTenants: []string{"*"},
	})
	if err != nil {
		return false, "", status.Errorf(codes.Internal, "create workload identity credential: %s", err)
	}
	azureCli = azureCli.WithCredential(cred)

	dummyKey := []byte(dummyAzureEncryptionKey)
	wraped, err := azureCli.WrapKeyByKeyvaultKey(ctx,
		azureCmek.KeyVaultName,
		azureCmek.KeyName,
		azureCmek.KeyVersion,
		azkeys.KeyOperationParameters{
			Algorithm: lo.ToPtr(encryptionAlgorithm),
			Value:     dummyKey,
		},
	)
	if err != nil {
		logger.Infof("wrap key failed: %s", err)
		return false, "Has no permissions to wrap a key", nil
	}
	unwraped, err := azureCli.UnwrapKeyByKeyvaultKey(ctx,
		azureCmek.KeyVaultName,
		azureCmek.KeyName,
		azureCmek.KeyVersion,
		azkeys.KeyOperationParameters{
			Algorithm: lo.ToPtr(encryptionAlgorithm),
			Value:     wraped.Result,
		},
	)
	if err != nil {
		logger.Infof("unwrap data: %s", err)
		return false, "Has no permissions to unwrap a key", nil
	}

	if !bytes.Equal(unwraped.Result, dummyKey) {
		return false, "", fmt.Errorf("decrypted is not equal to original")
	}

	return true, "OK", nil
}