package container

import (
	"iter"
	"sync"
)

// Map is a thread-safe generic map.
type Map[K comparable, V any] struct {
	m  map[K]V
	mu *sync.Mutex
}

func NewMap[K comparable, V any]() *Map[K, V] {
	return &Map[K, V]{
		m:  make(map[K]V),
		mu: &sync.Mutex{},
	}
}

func (m *Map[K, V]) Load(key K) (V, bool) {
	m.mu.Lock()
	defer m.mu.Unlock()

	v, ok := m.m[key]
	return v, ok
}

func (m *Map[K, V]) LoadAndDelete(key K) (V, bool) {
	m.mu.Lock()
	defer m.mu.Unlock()

	v, ok := m.m[key]
	if ok {
		delete(m.m, key)
	}
	return v, ok
}

func (m *Map[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
	m.mu.Lock()
	defer m.mu.Unlock()

	actual, loaded = m.m[key]
	if !loaded {
		m.m[key] = value
		actual = value
	}
	return
}

func (m *Map[K, V]) Store(key K, value V) {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.m[key] = value
}

func (m *Map[K, V]) Delete(key K) {
	m.mu.Lock()
	defer m.mu.Unlock()

	delete(m.m, key)
}

// All return a iterator of the map and the code in for...range is in lock.
//
// For example:
//
//	for k, v := range s.All() { // do something in lock }
func (m *Map[K, V]) All() iter.Seq2[K, V] {
	return func(yield func(K, V) bool) {
		m.mu.Lock()
		defer m.mu.Unlock()

		for k, v := range m.m {
			if !yield(k, v) {
				return
			}
		}
	}
}

// Keys is similar to All and returns keys of the map.
func (m *Map[K, V]) Keys() iter.Seq[K] {
	return func(yield func(K) bool) {
		m.mu.Lock()
		defer m.mu.Unlock()

		for k := range m.m {
			if !yield(k) {
				return
			}
		}
	}
}

// Values is similar to All and returns keys of the map.
func (m *Map[K, V]) Values() iter.Seq[V] {
	return func(yield func(V) bool) {
		m.mu.Lock()
		defer m.mu.Unlock()

		for _, v := range m.m {
			if !yield(v) {
				return
			}
		}
	}
}