package freelru

import (
	"errors"
	"fmt"
	"math/bits"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

// ShardedLRU is a thread-safe, sharded, growable LRU cache.
// Sharding is used to reduce lock contention on high concurrency.
// The downside is that exact LRU behavior is not given (as for the LRU and SynchedLRU types).
type ShardedLRU[K comparable, V comparable] struct {
	lrus   []LRU[K, V]
	mus    []sync.RWMutex
	hash   HashKeyCallback[K]
	shards uint32
	mask   uint32
	sweep  atomic.Uint32
}

// SetLifetime sets the default lifetime of LRU elements.
// Lifetime 0 means "forever".
func (lru *ShardedLRU[K, V]) SetLifetime(lifetime time.Duration) {
	for shard := range lru.lrus {
		lru.mus[shard].Lock()
		lru.lrus[shard].SetLifetime(lifetime)
		lru.mus[shard].Unlock()
	}
}

// SetOnEvict sets the OnEvict callback function.
// The onEvict function is called for each evicted lru entry.
func (lru *ShardedLRU[K, V]) SetOnEvict(onEvict OnEvictCallback[K, V]) {
	for shard := range lru.lrus {
		lru.mus[shard].Lock()
		lru.lrus[shard].SetOnEvict(onEvict)
		lru.mus[shard].Unlock()
	}
}

func (lru *ShardedLRU[K, V]) SetHealthCheck(healthCheck HealthCheckCallback[K, V]) {
	for shard := range lru.lrus {
		lru.mus[shard].Lock()
		lru.lrus[shard].SetHealthCheck(healthCheck)
		lru.mus[shard].Unlock()
	}
}

func nextPowerOfTwo(val uint32) uint32 {
	if bits.OnesCount32(val) != 1 {
		return 1 << bits.Len32(val)
	}
	return val
}

func newSharded[K comparable, V comparable](capacity uint32, hash HashKeyCallback[K]) (*ShardedLRU[K, V],
	error,
) {
	size := uint32(float64(capacity) * 1.25) // 25% extra space for fewer collisions

	return newShardedWithSize[K, V](uint32(runtime.GOMAXPROCS(0)*16), capacity, size, hash)
}

func newShardedWithSize[K comparable, V comparable](shards, capacity, size uint32,
	hash HashKeyCallback[K]) (
	*ShardedLRU[K, V], error,
) {
	if capacity == 0 {
		return nil, errors.New("capacity must be positive")
	}
	if size < capacity {
		return nil, fmt.Errorf("size (%d) is smaller than capacity (%d)", size, capacity)
	}

	if size < 1<<31 {
		size = nextPowerOfTwo(size) // next power of 2 so the LRUs can avoid costly divisions
	} else {
		size = 1 << 31 // the highest 2^N value that fits in a uint32
	}

	shards = nextPowerOfTwo(shards) // next power of 2 so we can avoid costly division for sharding

	for shards > size/16 {
		shards /= 16
	}
	if shards == 0 {
		shards = 1
	}

	maxSize := size / shards
	maxCapacity := (capacity + shards - 1) / shards
	initialCapacity := (min(capacity, uint32(defaultInitialCapacity)) + shards - 1) / shards
	initialCapacity = max(initialCapacity, 1)
	initialSize := sizeForCapacity(initialCapacity, maxCapacity, maxSize)

	lrus := make([]LRU[K, V], shards)
	buckets := make([]uint32, initialSize*shards)
	elements := make([]element[K, V], initialSize*shards)

	from := 0
	to := int(initialSize)
	for i := range lrus {
		initLRU(&lrus[i], initialCapacity, initialSize, maxCapacity, maxSize, hash, buckets[from:to], elements[from:to])
		from = to
		to += int(initialSize)
	}

	return &ShardedLRU[K, V]{
		lrus:   lrus,
		mus:    make([]sync.RWMutex, shards),
		hash:   hash,
		shards: shards,
		mask:   shards - 1,
	}, nil
}

// Len returns the number of elements stored in the cache.
func (lru *ShardedLRU[K, V]) Len() (length int) {
	for shard := range lru.lrus {
		lru.mus[shard].RLock()
		length += lru.lrus[shard].Len()
		lru.mus[shard].RUnlock()
	}
	return
}

// AddWithLifetime adds a key:value to the cache with a lifetime.
// Returns true, true if key was updated and eviction occurred.
func (lru *ShardedLRU[K, V]) AddWithLifetime(key K, value V,
	lifetime time.Duration,
) (evicted bool) {
	hash := lru.hash(key)
	shard := (hash >> 16) & lru.mask

	lru.mus[shard].Lock()
	evicted = lru.lrus[shard].addWithLifetime(hash, key, value, lifetime)
	lru.mus[shard].Unlock()

	return
}

// Add adds a key:value to the cache.
// Returns true, true if key was updated and eviction occurred.
func (lru *ShardedLRU[K, V]) Add(key K, value V) (evicted bool) {
	hash := lru.hash(key)
	shard := (hash >> 16) & lru.mask

	lru.mus[shard].Lock()
	evicted = lru.lrus[shard].add(hash, key, value)
	lru.mus[shard].Unlock()

	return
}

// Get returns the value associated with the key, setting it as the most
// recently used item.
// If the found cache item is already expired, the evict function is called
// and the return value indicates that the key was not found.
func (lru *ShardedLRU[K, V]) Get(key K) (value V, ok bool) {
	hash := lru.hash(key)
	shard := (hash >> 16) & lru.mask

	lru.mus[shard].Lock()
	value, ok = lru.lrus[shard].get(hash, key)
	lru.mus[shard].Unlock()

	return
}

func (lru *ShardedLRU[K, V]) GetWithLifetime(key K) (value V, lifetime time.Time, ok bool) {
	hash := lru.hash(key)
	shard := (hash >> 16) & lru.mask

	lru.mus[shard].Lock()
	value, lifetime, ok = lru.lrus[shard].getWithLifetime(hash, key)
	lru.mus[shard].Unlock()

	return
}

func (lru *ShardedLRU[K, V]) GetWithLifetimeNoExpire(key K) (value V, lifetime time.Time, ok bool) {
	hash := lru.hash(key)
	shard := (hash >> 16) & lru.mask

	lru.mus[shard].Lock()
	value, lifetime, ok = lru.lrus[shard].getWithLifetimeNoExpire(hash, key)
	lru.mus[shard].Unlock()

	return
}

// GetAndRefresh returns the value associated with the key, setting it as the most
// recently used item.
// The lifetime of the found cache item is refreshed, even if it was already expired.
func (lru *ShardedLRU[K, V]) GetAndRefresh(key K) (value V, ok bool) {
	hash := lru.hash(key)
	shard := (hash >> 16) & lru.mask

	lru.mus[shard].Lock()
	value, ok = lru.lrus[shard].getAndRefresh(hash, key)
	lru.mus[shard].Unlock()

	return
}

func (lru *ShardedLRU[K, V]) GetAndRefreshOrAdd(key K, constructor func() (V, bool)) (value V, updated bool, ok bool) {
	hash := lru.hash(key)
	shard := (hash >> 16) & lru.mask

	lru.mus[shard].Lock()
	var currentTime int64
	value, updated, ok, currentTime = lru.lrus[shard].getAndRefreshOrAdd(hash, key, constructor)
	lru.mus[shard].Unlock()

	if !updated && ok {
		lru.sweepExpired(shard, currentTime)
	}
	return
}

func (lru *ShardedLRU[K, V]) sweepExpired(skipShard uint32, currentTime int64) {
	if lru.shards <= 1 {
		return
	}
	// The insertion already swept its own shard. Sweep one more shard to keep
	// idle shards bounded without walking and locking every shard on each miss.
	var shard uint32
	for {
		shard = (lru.sweep.Add(1) - 1) & lru.mask
		if shard != skipShard {
			break
		}
	}
	lru.mus[shard].Lock()
	lru.lrus[shard].purgeExpiredAt(currentTime)
	lru.mus[shard].Unlock()
}

// Peek looks up a key's value from the cache, without changing its recent-ness.
// If the found entry is already expired, the evict function is called.
func (lru *ShardedLRU[K, V]) Peek(key K) (value V, ok bool) {
	hash := lru.hash(key)
	shard := (hash >> 16) & lru.mask

	lru.mus[shard].Lock()
	value, ok = lru.lrus[shard].peek(hash, key)
	lru.mus[shard].Unlock()

	return
}

func (lru *ShardedLRU[K, V]) PeekWithLifetime(key K) (value V, lifetime time.Time, ok bool) {
	hash := lru.hash(key)
	shard := (hash >> 16) & lru.mask

	lru.mus[shard].Lock()
	value, lifetime, ok = lru.lrus[shard].peekWithLifetime(hash, key)
	lru.mus[shard].Unlock()

	return
}

func (lru *ShardedLRU[K, V]) UpdateLifetime(key K, value V, lifetime time.Duration) (ok bool) {
	hash := lru.hash(key)
	shard := (hash >> 16) & lru.mask

	lru.mus[shard].Lock()
	ok = lru.lrus[shard].updateLifetime(hash, key, value, lifetime)
	lru.mus[shard].Unlock()

	return
}

// Contains checks for the existence of a key, without changing its recent-ness.
// If the found entry is already expired, the evict function is called.
func (lru *ShardedLRU[K, V]) Contains(key K) (ok bool) {
	hash := lru.hash(key)
	shard := (hash >> 16) & lru.mask

	lru.mus[shard].Lock()
	ok = lru.lrus[shard].contains(hash, key)
	lru.mus[shard].Unlock()

	return
}

// Remove removes the key from the cache.
// The return value indicates whether the key existed or not.
// The evict function is called for the removed entry.
func (lru *ShardedLRU[K, V]) Remove(key K) (removed bool) {
	hash := lru.hash(key)
	shard := (hash >> 16) & lru.mask

	lru.mus[shard].Lock()
	removed = lru.lrus[shard].remove(hash, key)
	lru.mus[shard].Unlock()

	return
}

// RemoveOldest removes the oldest entry from the cache.
// Key, value and an indicator of whether the entry has been removed is returned.
// The evict function is called for the removed entry.
func (lru *ShardedLRU[K, V]) RemoveOldest() (key K, value V, removed bool) {
	hash := lru.hash(key)
	shard := (hash >> 16) & lru.mask

	lru.mus[shard].Lock()
	key, value, removed = lru.lrus[shard].RemoveOldest()
	lru.mus[shard].Unlock()

	return
}

// Keys returns a slice of the keys in the cache, from oldest to newest.
// Expired entries are not included.
// The evict function is called for each expired item.
func (lru *ShardedLRU[K, V]) Keys() []K {
	keys := make([]K, 0, lru.Len())
	for shard := range lru.lrus {
		lru.mus[shard].Lock()
		keys = append(keys, lru.lrus[shard].Keys()...)
		lru.mus[shard].Unlock()
	}

	return keys
}

// Purge purges all data (key and value) from the LRU.
// The evict function is called for each expired item.
// The LRU metrics are reset.
func (lru *ShardedLRU[K, V]) Purge() {
	for shard := range lru.lrus {
		lru.mus[shard].Lock()
		lru.lrus[shard].Purge()
		lru.mus[shard].Unlock()
	}
}

// PurgeExpired purges all expired items from the LRU.
// The evict function is called for each expired item.
func (lru *ShardedLRU[K, V]) PurgeExpired() {
	currentTime := now()
	for shard := range lru.lrus {
		lru.mus[shard].Lock()
		lru.lrus[shard].purgeExpiredAt(currentTime)
		lru.mus[shard].Unlock()
	}
}

// Metrics returns the metrics of the cache.
func (lru *ShardedLRU[K, V]) Metrics() Metrics {
	metrics := Metrics{}

	for shard := range lru.lrus {
		lru.mus[shard].Lock()
		m := lru.lrus[shard].Metrics()
		lru.mus[shard].Unlock()

		addMetrics(&metrics, m)
	}

	return metrics
}

// ResetMetrics resets the metrics of the cache and returns the previous state.
func (lru *ShardedLRU[K, V]) ResetMetrics() Metrics {
	metrics := Metrics{}

	for shard := range lru.lrus {
		lru.mus[shard].Lock()
		m := lru.lrus[shard].ResetMetrics()
		lru.mus[shard].Unlock()

		addMetrics(&metrics, m)
	}

	return metrics
}

func addMetrics(dst *Metrics, src Metrics) {
	dst.Inserts += src.Inserts
	dst.Collisions += src.Collisions
	dst.Evictions += src.Evictions
	dst.Removals += src.Removals
	dst.Hits += src.Hits
	dst.Misses += src.Misses
}

// just used for debugging
func (lru *ShardedLRU[K, V]) dump() {
	for shard := range lru.lrus {
		fmt.Printf("Shard %d:\n", shard)
		lru.mus[shard].RLock()
		lru.lrus[shard].dump()
		lru.mus[shard].RUnlock()
		fmt.Println("")
	}
}

func (lru *ShardedLRU[K, V]) PrintStats() {
	for shard := range lru.lrus {
		fmt.Printf("Shard %d:\n", shard)
		lru.mus[shard].RLock()
		lru.lrus[shard].PrintStats()
		lru.mus[shard].RUnlock()
		fmt.Println("")
	}
}
