package container

import (
	"slices"
	"sync"
)

// M is ordered and concurent safe map implementation.
type M[K comparable, V any] struct {
	m  map[K]V
	k  []K
	mu sync.RWMutex
}

// Make returns new instance with `capacity` initial capacity.
func Make[K comparable, V any](capacity int) M[K, V] {
	return M[K, V]{
		m: make(map[K]V, capacity),
		k: make([]K, 0, capacity),
	}
}

// Get value `V` by key `k K`. Returns zero value if the key is absent.
func (m *M[K, V]) Get(k K) V {
	m.mu.RLock()
	defer m.mu.RUnlock()

	return m.m[k]
}

// Get2 value `v V` by key `k K` and a presence flag `ok bool`. Returns
// `(value, true)` if the key exists, and `(zero value, false)` if the key is
// absent.
func (m *M[K, V]) Get2(k K) (v V, ok bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	v, ok = m.m[k]

	return v, ok
}

// Has returns true if key `k K` exists in map.
func (m *M[K, V]) Has(k K) bool {
	m.mu.RLock()
	defer m.mu.RUnlock()

	_, ok := m.m[k]

	return ok
}

// Set value `v V` to key `k K`.
func (m *M[K, V]) Set(k K, v V) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if _, exists := m.m[k]; !exists {
		m.k = append(m.k, k)
	}

	m.m[k] = v
}

// Delete key `k K` from map.
func (m *M[K, V]) Delete(k K) {
	m.mu.Lock()
	defer m.mu.Unlock()

	if _, ok := m.m[k]; !ok {
		return
	}

	idx := slices.Index(m.k, k)
	m.k = slices.Delete(m.k, idx, idx+1)

	delete(m.m, k)
}

// Keys is iterator over keys.
func (m *M[K, V]) Keys(yield func(K) bool) {
	for _, k := range m.k {
		if !yield(k) {
			break
		}
	}
}

// Values is iterator over values.
func (m *M[K, V]) Values(yield func(V) bool) {
	for _, k := range m.k {
		if !yield(m.m[k]) {
			break
		}
	}
}

// Entries is iterator over key:value pairs.
func (m *M[K, V]) Entries(yield func(K, V) bool) {
	for _, k := range m.k {
		if !yield(k, m.m[k]) {
			break
		}
	}
}