package azure

import (
	"context"
	"fmt"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute"
	"github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/azkeys"
)

const (
	keyvaultDomainTemplate = "https://%s.vault.azure.net"
)

//go:generate mockgen -package azure -source=client.go -destination=client_mock.go
type Client interface {
	WithCredential(cred azcore.TokenCredential) Client
	RestartVM(ctx context.Context, subscription, rg string, vmss string, instance string) error
	QueryKeyvaultKeyVersions(ctx context.Context, vaultName, keyName string) ([]*azkeys.KeyProperties, error)
	GetKeyvaultKey(ctx context.Context, vaultName, keyName, keyVersion string) (*azkeys.KeyBundle, error)
	DecryptByKeyvaultKey(
		ctx context.Context, vaultName, keyName, keyVersion string, parameters azkeys.KeyOperationParameters,
	) (*azkeys.DecryptResponse, error)
	EncryptByKeyvaultKey(
		ctx context.Context, vaultName, keyName, keyVersion string, parameters azkeys.KeyOperationParameters,
	) (*azkeys.EncryptResponse, error)

	WrapKeyByKeyvaultKey(
		ctx context.Context, vaultName, keyName, keyVersion string, parameters azkeys.KeyOperationParameters,
	) (*azkeys.WrapKeyResponse, error)
	UnwrapKeyByKeyvaultKey(
		ctx context.Context, vaultName, keyName, keyVersion string, parameters azkeys.KeyOperationParameters,
	) (*azkeys.UnwrapKeyResponse, error)
}

type clientSet struct {
	cred azcore.TokenCredential
}

func NewClientSet(cred azcore.TokenCredential) (Client, error) {
	return &clientSet{cred: cred}, nil
}

func (c *clientSet) WithCredential(cred azcore.TokenCredential) Client {
	out := &clientSet{}
	out.cred = cred
	return out
}

func (c *clientSet) RestartVM(ctx context.Context, subscription, rg, vmss, instanceID string) error {
	vm, err := armcompute.NewVirtualMachineScaleSetVMsClient(subscription, c.cred, nil)
	if err != nil {
		return err
	}
	_, err = vm.BeginRestart(ctx, rg, vmss, instanceID, nil)
	return err
}

// QueryKeyvaultkeyVersions queries the key versions of a key in Azure Key Vault.
// It uses workload identity to authenticate and fetches the key versions.
//
//	clientID: The client ID of the Azure workload identity, for exmaple: client ID of Managed Identity, client ID of App registration.
func (c *clientSet) QueryKeyvaultKeyVersions(ctx context.Context, vaultName, keyName string) ([]*azkeys.KeyProperties, error) {
	client, err := azkeys.NewClient(fmt.Sprintf(keyvaultDomainTemplate, vaultName), c.cred, &azkeys.ClientOptions{})
	if err != nil {
		return nil, fmt.Errorf("create Key Vault client: %w", err)
	}

	var keys []*azkeys.KeyProperties
	pager := client.NewListKeyPropertiesVersionsPager(keyName, nil)
	for pager.More() {

		page, err := pager.NextPage(ctx)
		if err != nil {
			return nil, fmt.Errorf("get next page of key versions: %w", err)
		}
		keys = append(keys, page.Value...)
	}

	return keys, nil
}

func (c *clientSet) GetKeyvaultKey(ctx context.Context, vaultName, keyName, keyVersion string) (*azkeys.KeyBundle, error) {
	client, err := azkeys.NewClient(fmt.Sprintf(keyvaultDomainTemplate, vaultName), c.cred, &azkeys.ClientOptions{})
	if err != nil {
		return nil, fmt.Errorf("create Key Vault client: %w", err)
	}

	resp, err := client.GetKey(ctx, keyName, keyVersion, nil)
	if err != nil {
		return nil, fmt.Errorf("get key %s version %s: %w", keyName, keyVersion, err)
	}

	return &resp.KeyBundle, nil
}

func (c *clientSet) DecryptByKeyvaultKey(
	ctx context.Context,
	VaultName, keyName, keyVersion string,
	parameters azkeys.KeyOperationParameters,
) (*azkeys.DecryptResponse, error) {
	client, err := azkeys.NewClient(fmt.Sprintf(keyvaultDomainTemplate, VaultName), c.cred, &azkeys.ClientOptions{})
	if err != nil {
		return nil, fmt.Errorf("create Key Vault client: %w", err)
	}

	resp, err := client.Decrypt(ctx, keyName, keyVersion, parameters, nil)
	if err != nil {
		return nil, fmt.Errorf("decrypt failed: %w", err)
	}

	return &resp, nil
}

func (c *clientSet) EncryptByKeyvaultKey(
	ctx context.Context,
	VaultName, keyName, keyVersion string,
	parameters azkeys.KeyOperationParameters,
) (*azkeys.EncryptResponse, error) {
	client, err := azkeys.NewClient(fmt.Sprintf(keyvaultDomainTemplate, VaultName), c.cred, &azkeys.ClientOptions{})
	if err != nil {
		return nil, fmt.Errorf("create Key Vault client: %w", err)
	}

	resp, err := client.Encrypt(ctx, keyName, keyVersion, parameters, nil)
	if err != nil {
		return nil, fmt.Errorf("decrypt failed: %w", err)
	}

	return &resp, nil
}

func (c *clientSet) WrapKeyByKeyvaultKey(
	ctx context.Context,
	vaultName, keyName, keyVersion string,
	parameters azkeys.KeyOperationParameters,
) (*azkeys.WrapKeyResponse, error) {
	client, err := azkeys.NewClient(fmt.Sprintf(keyvaultDomainTemplate, vaultName), c.cred, &azkeys.ClientOptions{})
	if err != nil {
		return nil, fmt.Errorf("create Key Vault client: %w", err)
	}

	resp, err := client.WrapKey(ctx, keyName, keyVersion, parameters, nil)
	if err != nil {
		return nil, fmt.Errorf("decrypt failed: %w", err)
	}

	return &resp, nil
}

func (c *clientSet) UnwrapKeyByKeyvaultKey(
	ctx context.Context,
	vaultName, keyName, keyVersion string,
	parameters azkeys.KeyOperationParameters,
) (*azkeys.UnwrapKeyResponse, error) {
	client, err := azkeys.NewClient(fmt.Sprintf(keyvaultDomainTemplate, vaultName), c.cred, &azkeys.ClientOptions{})
	if err != nil {
		return nil, fmt.Errorf("create Key Vault client: %w", err)
	}

	resp, err := client.UnwrapKey(ctx, keyName, keyVersion, parameters, nil)
	if err != nil {
		return nil, fmt.Errorf("decrypt failed: %w", err)
	}

	return &resp, nil
}