package jwtutil

import (
	"strings"

	"github.com/lestrrat-go/jwx/v2/jwa"
	"github.com/lestrrat-go/jwx/v2/jwk"
	"github.com/pkg/errors"
)

func AddKeyWithSigningAlgo(keySet jwk.Set, key jwk.Key, signingAlgorithm jwa.SignatureAlgorithm) error {
	addedKey := key

	if !strings.HasPrefix(string(signingAlgorithm), "HS") {
		publicKey, err := key.PublicKey()
		if err != nil {
			return errors.WithStack(err)
		}

		addedKey = publicKey
	}

	if err := addedKey.Set(jwk.AlgorithmKey, signingAlgorithm); err != nil {
		return errors.WithStack(err)
	}

	if err := keySet.AddKey(addedKey); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func NewKeySet(keys ...jwk.Key) (jwk.Set, error) {
	set := jwk.NewSet()

	for _, k := range keys {
		if err := set.AddKey(k); err != nil {
			return nil, errors.WithStack(err)
		}
	}

	return set, nil
}

func NewSymmetricKey(secret []byte) (jwk.Key, error) {
	key, err := jwk.FromRaw(secret)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	if err := key.Set(jwk.AlgorithmKey, jwa.HS256); err != nil {
		return nil, errors.WithStack(err)
	}

	return key, nil
}

func NewSymmetricKeySet(secrets ...[]byte) (jwk.Set, error) {
	keys := make([]jwk.Key, len(secrets))

	for idx, sec := range secrets {
		key, err := NewSymmetricKey(sec)
		if err != nil {
			return nil, errors.WithStack(err)
		}

		keys[idx] = key
	}

	keySet, err := NewKeySet(keys...)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return keySet, nil
}