package ratelimit

import (
	"context"
	"fmt"
	"strconv"

	pkg_http "github.com/tidbcloud/account/pkg/grpc/http"
	"github.com/tidbcloud/pkgv2/auth/auth0"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"github.com/ulule/limiter/v3"
	"go.uber.org/zap"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
)

type LimiterWithName struct {
	limiter.Limiter
	Key     string
	Methods []Method
}

func UnaryServerInterceptorV2(limiters []*LimiterWithName) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
		if limiters == nil {
			return handler(ctx, req)
		}

		user := auth0.GetUserIDFromContext(ctx)
		if user == "" {
			return handler(ctx, req)
		}

		method := info.FullMethod
		for _, limiter := range limiters {
			if checkMethod(method, limiter.Methods) {
				s := fmt.Sprintf("%s%s", user, limiter.Key)
				key, err := FnvHash64(s)
				if err != nil {
					tilog.Error(ctx, "failed to hash key of ratelimit", zap.String("hashkey", s), zap.Error(err))
					return handler(ctx, req)
				}
				context, err := limiter.Get(ctx, key)
				if err != nil {
					tilog.Error(ctx, "failed to get ratelimit by key", zap.String("key", limiter.Key), zap.String("user", user))
					return handler(ctx, req)
				}

				m := make(map[string]string)
				rateLimitLimitKey := fmt.Sprintf("%s-%s", RateLimitLimitKey, limiter.Key)
				rateLimitRemainingKey := fmt.Sprintf("%s-%s", RateLimitRemainingKey, limiter.Key)
				m[rateLimitLimitKey] = strconv.FormatInt(context.Limit, 10)
				m[rateLimitRemainingKey] = strconv.FormatInt(context.Remaining, 10)
				err = grpc.SetHeader(ctx, metadata.New(m))
				if err != nil {
					tilog.Error(ctx, "failed to set ratelimit header", zap.Error(err))
				}
				if context.Reached {
					tilog.Info(ctx, "rate limit",
						zap.String("method", pkg_http.GetHttpMethod(ctx)),
						zap.String("url", pkg_http.GetHttpURI(ctx)),
						zap.String("user", user),
						zap.String("X-RateLimit-Key", limiter.Key),
						zap.Int64("X-RateLimit-Limit-Minute", context.Limit),
						zap.Int64("X-RateLimit-Remaining-Minute", context.Remaining),
					)
					return nil, status.Error(codes.ResourceExhausted, "Too many requests, please retry later.")
				}
			}
		}
		return handler(ctx, req)
	}
}