package auth

import (
	"context"
	"errors"
	"fmt"
	"slices"
	"strings"

	"github.com/MicahParks/keyfunc/v3"
	jwt "github.com/golang-jwt/jwt/v5"
)

type TokenValidator interface {
	ValidateToken(ctx context.Context, token string) (jwt.MapClaims, error)
}

type TokenValidatorOptions struct {
	JWKSURL  string
	Audience string
}

type tokenValidator struct {
	keyFunc  keyfunc.Keyfunc
	audience string
}

func NewTokenValidator(opts TokenValidatorOptions) (TokenValidator, error) {
	jwksURL := strings.TrimSpace(opts.JWKSURL)

	keyFunc, err := keyfunc.NewDefaultOverrideCtx(context.Background(), []string{jwksURL}, keyfunc.Override{
		// Now the jwks of tidb.auth0.com has a key with invalid 'x5t', it causes the failure of jwk validation.
		// So skip jwk validation for now, and we can remove this after the jwk is fixed.
		ValidationSkipAll: true,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to create JWKS parser: %w", err)
	}

	return &tokenValidator{
		keyFunc:  keyFunc,
		audience: opts.Audience,
	}, nil
}

func (tv *tokenValidator) ValidateToken(ctx context.Context, tokenString string) (jwt.MapClaims, error) {
	token, err := jwt.Parse(tokenString, tv.keyFunc.KeyfuncCtx(ctx))
	if err != nil {
		return nil, fmt.Errorf("failed to parse token: %w", err)
	}

	if !token.Valid {
		return nil, errors.New("invalid token")
	}

	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		return nil, errors.New("failed to parse claims")
	}

	if tv.audience != "" && !audienceMatches(claims["aud"], tv.audience) {
		return nil, errors.New("audience mismatch")
	}

	return claims, nil
}

func audienceMatches(audClaim any, expected string) bool {
	switch value := audClaim.(type) {
	case string:
		return value == expected
	case []string:
		return slices.Contains(value, expected)
	case []any:
		for _, item := range value {
			if aud, ok := item.(string); ok && aud == expected {
				return true
			}
		}
	}

	return false
}

// GetEmailClaim retrieves the "email" field from claims
func GetEmailClaim(claims jwt.MapClaims, key string) (string, error) {
	emailValue, ok := claims[key]
	if !ok || emailValue == nil {
		return "", errors.New("no email found")
	}

	email, ok := emailValue.(string)
	if !ok {
		return "", fmt.Errorf("email claim is not a string (type %T)", emailValue)
	}

	return email, nil
}