package signature

import (
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"net/http"
	"time"

	telemetry "code.justin.tv/amzn/TwitchTelemetry"
	"code.justin.tv/sse/malachai/pkg/internal/keystore"
	"code.justin.tv/sse/malachai/pkg/internal/stats/statsiface"
	"code.justin.tv/sse/malachai/pkg/jwtvalidation"
	"code.justin.tv/sse/malachai/pkg/log"
	"code.justin.tv/sse/malachai/pkg/registration"

	"github.com/SermoDigital/jose/jws"
	"github.com/SermoDigital/jose/jwt"
)

const (
	metricNameRequestValidated                    = "request.validated"
	metricNameRequestValidationAttempted          = "request.validation.attempted"
	metricNameRequestValidationSuccess            = "request.validation.%s.success"
	metricNameRequestValidationFailureExpired     = "request.validation.%s.failure.expired"
	metricNameRequestValidationFailureNBF         = "request.validation.%s.failure.nbf"
	metricNameRequestValidationUnknownFailure     = "request.validation.unknown.failure"
	metricNameRequestValidationFailureNoSignature = "request.validation.unsigned.failure"
)

// Validation errors
var (
	ErrCanonicalHashMismatch  = errors.New("sent canonical hash does not match request")
	ErrMissingSignatureHeader = errors.New("missing signature header")
)

// UnauthorizedRequestError is returned when the is invalid JWT
type UnauthorizedRequestError struct{}

// StatusCode returns the HTTP status code that should be served
func (te *UnauthorizedRequestError) StatusCode() int {
	return http.StatusUnauthorized
}

// TokenNotYetValidError is returned when a token is sent before its nbf time
type TokenNotYetValidError struct {
	UnauthorizedRequestError
	CheckTime time.Time
	NbfTime   time.Time
	NbfLeeway time.Duration
}

func (e *TokenNotYetValidError) Error() string {
	return fmt.Sprintf(
		"token is not yet valid: check time<%d> nbf<%d with %s leeway>",
		e.CheckTime.Unix(), e.NbfTime.Unix(), (e.NbfLeeway / time.Second).String())
}

// TokenExpiredError is returned when a token is sent before its nbf time
type TokenExpiredError struct {
	UnauthorizedRequestError
	CheckTime time.Time
	ExpTime   time.Time
	ExpLeeway time.Duration
}

func (e *TokenExpiredError) Error() string {
	return fmt.Sprintf(
		"token is expired: check time<%d> exp<%d with %s leeway>",
		e.CheckTime.Unix(), e.ExpTime.Unix(), (e.ExpLeeway / time.Second).String())
}

var requiredClaims = []string{
	"exp",
	"jti",
	CanonicalRequestHashClaim,
	CanonicalRequestHeadersClaim,
	CanonicalRequestHashAlgClaim,
}

// RequestValidatorConfig configures the RequestValidator
type RequestValidatorConfig struct {
	ExpLeeway       time.Duration
	NbfLeeway       time.Duration
	PassthroughMode bool
}

// RequestValidator validates requests
type RequestValidator interface {
	ValidateRequest(r *http.Request) (caller *jwtvalidation.SigningEntity, err error)
	ValidateRequestWithHashedBody(r *http.Request, hashedBody []byte) (caller *jwtvalidation.SigningEntity, err error)
}

// ServiceRegistrar implememnts GetServiceByID
type ServiceRegistrar interface {
	GetServiceByID(serviceID string) (svc *registration.Service, err error)
}

// RequestJWTValidator validates signatures
type RequestJWTValidator struct {
	*jwtvalidation.BaseValidator
	cfg       RequestValidatorConfig
	keystore  keystore.Storer
	logger    log.S2SLogger
	statter   statsiface.ReporterAPI
	registrar ServiceRegistrar
}

// NewRequestJWTValidator configures a new request validator
func NewRequestJWTValidator(cfg *RequestValidatorConfig, registrar ServiceRegistrar, pubkeystore keystore.Storer, logger log.S2SLogger, reporter statsiface.ReporterAPI) (validator *RequestJWTValidator) {
	if cfg == nil {
		cfg = new(RequestValidatorConfig)
	}

	return &RequestJWTValidator{
		BaseValidator: &jwtvalidation.BaseValidator{
			Now: func() time.Time {
				return time.Now().UTC()
			},
			RequiredClaims: requiredClaims,
		},
		cfg:      *cfg,
		keystore: pubkeystore,
		logger:   logger,
		statter:  reporter,
		registrar: &serviceCache{
			registrar: registrar,
			cache:     make(map[string]*registration.Service),
		},
	}
}

