package storage

import (
	"bytes"
	"context"
	"encoding/base64"
	"errors"
	"fmt"

	openapiutil "github.com/alibabacloud-go/darabonba-openapi/v2/utils"
	alikms "github.com/alibabacloud-go/kms-20160120/v3/client"
	"github.com/alibabacloud-go/tea/dara"
	alicredentials "github.com/aliyun/credentials-go/credentials"
	awskms "github.com/aws/aws-sdk-go-v2/service/kms"
	"github.com/aws/aws-sdk-go-v2/service/kms/types"

	"github.com/tidbcloud/service-proxy/pkg/checker/core"
	"github.com/tidbcloud/service-proxy/pkg/checker/util"
)

type aliCloudKMSClient interface {
	EncryptWithContext(context.Context, *alikms.EncryptRequest, *dara.RuntimeOptions) (*alikms.EncryptResponse, error)
	DecryptWithContext(context.Context, *alikms.DecryptRequest, *dara.RuntimeOptions) (*alikms.DecryptResponse, error)
}

var newAliCloudKMSClient = createAliCloudKMSClient
var newAliCloudCredential = alicredentials.NewCredential

func init() {
	core.DefaultRegistry.Register("kms", func(req *core.CheckRequest) (core.Checker, error) {
		if req.KMS == nil {
			return nil, errors.New("kms params required")
		}
		return NewKMSChecker(req.KMS)
	})
}

// KMSChecker checks KMS key encrypt/decrypt permissions.
type KMSChecker struct {
	params *core.KMSParams
}

// NewKMSChecker creates a new KMS checker.
func NewKMSChecker(params *core.KMSParams) (*KMSChecker, error) {
	if err := validateKMSParams(params); err != nil {
		return nil, err
	}
	return &KMSChecker{params: params}, nil
}

// Type returns the checker type.
func (c *KMSChecker) Type() string {
	return "kms"
}

// Check performs KMS encrypt/decrypt permission check.
func (c *KMSChecker) Check(ctx context.Context) error {
	switch c.params.Provider {
	case "aws":
		return c.checkAWS(ctx)
	case "gcp":
		return errors.New("gcp KMS check not yet implemented")
	case "azure":
		return errors.New("azure KMS check not yet implemented")
	case "alicloud":
		return c.checkAliCloud(ctx)
	default:
		return fmt.Errorf("unsupported provider: %s", c.params.Provider)
	}
}

func validateKMSParams(params *core.KMSParams) error {
	if params == nil {
		return errors.New("kms params are nil")
	}

	if params.Provider == "" {
		return errors.New("provider is required")
	}

	switch params.Provider {
	case "aws":
		if params.KeyID == "" {
			return errors.New("keyId is required for AWS provider")
		}
		if params.Region == "" {
			return errors.New("region is required for AWS provider")
		}
	case "alicloud":
		if params.KeyID == "" {
			return errors.New("keyId is required for AliCloud provider")
		}
		if params.Region == "" {
			return errors.New("region is required for AliCloud provider")
		}
	case "gcp", "azure":
		return nil
	default:
		return fmt.Errorf("unknown provider: %s (supported: aws, gcp, azure, alicloud)", params.Provider)
	}

	return nil
}

func (c *KMSChecker) checkAWS(ctx context.Context) error {
	client, err := c.createClient(ctx)
	if err != nil {
		return fmt.Errorf("failed to create kms client: %w", err)
	}

	// Test data for encrypt/decrypt verification
	testData := []byte("kms-permission-check-test-data")

	// Test encrypt permission
	encryptOutput, err := client.Encrypt(ctx, &awskms.EncryptInput{
		KeyId:               &c.params.KeyID,
		Plaintext:           testData,
		EncryptionAlgorithm: types.EncryptionAlgorithmSpecSymmetricDefault,
		EncryptionContext:   c.params.EncryptionContext,
	})
	if err != nil {
		return fmt.Errorf("failed to encrypt with key %s: %w", c.params.KeyID, err)
	}

	// Test decrypt permission
	decryptOutput, err := client.Decrypt(ctx, &awskms.DecryptInput{
		KeyId:               &c.params.KeyID,
		CiphertextBlob:      encryptOutput.CiphertextBlob,
		EncryptionAlgorithm: types.EncryptionAlgorithmSpecSymmetricDefault,
		EncryptionContext:   c.params.EncryptionContext,
	})
	if err != nil {
		return fmt.Errorf("failed to decrypt with key %s: %w", c.params.KeyID, err)
	}

	// Verify decrypted data matches original
	if !bytes.Equal(decryptOutput.Plaintext, testData) {
		return fmt.Errorf("decrypted data does not match original for key %s", c.params.KeyID)
	}

	return nil
}

