package grpctool

import (
	"context"
	"crypto/ed25519"
	"log/slog"

	"github.com/golang-jwt/jwt/v5"
	grpc_auth "github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/auth"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type JWTAuther struct {
	parser            *jwt.Parser
	loggerFromContext func(context.Context) *slog.Logger
	keyFunc           jwt.Keyfunc
}

// NewHMACJWTAuther creates a new JWT auth for HMAC signatures
// issuer may be empty to disable validation
// audience may be empty to disable validation
func NewHMACJWTAuther(secret []byte, issuer, audience string, loggerFromContext func(context.Context) *slog.Logger) *JWTAuther {
	return &JWTAuther{
		parser: jwt.NewParser(
			jwt.WithAudience(audience),
			jwt.WithIssuer(issuer),
			jwt.WithValidMethods([]string{jwt.SigningMethodHS256.Name, jwt.SigningMethodHS384.Name, jwt.SigningMethodHS512.Name}),
		),
		loggerFromContext: loggerFromContext,
		keyFunc: func(token *jwt.Token) (any, error) {
			return secret, nil
		},
	}
}

// NewEdDSAJWTAuther creates a new JWT auth for EdDSA signatures
// issuer may be empty to disable validation
// audience may be empty to disable validation
func NewEdDSAJWTAuther(publicKey ed25519.PublicKey, issuer, audience string, loggerFromContext func(context.Context) *slog.Logger) *JWTAuther {
	return &JWTAuther{
		parser: jwt.NewParser(
			jwt.WithAudience(audience),
			jwt.WithIssuer(issuer),
			jwt.WithValidMethods([]string{jwt.SigningMethodEdDSA.Alg()}),
		),
		loggerFromContext: loggerFromContext,
		keyFunc: func(token *jwt.Token) (any, error) {
			return publicKey, nil
		},
	}
}

// UnaryServerInterceptor returns a new unary server interceptors that performs per-request JWT auth.
func (a *JWTAuther) UnaryServerInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
	if err := a.doAuth(ctx); err != nil {
		return nil, err
	}
	return handler(ctx, req)
}

// StreamServerInterceptor returns a new stream server interceptors that performs per-request JWT auth.
func (a *JWTAuther) StreamServerInterceptor(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	if err := a.doAuth(stream.Context()); err != nil {
		return err
	}
	return handler(srv, stream)
}

func (a *JWTAuther) doAuth(ctx context.Context) error {
	token, err := grpc_auth.AuthFromMD(ctx, "bearer")
	if err != nil {
		return err // returns gRPC status error
	}
	_, err = a.parser.Parse(token, a.keyFunc)
	if err != nil {
		a.loggerFromContext(ctx).Debug("JWT validation failed", logz.Error(err))
		return status.Error(codes.Unauthenticated, "JWT validation failed")
	}
	return nil
}
