package throttle

import (
	"context"
	"sync"
	"time"
)

type bucket struct {
	tokens     float64
	lastRefill time.Time
}

// MemoryBucket is an in-memory TokenBucket implementation.
// Suitable for single-process deployments.
type MemoryBucket struct {
	mu      sync.Mutex
	buckets map[string]*bucket
}

func NewMemoryBucket() *MemoryBucket {
	return &MemoryBucket{
		buckets: make(map[string]*bucket),
	}
}

func (m *MemoryBucket) Take(_ context.Context, key string, capacity int, refillRate float64) (bool, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	now := time.Now()
	b, ok := m.buckets[key]
	if !ok {
		// First request: create bucket with capacity-1 tokens (this request uses one).
		m.buckets[key] = &bucket{
			tokens:     float64(capacity) - 1,
			lastRefill: now,
		}
		return true, nil
	}

	// Refill tokens based on elapsed time.
	elapsed := now.Sub(b.lastRefill).Seconds()
	b.tokens += elapsed * refillRate
	if b.tokens > float64(capacity) {
		b.tokens = float64(capacity)
	}
	b.lastRefill = now

	if b.tokens >= 1 {
		b.tokens--
		return true, nil
	}
	return false, nil
}

func (m *MemoryBucket) Reset(_ context.Context, key string) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.buckets, key)
	return nil
}

// Cleanup removes buckets that haven't been used within the given duration.
// Call periodically to prevent unbounded memory growth from stale entries.
func (m *MemoryBucket) Cleanup(stale time.Duration) {
	m.mu.Lock()
	defer m.mu.Unlock()
	now := time.Now()
	for key, b := range m.buckets {
		if now.Sub(b.lastRefill) > stale {
			delete(m.buckets, key)
		}
	}
}