// package auth provides auth middlewares.
package auth

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

	"github.com/golang-jwt/jwt/v5"
	"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 newJwtValidator = NewJwtValidator

// Middleware returns the authentication middleware.
func Middleware(handler http.Handler) http.Handler {
	// The original dbaas-proxy uses Auth0ClientSecret to generate a token
	// and then use the token to verify the authorization logic locally for testing
	// https://github.com/tidbcloud/dbaas/blob/a52a50887f1578385d3281b4bc31beb88074abe2/central/cmd/proxy/main.go#L144-L160
	// While cluster-proxy just skips authorization when Auth0ClientSecret is non-empty
	// This is also used for locally testing without having to provide a valid auth0 configuration
	isDev := *Auth0ClientSecret != ""
	var jwtValidator *JwtValidator
	if !isDev {
		var err error
		jwtValidator, err = newJwtValidator(*auth0DomainKey, *auth0AudienceKey)
		if err != nil {
			log.Fatal("Unable to create a JWT validator", 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 := jwtValidator.ValidateRequest(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)
	})
}

// JwtValidator handles JWT validation
type JwtValidator struct {
	auth0Domain   string
	auth0Audience string
	jwks          Jwks
}

// NewJwtValidator creates a new JWT validator
func NewJwtValidator(auth0Domain, auth0Audience string) (*JwtValidator, error) {
	jwks, err := GetJWKS(auth0Domain)
	if err != nil {
		return nil, err
	}

	return &JwtValidator{
		auth0Domain:   auth0Domain,
		auth0Audience: auth0Audience,
		jwks:          jwks,
	}, nil
}

// ValidateRequest validates the JWT token in the request
func (v *JwtValidator) ValidateRequest(r *http.Request) error {
	// Extract token from Authorization header
	// This is the default behavior of the previous `auth0/go-jwt-middleware`
	// ref: https://github.com/auth0/go-jwt-middleware/blob/5493cabe49f7bfa6e2ec444a09d334d90cd4e2bd/jwtmiddleware.go#L41
	authHeader := r.Header.Get("Authorization")
	if authHeader == "" {
		return errors.New("missing authorization header")
	}

	tokenString := strings.TrimPrefix(authHeader, "Bearer ")
	if tokenString == authHeader {
		// Try lowercase "bearer" if "Bearer" didn't match
		tokenString = strings.TrimPrefix(authHeader, "bearer ")
		if tokenString == authHeader {
			return errors.New("invalid authorization header format")
		}
	}

	// Parse and validate the token
	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		// Verify signing method
		if _, ok := token.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, errors.Errorf("unexpected signing method: %v", token.Header["alg"])
		}
		// Verify 'aud' claim
		auds, err := token.Claims.GetAudience()
		if err != nil {
			return nil, errors.New("invalid audience")
		}
		found := false
		for _, aud := range auds {
			if aud == v.auth0Audience {
				found = true
				break
			}
		}
		if !found {
			return nil, errors.New("invalid audience")
		}

		// Verify 'iss' claim
		iss := "https://" + v.auth0Domain + "/"
		issuer, err := token.Claims.GetIssuer()
		if err != nil || issuer != iss {
			return nil, errors.New("invalid issuer")
		}

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

		return jwt.ParseRSAPublicKeyFromPEM([]byte(cert))
	})

	if err != nil {
		return errors.Annotate(err, "token validation failed")
	}
	if !token.Valid {
		return errors.New("invalid token")
	}
	return 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
}