package aks

import (
	"context"
	"os"
	"time"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v6"
	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/storage/armstorage"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/sas"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/service"
	"github.com/tidbcloud/pkgv2/bizerr"
)

type AKSClient interface {
	ListClusterAdminCredentials(ctx context.Context, resourceGroupName, clusterName string, options *armcontainerservice.ManagedClustersClientListClusterAdminCredentialsOptions) (armcontainerservice.ManagedClustersClientListClusterAdminCredentialsResponse, error)
	ListClusterUserCredentials(ctx context.Context, resourceGroupName, clusterName string, options *armcontainerservice.ManagedClustersClientListClusterUserCredentialsOptions) (armcontainerservice.ManagedClustersClientListClusterUserCredentialsResponse, error)
	GenBlobContainerSASToken(ctx context.Context, resourceGroupName, storageAccountName, containerName string) (string, error)
	GetStorageAccountKey(ctx context.Context, resourceGroupName, storageAccountName string) (string, error)
}

type aksClient struct {
	SubscriptionID string

	cli *armcontainerservice.ManagedClustersClient
}

func NewAKSClient(subscriptionID string) (AKSClient, error) {
	cred, err := azidentity.NewDefaultAzureCredential(nil)
	if err != nil {
		return nil, err
	}
	containerserviceClientFactory, err := armcontainerservice.NewClientFactory(subscriptionID, cred, nil)
	if err != nil {
		return nil, err
	}
	managedClustersClient := containerserviceClientFactory.NewManagedClustersClient()
	return &aksClient{
		SubscriptionID: subscriptionID,
		cli:            managedClustersClient,
	}, nil
}

func (c *aksClient) ListClusterAdminCredentials(ctx context.Context, resourceGroupName, clusterName string, options *armcontainerservice.ManagedClustersClientListClusterAdminCredentialsOptions) (armcontainerservice.ManagedClustersClientListClusterAdminCredentialsResponse, error) {
	return c.cli.ListClusterAdminCredentials(ctx, resourceGroupName, clusterName, options)
}

func (c *aksClient) ListClusterUserCredentials(ctx context.Context, resourceGroupName, clusterName string, options *armcontainerservice.ManagedClustersClientListClusterUserCredentialsOptions) (armcontainerservice.ManagedClustersClientListClusterUserCredentialsResponse, error) {
	return c.cli.ListClusterUserCredentials(ctx, resourceGroupName, clusterName, options)
}

func getCredential() (azcore.TokenCredential, error) {
	tokenFile := os.Getenv("AZURE_FEDERATED_TOKEN_FILE")
	if tokenFile == "" { // no federated token, use default credential
		return azidentity.NewDefaultAzureCredential(nil)
	}
	opts := &azidentity.WorkloadIdentityCredentialOptions{
		ClientID:      os.Getenv("AZURE_CLIENT_ID"),
		TenantID:      os.Getenv("AZURE_TENANT_ID"),
		TokenFilePath: tokenFile,
	}
	return azidentity.NewWorkloadIdentityCredential(opts)
}

func getStorageClientFactory(subscriptionID string) (*armstorage.ClientFactory, error) {
	cred, err := getCredential()
	if err != nil {
		return nil, err
	}
	return armstorage.NewClientFactory(subscriptionID, cred, nil)
}

func (c *aksClient) GetStorageAccountKey(ctx context.Context, resourceGroupName, storageAccountName string) (string, error) {
	factory, err := getStorageClientFactory(c.SubscriptionID)
	if err != nil {
		return "", bizerr.WrapError(err, "get client factory")
	}
	accountsClient := factory.NewAccountsClient()
	resp, err := accountsClient.ListKeys(ctx, resourceGroupName, storageAccountName, nil)
	if err != nil {
		return "", bizerr.WrapError(err, "list keys")
	}
	for _, key := range resp.Keys {
		if key.Value != nil {
			return *key.Value, nil
		}
	}
	return "", bizerr.WrapError(err, "no storage account keys found")
}

func (c *aksClient) GenBlobContainerSASToken(ctx context.Context, resourceGroupName, storageAccountName, containerName string) (string, error) {
	factory, err := getStorageClientFactory(c.SubscriptionID)
	if err != nil {
		return "", bizerr.WrapError(err, "get client factory")
	}
	accountsClient := factory.NewAccountsClient()
	resp, err := accountsClient.ListKeys(ctx, resourceGroupName, storageAccountName, nil)
	if err != nil {
		return "", bizerr.WrapError(err, "list keys")
	}
	if len(resp.Keys) == 0 {
		return "", bizerr.WrapError(err, "no keys found")
	}
	// There are two keys for a storage account
	// The first key is used in normal case
	// The second key is used for rotation
	key := resp.Keys[0]
	cred, err := service.NewSharedKeyCredential(storageAccountName, *key.Value)
	if err != nil {
		return "", bizerr.WrapError(err, "new shared key credential")
	}

	// generate sas token with 30 days expiry
	start := time.Now().UTC()
	expiry := time.Now().UTC().Add(30 * 24 * time.Hour)

	// use container level permissions
	permissions := sas.BlobPermissions{
		Read:   true,
		Add:    true,
		Create: true,
		Write:  true,
		Delete: true,
		List:   true,
	}

	// generate container level sas token
	params := sas.BlobSignatureValues{
		Protocol:      sas.ProtocolHTTPS,
		StartTime:     start,
		ExpiryTime:    expiry,
		Permissions:   permissions.String(),
		ContainerName: containerName,
	}

	// generate container level sas token
	sas, err := params.SignWithSharedKey(cred)
	if err != nil {
		return "", bizerr.WrapError(err, "sign with shared key")
	}
	return sas.Encode(), nil
}