package util

import (
	"sync"
	"time"
)

type cacheEntry[V any] struct {
	value V
	size  int64
	// expiresAt is the zero time when the cache has no TTL.
	expiresAt time.Time
}

// Cache is a bounded, concurrency-safe cache that evicts in insertion order.
//
// ponytail: FIFO eviction, swap for LRU if hit rate matters.
type Cache[K comparable, V any] struct {
	mu       sync.Mutex
	max      int
	ttl      time.Duration
	maxBytes int64
	sizeOf   func(V) int64
	bytes    int64
	items    map[K]cacheEntry[V]
	order    []K
}

// NewCache returns a cache holding at most max entries. A zero ttl means
// entries never expire.
func NewCache[K comparable, V any](max int, ttl time.Duration) *Cache[K, V] {
	return &Cache[K, V]{max: max, ttl: ttl, items: make(map[K]cacheEntry[V], max)}
}

// NewSizedCache is NewCache with a second bound: the sum of sizeOf over all
// entries stays at or below maxBytes. Values larger than maxBytes are not cached.
func NewSizedCache[K comparable, V any](max int, maxBytes int64, ttl time.Duration, sizeOf func(V) int64) *Cache[K, V] {
	c := NewCache[K, V](max, ttl)
	c.maxBytes, c.sizeOf = maxBytes, sizeOf
	return c
}

// Get returns the value for key. Expired entries are dropped and report false.
func (c *Cache[K, V]) Get(key K) (V, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	e, ok := c.items[key]
	if !ok {
		return *new(V), false
	}
	if !e.expiresAt.IsZero() && !time.Now().Before(e.expiresAt) {
		c.remove(key)
		return *new(V), false
	}
	return e.value, true
}

// Set stores a value, evicting the oldest entries when the cache is full.
func (c *Cache[K, V]) Set(key K, v V) { c.SetTTL(key, v, c.ttl) }

// SetTTL is Set with a per-entry ttl. A zero ttl means the entry never expires.
func (c *Cache[K, V]) SetTTL(key K, v V, ttl time.Duration) {
	c.mu.Lock()
	defer c.mu.Unlock()
	e := cacheEntry[V]{value: v}
	if ttl > 0 {
		e.expiresAt = time.Now().Add(ttl)
	}
	c.remove(key)
	if c.sizeOf != nil {
		e.size = c.sizeOf(v)
		if e.size > c.maxBytes {
			return
		}
	}
	for len(c.order) > 0 && (len(c.order) >= c.max || c.bytes+e.size > c.maxBytes && c.sizeOf != nil) {
		c.bytes -= c.items[c.order[0]].size
		delete(c.items, c.order[0])
		c.order = c.order[1:]
	}
	c.items[key] = e
	c.order = append(c.order, key)
	c.bytes += e.size
}

// Delete drops one entry.
func (c *Cache[K, V]) Delete(key K) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.remove(key)
}

// Len reports how many entries are held, expired ones included.
func (c *Cache[K, V]) Len() int {
	c.mu.Lock()
	defer c.mu.Unlock()
	return len(c.items)
}

// remove deletes key from both the map and the order. The caller holds the lock.
func (c *Cache[K, V]) remove(key K) {
	e, ok := c.items[key]
	if !ok {
		return
	}
	c.bytes -= e.size
	delete(c.items, key)
	for i, k := range c.order {
		if k == key {
			c.order = append(c.order[:i], c.order[i+1:]...)
			return
		}
	}
}