package cooler

import (
	"context"
	"fmt"
	"log/slog"
	"time"

	"github.com/redis/go-redis/v9"
)

// attemptScript is one atomic fixed-window admission with the semantics of
// surf.MemoryStore: a missing or expired key opens a window of decay
// milliseconds, a count at max denies without incrementing, otherwise the
// count is incremented. INCR keeps the TTL, so the window never extends.
// It returns {allowed (0|1), count, remaining window in ms}.
//
// KEYS[1] = <prefix>rl:<key>; ARGV[1] = max; ARGV[2] = decay in ms (>= 1).
var attemptScript = redis.NewScript(`
local max = tonumber(ARGV[1])
local decay = tonumber(ARGV[2])
local count = tonumber(redis.call('GET', KEYS[1]) or '0')
local ttl = redis.call('PTTL', KEYS[1])
if ttl < 0 then
  count = 0
  redis.call('SET', KEYS[1], 0, 'PX', decay)
  ttl = decay
end
if count >= max then
  return {0, count, ttl}
end
count = redis.call('INCR', KEYS[1])
return {1, count, ttl}
`)

// RateLimitStore is a fixed-window rate-limit store on Redis, shared by every
// application instance that points at the same Redis. Its Attempt has the
// exact signature of surf.Store, so surf uses it without cooler importing
// surf. When Redis does not answer, Attempt counts in the fallback store
// (per instance) and logs a warning at most once per minute.
type RateLimitStore struct {
	client   *redis.Client
	prefix   string
	fallback Attempter
	warn     *warner
}

// NewRateLimitStore returns a store that keeps its counters under
// <prefix>rl:<key> on client. fallback answers while Redis fails; surf passes
// a surf.MemoryStore. A nil fallback admits every request during an outage.
func NewRateLimitStore(client *redis.Client, prefix string, fallback Attempter) *RateLimitStore {
	return &RateLimitStore{client: client, prefix: prefix, fallback: fallback, warn: newWarner(slog.Default())}
}

// Attempt admits and counts one request when key is below max, in one
// atomic Redis script bounded by AttemptTimeout. The first attempt opens the
// window, a denied attempt does not increment, and retryAfter is the time
// left in the window.
func (s *RateLimitStore) Attempt(key string, max int, decay time.Duration) (allowed bool, attempts int, retryAfter time.Duration) {
	ms := decay.Milliseconds()
	if ms < 1 {
		ms = 1
	}
	ctx, cancel := context.WithTimeout(context.Background(), AttemptTimeout)
	defer cancel()
	res, err := s.run(ctx, s.prefix+"rl:"+key, max, ms)
	if err != nil {
		s.warn.warn("cooler: redis rate-limit store failed; counting per instance until it answers",
			"endpoint", redisEndpoint(s.client), "error", err)
		if s.fallback == nil {
			return true, 0, 0
		}
		return s.fallback.Attempt(key, max, decay)
	}
	ttl := time.Duration(res[2]) * time.Millisecond
	if ttl < 0 {
		ttl = 0
	}
	return res[0] == 1, int(res[1]), ttl
}

func (s *RateLimitStore) run(ctx context.Context, key string, max int, ms int64) ([]int64, error) {
	if s == nil || s.client == nil {
		return nil, fmt.Errorf("cooler: rate-limit store has no Redis client")
	}
	res, err := attemptScript.Run(ctx, s.client, []string{key}, max, ms).Int64Slice()
	if err != nil {
		return nil, err
	}
	if len(res) != 3 {
		return nil, fmt.Errorf("cooler: rate-limit script returned %d values, want 3", len(res))
	}
	return res, nil
}