package middleware

import (
	"context"
	"strings"

	"ky-go-kratos/pkg/jwt"

	"github.com/go-kratos/kratos/v2/errors"
	"github.com/go-kratos/kratos/v2/middleware"
	"github.com/go-kratos/kratos/v2/transport"
)

const (
	authorizationKey = "Authorization"
	bearerWord       = "Bearer"
)

type contextKey string

const (
	userIDKey   contextKey = "user_id"
	usernameKey contextKey = "username"
)

func JWTAuth(secretKey string) middleware.Middleware {
	return func(handler middleware.Handler) middleware.Handler {
		return func(ctx context.Context, req interface{}) (interface{}, error) {
			if tr, ok := transport.FromServerContext(ctx); ok {
				auths := strings.Split(tr.RequestHeader().Get(authorizationKey), " ")
				if len(auths) != 2 || !strings.EqualFold(auths[0], bearerWord) {
					return nil, errors.Unauthorized("Unauthorized", "invalid token")
				}
				claims, err := jwt.ParseToken(auths[1], secretKey)
				if err != nil {
					return nil, errors.Unauthorized("Unauthorized", err.Error())
				}
				ctx = context.WithValue(ctx, userIDKey, claims.UserID)
				ctx = context.WithValue(ctx, usernameKey, claims.Username)
			}
			return handler(ctx, req)
		}
	}
}