package middleware

import (
	"strings"

	"github.com/golang-jwt/jwt"
	"github.com/labstack/echo/v4"
	mw "github.com/labstack/echo/v4/middleware"

	"github.com/pingcap/test-infra/common/model"
	"github.com/pingcap/test-infra/services/tcms/service/common"
	"github.com/pingcap/test-infra/services/tcms/utils"
)

func init() {
	mw.DefaultJWTConfig.Claims = &utils.JWTClaims{}
	mw.DefaultJWTConfig.ContextKey = utils.CONTEXT_KEY_JWT
}

// Auth supports both JWT Token and PAT(Personal Access Token).
func Auth(jwtKey []byte, validatePAT func(ctx utils.Context, pat string) (*model.PersonalToken, error)) echo.MiddlewareFunc {
	validateJWT := mw.JWT(jwtKey)(func(echo.Context) error { return nil })

	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			ctx := c.(utils.Context)

			var token string
			auth := c.Request().Header.Get(echo.HeaderAuthorization)
			l := len("Bearer")
			if len(auth) > l+1 {
				token = auth[l+1:]
			}

			var user *utils.SessionUser
			if strings.HasPrefix(token, common.TOKEN_PREFIX) {
				// 1. Personal Access Token
				info, err := validatePAT(ctx, token)
				if err != nil {
					return err
				}
				user = &utils.SessionUser{ID: info.UserID}
			} else {
				// 2. JWT Token
				if err := validateJWT(c); err != nil {
					return err
				}
				user = &ctx.Get(utils.CONTEXT_KEY_JWT).(*jwt.Token).Claims.(*utils.JWTClaims).User
			}

			ctx.SetUser(user)

			return next(c)
		}
	}
}