// Package ratelimit implements a per-IP sliding-window failure counter used to
// hide the database behind a cap on failed lookups.
package ratelimit

import (
	"sync"
	"time"
)

type bucket struct {
	index int64
	count int
}

type ipState struct {
	current  bucket
	previous bucket
}

// Limiter is a two-bucket sliding-window counter per IP. Allowed is a pure read
// (RLock) that estimates the failure count in the trailing window from whatever
// is already stored, without mutating it. RecordFailure takes the write lock and
// performs the actual window roll and increment. Keeping the previous window's
// count means a client blocked near the end of one window stays throttled into
// the next, instead of getting an instant reset at the window boundary.
type Limiter struct {
	mu     sync.RWMutex
	state  map[string]ipState
	window time.Duration
	max    int
}

func NewLimiter(max int, window time.Duration) *Limiter {
	return &Limiter{
		state:  make(map[string]ipState),
		window: window,
		max:    max,
	}
}

func (l *Limiter) windowIndex(t time.Time) int64 {
	return t.UnixNano() / int64(l.window)
}

// Allowed reports whether ip is currently under the failure threshold. It never
// mutates the limiter's state.
func (l *Limiter) Allowed(ip string) bool {
	now := time.Now()
	idx := l.windowIndex(now)

	l.mu.RLock()
	state, ok := l.state[ip]
	l.mu.RUnlock()
	if !ok {
		return true
	}

	current, previous := effectiveBuckets(state, idx)

	overlap := overlapFraction(now, l.window, previous.index)
	estimate := float64(current.count) + float64(previous.count)*overlap

	return estimate < float64(l.max)
}

// RecordFailure increments the failure count for ip in the current window,
// rolling the previous stored window forward if it has become stale.
func (l *Limiter) RecordFailure(ip string) {
	now := time.Now()
	idx := l.windowIndex(now)

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

	state := l.state[ip]
	current, previous := effectiveBuckets(state, idx)
	current.count++
	l.state[ip] = ipState{current: current, previous: previous}
}

// Cleanup removes IPs whose state is more than 2 windows stale, bounding memory.
func (l *Limiter) Cleanup() {
	idx := l.windowIndex(time.Now())

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

	for ip, state := range l.state {
		if idx-state.current.index > 2 {
			delete(l.state, ip)
		}
	}
}

// Run periodically calls Cleanup until ctx is done.
func (l *Limiter) Run(done <-chan struct{}, interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-done:
			return
		case <-ticker.C:
			l.Cleanup()
		}
	}
}

// effectiveBuckets derives the current/previous buckets as seen from window idx,
// without mutating the stored state: if the stored current bucket is stale by
// exactly one window it becomes the effective previous bucket; if it is older
// than that, both buckets are treated as empty.
func effectiveBuckets(state ipState, idx int64) (current, previous bucket) {
	switch {
	case state.current.index == idx:
		current = state.current
		if state.previous.index == idx-1 {
			previous = state.previous
		} else {
			previous = bucket{index: idx - 1}
		}
	case state.current.index == idx-1:
		current = bucket{index: idx}
		previous = state.current
	default:
		current = bucket{index: idx}
		previous = bucket{index: idx - 1}
	}
	return current, previous
}

// overlapFraction returns how much of the previous window (identified by
// previousIdx) still falls within the trailing lookback window of length
// windowSize measured back from now.
func overlapFraction(now time.Time, windowSize time.Duration, previousIdx int64) float64 {
	previousWindowEnd := time.Unix(0, (previousIdx+1)*int64(windowSize))
	elapsedSincePreviousEnd := now.Sub(previousWindowEnd)
	if elapsedSincePreviousEnd >= windowSize {
		return 0
	}
	if elapsedSincePreviousEnd <= 0 {
		return 1
	}
	return 1 - float64(elapsedSincePreviousEnd)/float64(windowSize)
}