package jwtutil

import (
	"time"

	"github.com/lestrrat-go/jwx/v2/jwa"
	"github.com/lestrrat-go/jwx/v2/jwk"
	"github.com/lestrrat-go/jwx/v2/jws"
	"github.com/lestrrat-go/jwx/v2/jwt"
	"github.com/oklog/ulid/v2"
	"github.com/pkg/errors"
)

func SignedToken(key jwk.Key, signingAlgorithm jwa.SignatureAlgorithm, claims map[string]any) ([]byte, error) {
	token := jwt.New()

	if err := token.Set(jwt.NotBeforeKey, time.Now()); err != nil {
		return nil, errors.WithStack(err)
	}

	if err := token.Set(jwt.JwtIDKey, ulid.Make().String()); err != nil {
		return nil, errors.WithStack(err)
	}

	for key, value := range claims {
		if err := token.Set(key, value); err != nil {
			return nil, errors.Wrapf(err, "could not set claim '%s' with value '%v'", key, value)
		}
	}

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

	rawToken, err := jwt.Sign(token, jwt.WithKey(signingAlgorithm, key))
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return rawToken, nil
}

func Parse(rawToken []byte, keySet jwk.Set) (jwt.Token, error) {
	token, err := jwt.Parse(rawToken,
		jwt.WithKeySet(keySet, jws.WithRequireKid(false)),
		jwt.WithValidate(true),
	)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return token, nil
}