// ValidateRequestWithHashedBody validates a request signature using a hash
// of the body of the body itself
func (s *RequestJWTValidator) ValidateRequestWithHashedBody(r *http.Request, hashedBody []byte) (caller *jwtvalidation.SigningEntity, err error) {
	sig := []byte(r.Header.Get(SignatureHeader))
	if len(sig) == 0 {
		s.statter.Report(metricNameRequestValidationFailureNoSignature, 1, telemetry.UnitCount)
		err = ErrMissingSignatureHeader
		return
	}
	s.statter.Report(metricNameRequestValidationAttempted, 1, telemetry.UnitCount)
	startTime := time.Now()
	requestJWT, err := jws.ParseJWT(sig)
	if err != nil {
		return
	}

	// retrieve signing method and validate user-provided signing method
	// against whitelist
	signingMethod, err := jwtvalidation.ParseSigningMethod(requestJWT)
	if err != nil {
		return
	}

	kid, err := jwtvalidation.GetKID(requestJWT)
	if err != nil {
		return
	}

	claimedCaller, err := jwtvalidation.DecodeSigningEntityString(kid)
	if err != nil {
		return
	}

	svc, err := s.registrar.GetServiceByID(claimedCaller.Caller)
	if err != nil {
		return
	}

	publicKey, err := s.keystore.GetRSAPublicKey(claimedCaller)
	if err != nil {
		return
	}

	validator := jws.NewValidator(
		jws.Claims{
			CanonicalRequestHashAlgClaim: CanonicalRequestHashAlg,
		},
		s.cfg.ExpLeeway,
		s.cfg.NbfLeeway,
		s.getJWTClaimsValidator(r, hashedBody),
	)

	var metricName string
	err = requestJWT.Validate(publicKey, signingMethod, validator)
	if err != nil {
		if err == jwt.ErrTokenNotYetValid {
			metricName = fmt.Sprintf(metricNameRequestValidationFailureNBF, svc.Name)
			nbf, _ := requestJWT.Claims().NotBefore()
			err = &TokenNotYetValidError{
				CheckTime: time.Now(),
				NbfTime:   nbf,
				NbfLeeway: s.cfg.NbfLeeway,
			}
		} else if err == jwt.ErrTokenIsExpired {
			metricName = fmt.Sprintf(metricNameRequestValidationFailureExpired, svc.Name)
			exp, _ := requestJWT.Claims().Expiration()
			err = &TokenExpiredError{
				CheckTime: time.Now(),
				ExpTime:   exp,
				ExpLeeway: s.cfg.ExpLeeway,
			}
		} else {
			metricName = metricNameRequestValidationUnknownFailure

		}
		s.statter.Report(metricName, 1, telemetry.UnitCount)
		return
	}

	//No Sampling for now.
	s.statter.ReportDurationSample(metricNameRequestValidated, time.Since(startTime))
	s.statter.Report(fmt.Sprintf(metricNameRequestValidationSuccess, svc.Name), 1, telemetry.UnitCount)

	caller = claimedCaller
	return
}

// ValidateRequest returns an error if request signature is not valid
func (s *RequestJWTValidator) ValidateRequest(r *http.Request) (caller *jwtvalidation.SigningEntity, err error) {
	return s.ValidateRequestWithHashedBody(r, nil)
}

func (s *RequestJWTValidator) getJWTClaimsValidator(r *http.Request, hashedBody []byte) func(c jws.Claims) error {
	return func(c jws.Claims) (err error) {
		err = s.ValidateRequiredClaims(c)
		if err != nil {
			return
		}

		err = s.validateHash(r, hashedBody, c)
		return
	}
}

func (s *RequestJWTValidator) validateHash(r *http.Request, hashedBody []byte, c jws.Claims) (err error) {
	claimedCanonicalHashEncoded, ok := c[CanonicalRequestHashClaim].(string)
	if !ok {
		err = jwtvalidation.ClaimWrongType{
			FieldName: CanonicalRequestHashClaim,
			TypeName:  "string",
		}
		return
	}

	claimedCanonicalHash, err := base64.StdEncoding.DecodeString(claimedCanonicalHashEncoded)
	if err != nil {
		return
	}

	var headersList []string
	switch t := c[CanonicalRequestHeadersClaim].(type) {
	default:
		// this catches everything else
		err = jwtvalidation.ClaimWrongType{
			FieldName: CanonicalRequestHeadersClaim,
			TypeName:  "[]string",
		}
		return
	case []interface{}:
		// when len(list) > 0 with strings
		headersList = make([]string, len(t))
		for i, v := range t {
			var ok bool
			headersList[i], ok = v.(string)
			if !ok {
				err = jwtvalidation.ClaimWrongType{
					FieldName: CanonicalRequestHeadersClaim,
					TypeName:  "[]string",
				}
				return
			}
		}
	case []string:
		// when len(list) == 0
		headersList = t
	}

	requestHash, err := generateCanonicalRequestHash(r, hashedBody, headersList)
	if err != nil {
		return
	}

	if !bytes.Equal(requestHash, claimedCanonicalHash) {
		err = ErrCanonicalHashMismatch
	}
	return
}
