package jwt

import (
	"context"
	"time"

	"forge.cadoles.com/cadoles/bouncer/internal/jwk"
	"github.com/lestrrat-go/jwx/v2/jwa"
	"github.com/lestrrat-go/jwx/v2/jws"
	"github.com/lestrrat-go/jwx/v2/jwt"
	"github.com/pkg/errors"
)

const keyRole = "role"

func parseToken(ctx context.Context, keys jwk.Set, issuer string, rawToken string, acceptableSkew time.Duration) (jwt.Token, error) {
	token, err := jwt.Parse(
		[]byte(rawToken),
		jwt.WithContext(ctx),
		jwt.WithKeySet(keys, jws.WithRequireKid(false)),
		jwt.WithIssuer(issuer),
		jwt.WithValidate(true),
		jwt.WithAcceptableSkew(acceptableSkew),
	)
	if err != nil {
		return nil, errors.WithStack(err)
	}

	return token, nil
}

func GenerateToken(ctx context.Context, key jwk.Key, issuer, subject string, role Role) (string, error) {
	token := jwt.New()

	if err := token.Set(jwt.SubjectKey, subject); err != nil {
		return "", errors.WithStack(err)
	}

	if err := token.Set(jwt.IssuerKey, issuer); err != nil {
		return "", errors.WithStack(err)
	}

	if err := token.Set(keyRole, role); err != nil {
		return "", errors.WithStack(err)
	}

	now := time.Now().UTC()

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

	if err := token.Set(jwt.IssuedAtKey, now); err != nil {
		return "", errors.WithStack(err)
	}

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

	return string(rawToken), nil
}

func GenerateTokenWithPrivateKey(ctx context.Context, privateKeyFile string, issuer string, subject string, role Role) (string, jwk.Key, error) {
	key, err := jwk.LoadOrGenerate(privateKeyFile, jwk.DefaultKeySize)
	if err != nil {
		return "", nil, errors.WithStack(err)
	}

	token, err := GenerateToken(ctx, key, issuer, subject, role)
	if err != nil {
		return "", nil, errors.WithStack(err)
	}

	return token, key, nil
}