package ozon

import (
	"fmt"
	"net/http"
	"time"

	"github.com/redis/rueidis"
)

const (
	windowSize = time.Second
	rps        = 50 // requests per second
)

var (
	rateLimiterScript = rueidis.NewLuaScript(`
			local key = KEYS[1]
			local now = tonumber(ARGV[1])
			local window = tonumber(ARGV[2])
			local limit = tonumber(ARGV[3])
			
			-- Удаляем старые записи вне окна времени
			redis.call('ZREMRANGEBYSCORE', key, '-inf', now - window)
			local count = redis.call('ZCARD', key)
			
			if count < limit then
				-- Лимит не превышен, добавляем новую метку и устанавливаем TTL
				redis.call('ZADD', key, now, now)
				redis.call('EXPIRE', key, math.ceil(window / 1000000000))
				return 0 -- Можно выполнять запрос сразу
			else
				-- Лимит превышен, находим самую старую метку
				local oldest = redis.call('ZRANGE', key, 0, 0, 'WITHSCORES')[2]
				-- Возвращаем время, которое нужно подождать до освобождения слота
				return (tonumber(oldest) + window) - now
			end
	`)
)

type RateLimitTransport struct {
	http.RoundTripper
	redis rueidis.Client
}

func (t *RateLimitTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	ctx := req.Context()
	clientId := req.Header.Get("Client-Id")
	now := time.Now().UnixNano()

	waitTime, err := rateLimiterScript.Exec(ctx, t.redis, []string{clientId}, []string{
		fmt.Sprintf("%d", now),
		fmt.Sprintf("%d", int64(windowSize)),
		fmt.Sprintf("%d", rps),
	}).ToInt64()
	if err != nil {
		return nil, fmt.Errorf("failed to execute rate limit script: %w", err)
	}
	if waitTime > 0 {
		fmt.Printf("Rate limit exceeded for client %s, waiting for %d nanoseconds\n", clientId, waitTime)
		time.Sleep(time.Duration(waitTime))
	}
	return t.RoundTripper.RoundTrip(req)
}
func NewRateLimitTransport(redis rueidis.Client) *RateLimitTransport {

	return &RateLimitTransport{RoundTripper: http.DefaultTransport, redis: redis}
}