// Package util provide some utility functions
package util

import (
	"iter"
	"sync"
)

// TypedMap define a type-safe generic wrapper around sync.Map.
// It keeps the concurrency guarantees of sync.Map while removing the
// per-call-site type assertions previously required for every Load.
type TypedMap[K comparable, V any] struct {
	m sync.Map
}

// Load define func of return the value stored for key, and whether it was found.
func (t *TypedMap[K, V]) Load(key K) (V, bool) {
	value, ok := t.m.Load(key)
	if !ok {
		var zero V
		return zero, false
	}
	// safe: only values of type V are ever stored through this wrapper
	return value.(V), true
}

// Store define func of set the value for key.
func (t *TypedMap[K, V]) Store(key K, value V) {
	t.m.Store(key, value)
}

// LoadOrStore define func of return the existing value for key if present.
// Otherwise it stores and returns the given value. loaded reports whether the
// value was already present.
func (t *TypedMap[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
	v, loaded := t.m.LoadOrStore(key, value)
	return v.(V), loaded
}

// Swap define func of store value for key and return the previous value (if any).
// loaded reports whether a value was already present for key.
func (t *TypedMap[K, V]) Swap(key K, value V) (previous V, loaded bool) {
	prev, loaded := t.m.Swap(key, value)
	if !loaded {
		var zero V
		return zero, false
	}
	return prev.(V), true
}

// Delete define func of remove the value for key.
func (t *TypedMap[K, V]) Delete(key K) {
	t.m.Delete(key)
}

// Range define func of iterate over all key/value pairs.
// Iteration stops early if f returns false.
func (t *TypedMap[K, V]) Range(f func(key K, value V) bool) {
	t.m.Range(func(key, value any) bool {
		return f(key.(K), value.(V))
	})
}

// Len define func of return the number of key/value pairs currently stored.
// It walks the map, so the cost is O(n).
func (t *TypedMap[K, V]) Len() int {
	count := 0
	t.m.Range(func(_, _ any) bool {
		count++
		return true
	})
	return count
}

// All define func of return an iterator over all key/value pairs,
// usable directly with range-over-func. Iteration stops early if the
// consumer breaks out of the loop.
func (t *TypedMap[K, V]) All() iter.Seq2[K, V] {
	return func(yield func(K, V) bool) {
		t.m.Range(func(key, value any) bool {
			return yield(key.(K), value.(V))
		})
	}
}