package utils

import (
	"context"
	"crypto"
	"encoding/json"
	"encoding/pem"
	"sync"

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

	"github.com/tidbcloud/account/commons/config"
	"github.com/tidbcloud/pkgv2/log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/kms"
	"github.com/lestrrat-go/jwx/v2/jwa"
	"github.com/lestrrat-go/jwx/v2/jwk"

	"go.uber.org/zap"
)

type KeyBean struct {
	KeyID     string
	KeyName   string
	Desc      string
	PublicKey []byte
}

type KeyService struct {
	kms *kms.KMS
}

var keyService *KeyService
var keyOnce sync.Once

func GetKeyServiceSingleton() *KeyService {
	var err error
	keyOnce.Do(func() {
		keyService, err = getKeyService()
		if err != nil {
			tilog.Info(context.Background(), "init keyService obj err", zap.Any("err", err.Error()))
		}
		tilog.Info(context.Background(), "init keyService obj success")

	})
	return keyService
}

func Hash(hash crypto.Hash, message []byte) []byte {
	h := hash.New()
	h.Write(message)
	return h.Sum(nil)
}

/*
func InitKms(region string) (*kms.KMS, error) {
	awsConfig := aws.NewConfig().WithMaxRetries(3).WithRegion(region)
	opts := session.Options{
		Config: *awsConfig,
	}

	sess, err := session.NewSessionWithOptions(opts)
	if err != nil {
		return nil, err
	}

	k := kms.New(sess)

	return k, nil
}
*/

func New(region string) (*KeyService, error) {
	awsConfig := aws.NewConfig().WithMaxRetries(3).WithRegion(region)
	opts := session.Options{
		Config: *awsConfig,
	}

	sess, err := session.NewSessionWithOptions(opts)
	if err != nil {
		return nil, err
	}

	k := kms.New(sess)
	svc := &KeyService{
		kms: k,
	}

	return svc, nil
}

// KMS returns raw *kms.KMS instance for direct API calls
func (s *KeyService) KMS() *kms.KMS {
	return s.kms
}

func (s *KeyService) CreateSignKey(description string) (string, string, error) {
	input := &kms.CreateKeyInput{
		KeyUsage: aws.String(kms.KeyUsageTypeSignVerify),
		KeySpec:  aws.String(kms.KeySpecRsa4096),
	}

	if len(description) > 0 {
		input.SetDescription(description)
	}

	output, err := s.kms.CreateKey(input)
	if err != nil {
		return "", "", err
	}

	keyID := output.KeyMetadata.KeyId
	arn := output.KeyMetadata.Arn

	ctx := context.Background()
	tilog.Info(ctx, "createKeyResult", zap.Any("keyID", keyID), zap.Any("arn", arn))
	return *keyID, *arn, err
}

func (s *KeyService) GetSignKeyById(keyID string) (*kms.KeyMetadata, error) {
	input := &kms.DescribeKeyInput{
		KeyId: aws.String(keyID),
	}

	output, err := s.kms.DescribeKey(input)
	if err != nil {
		return nil, err
	}

	return output.KeyMetadata, nil
}

func (s *KeyService) DeleteSignKeyById(keyID string) error {
	input := &kms.ScheduleKeyDeletionInput{
		KeyId: aws.String(keyID),
	}

	_, err := s.kms.ScheduleKeyDeletion(input)
	if err != nil {
		return err
	}

	return nil
}

func (s *KeyService) SignRaw(keyID string, message []byte) ([]byte, error) {
	input := &kms.SignInput{
		KeyId:            aws.String(keyID),
		Message:          message,
		MessageType:      aws.String(kms.MessageTypeRaw),
		SigningAlgorithm: aws.String(kms.SigningAlgorithmSpecRsassaPkcs1V15Sha256),
	}

	output, err := s.kms.Sign(input)
	if err != nil {
		return nil, err
	}

	return output.Signature, nil
}

func (s *KeyService) ExportPublicKey(keyID string) ([]byte, error) {
	input := &kms.GetPublicKeyInput{
		KeyId: aws.String(keyID),
	}

	output, err := s.kms.GetPublicKey(input)
	if err != nil {
		return nil, err
	}

	return output.PublicKey, nil
}

