package firewall

import (
	"net/netip"
	"sync"
	"sync/atomic"
	"time"
)

type rateLimitEntry struct {
	count atomic.Int64
	start atomic.Int64 // UnixNano
}

type rateLimiter struct {
	limit   int64
	window  time.Duration
	entries sync.Map // netip.Addr → *rateLimitEntry
	now     func() time.Time
	done    chan struct{}
}

func newRateLimiter(limit int64, window time.Duration) *rateLimiter {
	rl := &rateLimiter{
		limit:  limit,
		window: window,
		now:    time.Now,
		done:   make(chan struct{}),
	}

	go rl.cleanup()
	return rl
}

// Allow checks whether the given address is within its rate limit window.
// Returns false if the address has exceeded the limit.
func (rl *rateLimiter) Allow(addr netip.Addr) bool {
	now := rl.now().UnixNano()

	val, _ := rl.entries.LoadOrStore(addr, &rateLimitEntry{})
	entry := val.(*rateLimitEntry)

	windowStart := entry.start.Load()
	if now-windowStart >= rl.window.Nanoseconds() {
		// Window expired — reset. Intentionally non-atomic across the two
		// stores; worst case a few extra connections slip through at the
		// boundary, which is acceptable for a connection rate limiter.
		entry.start.Store(now)
		entry.count.Store(1)
		return true
	}

	n := entry.count.Add(1)
	return n <= rl.limit
}

// Stop terminates the cleanup goroutine.
func (rl *rateLimiter) Stop() {
	close(rl.done)
}

func (rl *rateLimiter) cleanup() {
	ticker := time.NewTicker(rl.window)
	defer ticker.Stop()

	for {
		select {
		case <-rl.done:
			return
		case <-ticker.C:
			cutoff := rl.now().Add(-2 * rl.window).UnixNano()
			rl.entries.Range(func(key, value any) bool {
				entry := value.(*rateLimitEntry)
				if entry.start.Load() < cutoff {
					rl.entries.Delete(key)
				}
				return true
			})
		}
	}
}