package lfu

import (
	"sync"
	"sync/atomic"
)

type Map[K comparable, V any] struct {
	size  atomic.Int32
	inner sync.Map
}

func (m *Map[K, V]) Get(key K) (V, bool) {
	raw, exists := m.inner.Load(key)
	if !exists {
		return *new(V), false
	}

	value, ok := raw.(V)
	if !ok {
		return *new(V), false
	}

	return value, true
}

func (m *Map[K, V]) GetOrSet(key K, defaultValue V) (V, bool) {
	raw, loaded := m.inner.LoadOrStore(key, defaultValue)
	if !loaded {
		m.size.Add(1)
	}

	value, ok := raw.(V)
	if !ok {
		return *new(V), loaded
	}

	return value, loaded
}

func (m *Map[K, V]) Set(key K, value V) {
	_, loaded := m.inner.Swap(key, value)
	if !loaded {
		m.size.Add(1)
	}
}

func (m *Map[K, V]) Delete(key K) {
	_, existed := m.inner.LoadAndDelete(key)
	if existed {
		m.size.Add(-1)
	}
}

func (m *Map[K, V]) Range(fn func(key, value any) bool) {
	m.inner.Range(fn)
}

func (m *Map[K, V]) Len() int {
	return int(m.size.Load())
}

func NewMap[K comparable, V any]() *Map[K, V] {
	return &Map[K, V]{
		inner: sync.Map{},
	}
}