// package auth provides auth middlewares.
package auth

import (
	"encoding/json"
	"flag"
	"net/http"
	"os"

	jwtmiddleware "github.com/auth0/go-jwt-middleware"
	"github.com/dgrijalva/jwt-go"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"go.uber.org/zap"
)

var (
	Auth0ClientSecret = flag.String("auth0-client-secret", os.Getenv("ProxyAuth0ClientSecret"), "generate auth0 token used for develop in local")
	auth0DomainKey    = flag.String("auth0-domain", os.Getenv("AUTH0-DOMAIN"), "The domain of the Auth0 JWT.")
	auth0AudienceKey  = flag.String("auth0-audience", os.Getenv("AUTH0-AUDIENCE"), "The audience of the Auth0 JWT")
)

var newJwtMiddleware = NewJwtMiddleware

// Middleware returns the authentication middleware.
func Middleware(handler http.Handler) http.Handler {
	isDev := *Auth0ClientSecret != ""
	var jwtMiddleware *jwtmiddleware.JWTMiddleware
	if !isDev {
		var err error
		jwtMiddleware, err = newJwtMiddleware(*auth0DomainKey, *auth0AudienceKey)
		if err != nil {
			log.Fatal("Unable to create a JWT middleware", zap.Error(err))
		}
	}

	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if isDev {
			log.Debug("Dev mode is on. Skip authorization.")
		} else {
			if err := jwtMiddleware.CheckJWT(w, r); err != nil {
				log.Error("invalid jwt", zap.Stringer("URL", r.URL), zap.Error(err))
				http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
				return
			}
		}
		handler.ServeHTTP(w, r)
	})
}

// NewJwtMiddleware return a handler to check the client auth
func NewJwtMiddleware(auth0Domain, auth0Audience string) (*jwtmiddleware.JWTMiddleware, error) {
	jwks, err := GetJWKS(auth0Domain)
	if err != nil {
		return nil, err
	}

	jwtMiddleware := jwtmiddleware.New(jwtmiddleware.Options{
		ValidationKeyGetter: func(token *jwt.Token) (interface{}, error) {
			// Verify 'aud' claim
			checkAud := token.Claims.(jwt.MapClaims).VerifyAudience(auth0Audience, false)
			if !checkAud {
				return token, errors.New("invalid audience")
			}
			// Verify 'iss' claim
			iss := "https://" + auth0Domain + "/"
			checkIss := token.Claims.(jwt.MapClaims).VerifyIssuer(iss, false)
			if !checkIss {
				return token, errors.New("invalid issuer")
			}

			cert, err := MakePemCert(token, jwks)
			if err != nil {
				panic(err.Error())
			}

			return jwt.ParseRSAPublicKeyFromPEM([]byte(cert))
		},
		SigningMethod: jwt.SigningMethodRS256,
	})
	return jwtMiddleware, nil
}

// GetJWKS retrieves the public key from Auth0
// You should call this function at startup
func GetJWKS(domain string) (Jwks, error) {
	var jwks = Jwks{}
	var err error
	resp, err := http.Get("https://" + domain + "/.well-known/jwks.json")
	if err != nil {
		return jwks, errors.Annotate(err, "Get jwks.json")
	}
	defer resp.Body.Close()

	err = json.NewDecoder(resp.Body).Decode(&jwks)
	if err != nil {
		return jwks, errors.Annotate(err, "json decode")
	}

	return jwks, nil
}

// Jwks is the return from Auth0 api call for keys
type Jwks struct {
	Keys []JSONWebKeys `json:"keys"`
}

// JSONWebKeys is the return from Auth0 api call for keys
type JSONWebKeys struct {
	Kty string   `json:"kty"`
	Kid string   `json:"kid"`
	Use string   `json:"use"`
	N   string   `json:"n"`
	E   string   `json:"e"`
	X5c []string `json:"x5c"`
}

// MakePemCert makes a cert
func MakePemCert(token *jwt.Token, jwks Jwks) (string, error) {
	cert := ""

	for k := range jwks.Keys {
		if token.Header["kid"] == jwks.Keys[k].Kid {
			cert = "-----BEGIN CERTIFICATE-----\n" + jwks.Keys[k].X5c[0] + "\n-----END CERTIFICATE-----"
		}
	}

	if cert == "" {
		err := errors.New("unable to find appropriate key")
		return cert, err
	}

	return cert, nil
}