package main

import (
	"log"
	"net/http"
	"sync/atomic"
	"time"

	ratelimit "github.com/JGLTechnologies/gin-rate-limit"
	"github.com/gin-gonic/gin"
)

func mgRateLimit() gin.HandlerFunc {
	var totalLimited uint32
	store := ratelimit.InMemoryStore(&ratelimit.InMemoryOptions{
		Rate:  time.Second,
		Limit: 50, // TODO: This must be changed to the 1000rps.
	})
	return ratelimit.RateLimiter(store, &ratelimit.Options{
		ErrorHandler: func(c *gin.Context, info ratelimit.Info) {
			cur := atomic.AddUint32(&totalLimited, 1)
			log.Printf("%s: [%d] rate limited > %t, will reset at %s\n", c.ClientIP(), cur, info.RateLimited, info.ResetTime)
			c.AbortWithStatus(http.StatusTooManyRequests)
		},
		KeyFunc: keyFunc,
	})
}

func vkRateLimit() gin.HandlerFunc {
	var totalLimited uint32
	store := ratelimit.InMemoryStore(&ratelimit.InMemoryOptions{
		Rate:  time.Second,
		Limit: 20,
	})
	return ratelimit.RateLimiter(store, &ratelimit.Options{
		ErrorHandler: func(c *gin.Context, info ratelimit.Info) {
			cur := atomic.AddUint32(&totalLimited, 1)
			log.Printf("%s: [%d] VK request rate limited > %t, will reset at %s\n", c.ClientIP(), cur, info.RateLimited, info.ResetTime)
			c.AbortWithStatus(http.StatusTooManyRequests)
		},
		KeyFunc: keyFunc,
	})
}

func keyFunc(c *gin.Context) string {
	return c.ClientIP()
}