func (s *KeyService) ListKeys() ([]*kms.KeyListEntry, error) {
	input := &kms.ListKeysInput{}
	output, err := s.kms.ListKeys(input)
	if err != nil {
		return nil, err
	}

	return output.Keys, nil
}

func (s *KeyService) FindKey(name string) (string, error) {
	input := &kms.ListAliasesInput{
		Limit: aws.Int64(100),
	}

	output, err := s.kms.ListAliases(input)
	if err != nil {
		return "", err
	}

	for _, alias := range output.Aliases {
		if *alias.AliasName == name {
			return *alias.TargetKeyId, nil
		}
	}

	return "", nil
}

func (s *KeyService) SignDigest(keyID string, digest []byte) ([]byte, error) {
	input := &kms.SignInput{
		KeyId:            aws.String(keyID),
		Message:          digest,
		MessageType:      aws.String(kms.MessageTypeDigest),
		SigningAlgorithm: aws.String(kms.SigningAlgorithmSpecRsassaPkcs1V15Sha256),
	}

	output, err := s.kms.Sign(input)
	if err != nil {
		return nil, err
	}

	return output.Signature, nil
}

func (s *KeyService) VerifyRaw(keyID string, message []byte, signature []byte) (bool, error) {
	input := &kms.VerifyInput{
		KeyId:            aws.String(keyID),
		Message:          message,
		Signature:        signature,
		MessageType:      aws.String(kms.MessageTypeRaw),
		SigningAlgorithm: aws.String(kms.SigningAlgorithmSpecRsassaPkcs1V15Sha256),
	}

	output, err := s.kms.Verify(input)
	if err != nil {
		return false, err
	}

	return *output.SignatureValid, nil
}

func (s *KeyService) VerifyDigest(keyID string, digest []byte, signature []byte) (bool, error) {
	input := &kms.VerifyInput{
		KeyId:            aws.String(keyID),
		Message:          digest,
		Signature:        signature,
		MessageType:      aws.String(kms.MessageTypeDigest),
		SigningAlgorithm: aws.String(kms.SigningAlgorithmSpecRsassaPkcs1V15Sha256),
	}

	output, err := s.kms.Verify(input)
	if err != nil {
		return false, err
	}

	return *output.SignatureValid, nil
}

//	func getLogger() *zap.Logger {
//		ctx := context.Background()
//		logger := log.WithContextL(ctx)
//		return logger
//	}
func PublicKeyToPEM(der []byte) []byte {
	block := &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: der,
	}

	return pem.EncodeToMemory(block)
}

func DerToJWK(der []byte) (jwk.Key, error) {
	pem := PublicKeyToPEM(der)
	key, err := jwk.ParseKey(pem, jwk.WithPEM(true))
	if err != nil {
		return nil, err
	}

	return key, nil
}

func (s *KeyService) PublicKeyToJWK(keyArr []KeyBean) ([]byte, error) {
	jwks := jwk.NewSet()
	for _, keyBean := range keyArr {
		key, err := DerToJWK(keyBean.PublicKey)
		if err != nil {
			return nil, err
		}

		_ = key.Set(jwk.KeyIDKey, keyBean.KeyID)
		_ = key.Set(jwk.KeyUsageKey, jwk.ForSignature)
		_ = key.Set(jwk.AlgorithmKey, jwa.RS256)
		_ = jwks.AddKey(key)
	}

	data, _ := json.Marshal(jwks)
	return data, nil
}

func getKeyService() (*KeyService, error) {
	region := config.Global().ServerConfig.SecretStore.Region
	ctx := context.Background()
	log.Info(ctx, "will init KeyService in region="+region)

	//region := "us-west-2"
	keyService, err := New(region)
	if err != nil {
		return nil, err
	}

	return keyService, nil
}

func GetPublicKeyByKeyId(keyIds []string) string {
	keyArr := make([]KeyBean, len(keyIds))
	keyService := GetKeyServiceSingleton()
	for i, k := range keyIds {

		byteArr, _ := keyService.ExportPublicKey(k)
		keyBean := &KeyBean{
			KeyID:     k,
			PublicKey: byteArr,
		}

		keyArr[i] = *keyBean
	}
	//arry, _ := keyService.PublicKeyToJWK2(byteArr, keyId)
	arry, _ := keyService.PublicKeyToJWK(keyArr)
	keyStr := string(arry)
	return keyStr

}