package impl

import (
	"crypto/rsa"
	"errors"
	"github.com/google/uuid"
	"github.com/lestrrat-go/jwx/v2/jwa"
	"github.com/lestrrat-go/jwx/v2/jwk"
	"github.com/lestrrat-go/jwx/v2/jwt"
	"opensavecloudserver/data/repository/user"
	"opensavecloudserver/server/authentication"
	"time"
)

type (
	JWTAuthenticator struct {
		userRepo user.UserRepository
		key      jwk.Key
		pubKey   jwk.Key
	}

	JWTSession struct {
		id     user.ID
		scopes []authentication.Scope
		roles  []user.Role
	}

	claimKey string
)

func (J JWTSession) UserID() user.ID {
	return J.id
}

func (J JWTSession) Scopes() []authentication.Scope {
	return J.scopes
}

func (J JWTSession) Roles() []user.Role {
	return J.roles
}

const (
	userScopesClaimKey claimKey = "user.scopes"
	userRolesClaimKey  claimKey = "user.roles"

	issuerName string = "oscs"
)

var (
	ErrScopesNotFound = errors.New("scopes not found in JWT")
	ErrRolesNotFound  = errors.New("roles not found in JWT")
)

func (J *JWTAuthenticator) Authenticate(username, password string) ([]byte, error) {
	u, err := J.userRepo.UserByUsername(username)
	if err != nil {
		return nil, err
	}
	if u.CheckPassword(password) {
		// Build a JWT!
		tok, err := jwt.NewBuilder().
			Issuer(issuerName).
			IssuedAt(time.Now()).
			Subject(u.ID().String()).
			Claim(string(userScopesClaimKey), "").
			Claim(string(userRolesClaimKey), u.Roles()).
			Build()
		if err != nil {
			return nil, err
		}

		// Sign a JWT!
		return jwt.Sign(tok, jwt.WithKey(jwa.RS256, J.key))
	}
	return nil, authentication.ErrBadPassword
}

func (J *JWTAuthenticator) Validate(token string) (authentication.Session, error) {
	verifiedToken, err := jwt.Parse([]byte(token), jwt.WithKey(jwa.RS256, J.pubKey), jwt.WithValidate(true), jwt.WithIssuer(issuerName))
	if err != nil {
		return nil, err
	}
	id, err := uuid.Parse(verifiedToken.Subject())
	if err != nil {
		return nil, err
	}
	value, ok := verifiedToken.Get(string(userScopesClaimKey))
	if !ok {
		return nil, ErrScopesNotFound
	}
	scopes, ok := value.([]authentication.Scope)
	if !ok {
		return nil, ErrScopesNotFound
	}
	value, ok = verifiedToken.Get(string(userRolesClaimKey))
	if !ok {
		return nil, ErrRolesNotFound
	}
	roles, ok := value.([]user.Role)
	if !ok {
		return nil, ErrRolesNotFound
	}
	return JWTSession{
		id:     user.ID(id),
		scopes: scopes,
		roles:  roles,
	}, nil
}

func NewJWTAuthenticator(key *rsa.PrivateKey, userRepo user.UserRepository) (authentication.Authenticator, error) {
	a := new(JWTAuthenticator)
	a.userRepo = userRepo
	// Parse, serialize, slice and dice JWKs!
	privkey, err := jwk.FromRaw(key)
	if err != nil {
		return nil, err
	}

	pubkey, err := jwk.PublicKeyOf(privkey)
	if err != nil {
		return nil, err
	}

	a.key = privkey
	a.pubKey = pubkey
	return a, nil
}