package auth

import (
	"context"
	"fmt"
	"time"

	"code.crute.us/mcrute/cloud-identity-broker/app/models"

	"gopkg.in/square/go-jose.v2"
	"gopkg.in/square/go-jose.v2/jwt"
)

const keyIdHeader = jose.HeaderKey("kid")

type JWTManager struct {
	Store        models.UserStore
	Audience     string
	TokenExpires time.Duration
}

// Validate performs validate of a JWT for authentication. If the
// authentication was successful it will return the user model for the
// authenticated user.
//
// JWTs are signed by a per-user key on the service side. These keys may be
// generated automatically when issuing a token to the user or, in the case of
// service accounts, they may be generated by the user and validated on the
// service side with the user's public key.
func (m *JWTManager) Validate(j string) (*models.User, error) {
	ctx := context.Background()

	parsed, err := jwt.ParseSigned(j)
	if err != nil {
		return nil, err
	}

	// With a single signature there should be exactly one header. Without that
	// we can't get the key id which means we won't be able to get the key and
	// validate the signature.
	if len(parsed.Headers) != 1 {
		return nil, fmt.Errorf("Expected exactly 1 JWT header, got %d", len(parsed.Headers))
	}

	kid := parsed.Headers[0].KeyID
	if kid == "" {
		return nil, fmt.Errorf("No key ID in token header")
	}

	// We need the subject claim to lookup the user so we can get their key and
	// validate the token signature.
	untrustedClaims := jwt.Claims{}
	if err := parsed.UnsafeClaimsWithoutVerification(&untrustedClaims); err != nil {
		return nil, err
	}

	user, err := m.Store.Get(ctx, untrustedClaims.Subject)
	if err != nil {
		return nil, err
	}

	key := user.GetKey(kid) // Will not return invalid keys
	if key == nil {
		return nil, fmt.Errorf("No key in user record for key id %s", kid)
	}

	claims := jwt.Claims{}
	if err = parsed.Claims(key.PublicKey, &claims); err != nil {
		return nil, err
	}

	if err = claims.Validate(jwt.Expected{
		Audience: jwt.Audience{m.Audience},
		Time:     time.Now(), // +/- 1 minute
	}); err != nil {
		return nil, err
	}

	// If we made it here then the user matches the public key used to sign the
	// token and the claims are verified so just return the user record we
	// fetched earlier.

	return user, nil
}

// CreateForUser creates a new JWT based on the passed user. In some cases it
// may also need to generate a new SessionKey for the purpose of encrypting
// those JWTs. Callers should be sure to save the returned SessionKey, if there
// is one, to the user when setting the token otherwise the token signature
// will not be able to be validated on subsequent requests.
func (m *JWTManager) CreateForUser(u *models.User) (string, *models.SessionKey, error) {
	pk, err := models.GenerateSessionKey(m.TokenExpires)
	if err != nil {
		return "", nil, err
	}

	// The ExtraHeaders bit is kind of an  ugly hack but there's no way in the
	// official API to set the KeyID for single-signer tokens. However, if the
	// magic "kid" label is set in the extra headers it will be set as the
	// KeyID in the header when that's deserialized.
	signer, err := jose.NewSigner(jose.SigningKey{
		Algorithm: jose.ES256,
		Key:       pk.PrivateKey,
	}, &jose.SignerOptions{
		ExtraHeaders: map[jose.HeaderKey]interface{}{
			keyIdHeader: pk.KeyId,
		},
	})
	if err != nil {
		return "", nil, err
	}

	now := time.Now()
	j, err := jwt.Signed(signer).Claims(jwt.Claims{
		Issuer:   m.Audience,
		Subject:  u.Username,
		Audience: jwt.Audience{m.Audience},
		Expiry:   jwt.NewNumericDate(now.Add(m.TokenExpires)),
		IssuedAt: jwt.NewNumericDate(now),
	}).Claims(map[string]interface{}{
		"admin": u.IsAdmin, // Advisory, for UI, the server must never trust this
	}).CompactSerialize()

	return j, pk, err
}