import (
    "golang.org/x/time/rate"
    "google.golang.org/grpc"
    "google.golang.org/grpc/status"
    "google.golang.org/grpc/codes"
)

var limiter = rate.NewLimiter(100, 200) // 每秒100个令牌，突发容量200

func RateLimitInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
    if !limiter.Allow() {
        return nil, status.Errorf(codes.ResourceExhausted, "rate limit exceeded")
    }
    return handler(ctx, req)
}

// 启动服务时注册拦截器
server := grpc.NewServer(grpc.UnaryInterceptor(RateLimitInterceptor))

import (
    "net"
    "google.golang.org/grpc/peer"
)

var clients = sync.Map{} // map[string]*rate.Limiter

func getLimiterByIP(ctx context.Context) *rate.Limiter {
    peer, ok := peer.FromContext(ctx)
    if !ok {
        return rate.NewLimiter(10, 20)
    }
    host, _, _ := net.SplitHostPort(peer.Addr.String())
    
    if limiter, exists := clients.Load(host); exists {
        return limiter.(*rate.Limiter)
    }
    
    newLimiter := rate.NewLimiter(50, 100)
    clients.Store(host, newLimiter)
    return newLimiter
}

func redisRateLimit(client *redis.Client, key string, max int, window time.Duration) bool {
    script := `
        local count = redis.call("GET", KEYS[1])
        if not count then
            redis.call("SET", KEYS[1], 1, "EX", ARGV[1])
            return 1
        end
        if tonumber(count) < tonumber(ARGV[2]) then
            redis.call("INCR", KEYS[1])
            return tonumber(count) + 1
        else
            return 0
        end
    `
    result, err := client.Eval(ctx, script, []string{key}, []string{fmt.Sprintf("%d", int(window.Seconds())), fmt.Sprintf("%d", max)}).Result()
    if err != nil || result == 0 {
        return false
    }
    return true
}