package utils

import "sync"

type Key interface {
	int | int32 | int64 | string
}

type SafeMap[K Key, V any] struct {
	m    map[K]V
	lock sync.RWMutex
}

func NewSafeMap[K Key, V any]() *SafeMap[K, V] {
	return &SafeMap[K, V]{
		m: make(map[K]V),
	}
}

func (sm *SafeMap[K, V]) Set(key K, value V) {
	sm.lock.Lock()
	defer sm.lock.Unlock()
	sm.m[key] = value
}

func (sm *SafeMap[K, V]) Get(key K) (V, bool) {
	sm.lock.RLock()
	defer sm.lock.RUnlock()
	value, ok := sm.m[key]
	return value, ok
}

func (sm *SafeMap[K, V]) GetWithSetDefault(key K, defaultValue V) V {
	sm.lock.Lock()
	defer sm.lock.Unlock()
	if value, ok := sm.m[key]; ok {
		return value
	}
	sm.m[key] = defaultValue
	return defaultValue
}

func (sm *SafeMap[K, V]) Delete(key K) {
	sm.lock.Lock()
	defer sm.lock.Unlock()
	delete(sm.m, key)
}

func (sm *SafeMap[K, V]) Range(f func(key K, value V) bool) {
	sm.lock.RLock()
	defer sm.lock.RUnlock()
	for k, v := range sm.m {
		if !f(k, v) {
			break
		}
	}
}

func (sm *SafeMap[K, V]) Len() int {
	sm.lock.RLock()
	defer sm.lock.RUnlock()
	return len(sm.m)
}

func (sm *SafeMap[K, V]) Keys() []K {
	sm.lock.RLock()
	defer sm.lock.RUnlock()
	keys := make([]K, 0, len(sm.m))
	for k := range sm.m {
		keys = append(keys, k)
	}
	return keys
}

func (sm *SafeMap[K, V]) Values() []V {
	sm.lock.RLock()
	defer sm.lock.RUnlock()
	values := make([]V, 0, len(sm.m))
	for _, v := range sm.m {
		values = append(values, v)
	}
	return values
}

func (sm *SafeMap[K, V]) Clear() {
	sm.lock.Lock()
	defer sm.lock.Unlock()
	sm.m = make(map[K]V)
}

func (sm *SafeMap[K, V]) Copy() map[K]V {
	sm.lock.RLock()
	defer sm.lock.RUnlock()
	m := make(map[K]V, len(sm.m))
	for k, v := range sm.m {
		m[k] = v
	}
	return m
}

func (sm *SafeMap[K, V]) Merge(m map[K]V) {
	sm.lock.Lock()
	defer sm.lock.Unlock()
	for k, v := range m {
		sm.m[k] = v
	}
}
