package util

import (
	"sync"
	"sync/atomic"
)

type SafeMap[K comparable, V any] struct {
	m      sync.Map
	length int64 // 用于记录映射的长度
}

func (sm *SafeMap[K, V]) Set(key K, value V) {
	sm.m.Store(key, value)
	atomic.AddInt64(&sm.length, 1)
}

func (sm *SafeMap[K, V]) Get(key K) (V, bool) {
	value, ok := sm.m.Load(key)
	if !ok {
		var zero V
		return zero, false
	}
	return value.(V), true
}

func (sm *SafeMap[K, V]) Delete(key K) {
	sm.m.Delete(key)
	atomic.AddInt64(&sm.length, -1)
}

func (sm *SafeMap[K, V]) Length() int {
	return int(atomic.LoadInt64(&sm.length))
}

func (sm *SafeMap[K, V]) Range(fn func(V) bool) {
	sm.m.Range(func(key, value any) bool {
		return fn(value.(V))
	})
}

func (sm *SafeMap[K, V]) Clear() {
	sm.m = sync.Map{}
	sm.length = 0
}