// Package ratelimit provides small process-local rate-limit primitives used by
// the gateway edge policy layers.
package ratelimit

import (
	"sync"
	"time"

	"golang.org/x/time/rate"
)

// Policy describes one token-bucket budget enforced for a concrete key.
type Policy struct {
	// Requests is the number of accepted requests replenished per Window.
	Requests int

	// Window is the interval over which Requests are replenished.
	Window time.Duration

	// Burst is the maximum number of immediately available tokens.
	Burst int
}

// Decision describes the result of one limiter reservation attempt.
type Decision struct {
	// Allowed reports whether the request may proceed immediately.
	Allowed bool

	// RetryAfter is the minimum delay the caller should wait before retrying
	// when Allowed is false.
	RetryAfter time.Duration
}

// Limiter applies a policy to one concrete key.
type Limiter interface {
	// Reserve evaluates key under policy and reports whether the request may
	// proceed immediately.
	Reserve(key string, policy Policy) Decision
}

// InMemory is a process-local Limiter backed by x/time/rate token buckets.
type InMemory struct {
	now             func() time.Time
	cleanupInterval time.Duration

	mu          sync.Mutex
	entries     map[string]*entry
	nextCleanup time.Time
}

type entry struct {
	limiter   *rate.Limiter
	limit     rate.Limit
	burst     int
	expiresAt time.Time
}

// NewInMemory constructs a process-local limiter suitable for one gateway
// process instance.
func NewInMemory() *InMemory {
	return &InMemory{
		now:             time.Now,
		cleanupInterval: time.Minute,
		entries:         make(map[string]*entry),
	}
}

// Reserve evaluates key against policy and reports whether the request may
// proceed immediately.
func (l *InMemory) Reserve(key string, policy Policy) Decision {
	if policy.Requests <= 0 || policy.Window <= 0 || policy.Burst <= 0 {
		return Decision{}
	}

	now := l.now()
	limit := rate.Limit(float64(policy.Requests) / policy.Window.Seconds())

	l.mu.Lock()
	defer l.mu.Unlock()

	l.cleanupExpiredBucketsLocked(now)

	current, ok := l.entries[key]
	if !ok || current.limit != limit || current.burst != policy.Burst {
		current = &entry{
			limiter: rate.NewLimiter(limit, policy.Burst),
			limit:   limit,
			burst:   policy.Burst,
		}
		l.entries[key] = current
	}

	current.expiresAt = now.Add(entryTTL(policy.Window))

	reservation := current.limiter.ReserveN(now, 1)
	if !reservation.OK() {
		return Decision{
			Allowed:    false,
			RetryAfter: policy.Window,
		}
	}

	retryAfter := reservation.DelayFrom(now)
	if retryAfter > 0 {
		return Decision{
			Allowed:    false,
			RetryAfter: retryAfter,
		}
	}

	return Decision{Allowed: true}
}

func (l *InMemory) cleanupExpiredBucketsLocked(now time.Time) {
	if !l.nextCleanup.IsZero() && now.Before(l.nextCleanup) {
		return
	}

	for key, current := range l.entries {
		if !current.expiresAt.After(now) {
			delete(l.entries, key)
		}
	}

	l.nextCleanup = now.Add(l.cleanupInterval)
}

func entryTTL(window time.Duration) time.Duration {
	if window < time.Minute {
		return time.Minute
	}

	return 2 * window
}

var _ Limiter = (*InMemory)(nil)