package lfu

import (
	"sync"

	"github.com/pkg/errors"
)

type Synchronizer[K comparable] struct {
	index *Map[K, *sync.RWMutex]
}

func (s *Synchronizer[K]) Remove(key K) {
	s.index.Delete(key)
}

func (s *Synchronizer[K]) ReadTx(key K, fn func(upgrade func(fn func())) error) error {
	mutex, _ := s.index.GetOrSet(key, &sync.RWMutex{})
	mutex.RLock()
	defer mutex.RUnlock()

	upgrade := func(fn func()) {
		mutex.RUnlock()
		mutex.Lock()
		defer func() {
			mutex.Unlock()
			mutex.RLock()
		}()

		fn()
	}

	if err := fn(upgrade); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func (s *Synchronizer[K]) WriteTx(key K, fn func() error) error {
	mutex, _ := s.index.GetOrSet(key, &sync.RWMutex{})
	mutex.Lock()
	defer mutex.Unlock()

	if err := fn(); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func NewSynchronizer[K comparable]() *Synchronizer[K] {
	return &Synchronizer[K]{
		index: NewMap[K, *sync.RWMutex](),
	}
}