func (c *KMSChecker) checkAliCloud(ctx context.Context) error {
	client, err := newAliCloudKMSClient(c.params)
	if err != nil {
		return fmt.Errorf("failed to create alicloud kms client: %w", err)
	}

	testData := []byte("kms-permission-check-test-data")
	encodedPlaintext := base64.StdEncoding.EncodeToString(testData)

	encryptOutput, err := client.EncryptWithContext(ctx, &alikms.EncryptRequest{
		KeyId:             &c.params.KeyID,
		Plaintext:         &encodedPlaintext,
		EncryptionContext: buildAliCloudEncryptionContext(c.params.EncryptionContext),
	}, &dara.RuntimeOptions{})
	if err != nil {
		return fmt.Errorf("failed to encrypt with alicloud key %s: %w", c.params.KeyID, err)
	}
	if encryptOutput == nil || encryptOutput.Body == nil || encryptOutput.Body.CiphertextBlob == nil {
		return fmt.Errorf("alicloud kms encrypt response missing ciphertext for key %s", c.params.KeyID)
	}

	decryptOutput, err := client.DecryptWithContext(ctx, &alikms.DecryptRequest{
		CiphertextBlob:    encryptOutput.Body.CiphertextBlob,
		EncryptionContext: buildAliCloudEncryptionContext(c.params.EncryptionContext),
	}, &dara.RuntimeOptions{})
	if err != nil {
		return fmt.Errorf("failed to decrypt with alicloud key %s: %w", c.params.KeyID, err)
	}
	if decryptOutput == nil || decryptOutput.Body == nil || decryptOutput.Body.Plaintext == nil {
		return fmt.Errorf("alicloud kms decrypt response missing plaintext for key %s", c.params.KeyID)
	}

	plaintext, err := base64.StdEncoding.DecodeString(*decryptOutput.Body.Plaintext)
	if err != nil {
		return fmt.Errorf("failed to decode alicloud kms plaintext for key %s: %w", c.params.KeyID, err)
	}
	if !bytes.Equal(plaintext, testData) {
		return fmt.Errorf("decrypted data does not match original for alicloud key %s", c.params.KeyID)
	}

	return nil
}

func (c *KMSChecker) createClient(ctx context.Context) (*awskms.Client, error) {
	auth := buildAWSAssumeRole(c.params.AWS)
	return util.NewKMSClientWithAssumeRole(ctx, c.params.Region, []util.AssumeRole{
		auth,
	})
}

func buildAWSAssumeRole(auth *core.AWSAuth) util.AssumeRole {
	if auth == nil {
		return util.AssumeRole{}
	}

	return util.AssumeRole{
		RoleARN:    auth.RoleARN,
		ExternalID: auth.ExternalID,
	}
}

func buildAliCloudEncryptionContext(ctx map[string]string) map[string]interface{} {
	if len(ctx) == 0 {
		return nil
	}

	result := make(map[string]interface{}, len(ctx))
	for k, v := range ctx {
		result[k] = v
	}

	return result
}

func createAliCloudKMSClient(params *core.KMSParams) (aliCloudKMSClient, error) {
	credential, err := buildAliCloudCredential(params, "kms-check")
	if err != nil {
		return nil, err
	}

	return alikms.NewClient(new(openapiutil.Config).
		SetCredential(credential).
		SetRegionId(params.Region))
}

func buildAliCloudCredential(params *core.KMSParams, defaultSessionName string) (alicredentials.Credential, error) {
	baseProvider, err := newAliCloudCredential(new(alicredentials.Config).SetType("ecs_ram_role"))
	if err != nil {
		return nil, fmt.Errorf("failed to create ecs role credential: %w", err)
	}

	if params == nil || params.AliCloud == nil || params.AliCloud.RoleARN == "" {
		return baseProvider, nil
	}

	baseCred, err := baseProvider.GetCredential()
	if err != nil {
		return nil, fmt.Errorf("failed to get ecs role credential: %w", err)
	}
	accessKeyID, err := credentialValue(baseCred.AccessKeyId, "access key id")
	if err != nil {
		return nil, err
	}
	accessKeySecret, err := credentialValue(baseCred.AccessKeySecret, "access key secret")
	if err != nil {
		return nil, err
	}
	securityToken, err := credentialValue(baseCred.SecurityToken, "security token")
	if err != nil {
		return nil, err
	}

	sessionName := defaultSessionName
	if params.AliCloud.RoleSessionName != "" {
		sessionName = params.AliCloud.RoleSessionName
	}

	assumedProvider, err := newAliCloudCredential(new(alicredentials.Config).
		SetType("ram_role_arn").
		SetAccessKeyId(accessKeyID).
		SetAccessKeySecret(accessKeySecret).
		SetSecurityToken(securityToken).
		SetRoleArn(params.AliCloud.RoleARN).
		SetExternalId(params.AliCloud.ExternalID).
		SetRoleSessionName(sessionName))
	if err != nil {
		return nil, fmt.Errorf("failed to create ram role credential: %w", err)
	}

	return assumedProvider, nil
}

func credentialValue(value *string, name string) (string, error) {
	if value == nil || *value == "" {
		return "", fmt.Errorf("ecs role credential missing %s", name)
	}
	return *value, nil
}