package keymutex

import "sync"

type entry struct {
	mu   sync.Mutex
	refs int
}

type KeyMutex[K comparable] struct {
	mu      sync.Mutex
	entries map[K]*entry
}

func NewKeyMutex[K comparable]() *KeyMutex[K] {
	return &KeyMutex[K]{entries: make(map[K]*entry)}
}

func (k *KeyMutex[K]) Lock(key K) {
	k.mu.Lock()
	e := k.entries[key]
	if e == nil {
		e = &entry{}
		k.entries[key] = e
	}
	e.refs++
	k.mu.Unlock()

	e.mu.Lock()
}

func (k *KeyMutex[K]) Unlock(key K) {
	k.mu.Lock()
	e := k.entries[key]
	k.mu.Unlock()

	e.mu.Unlock()
	k.release(key, e)
}

func (k *KeyMutex[K]) TryLock(key K) bool {
	k.mu.Lock()
	e := k.entries[key]
	if e == nil {
		e = &entry{}
		k.entries[key] = e
	}
	e.refs++
	k.mu.Unlock()

	if e.mu.TryLock() {
		return true
	}
	k.release(key, e)
	return false
}

func (k *KeyMutex[K]) Do(key K, f func()) {
	k.Lock(key)
	defer k.Unlock(key)
	f()
}

func (k *KeyMutex[K]) release(key K, e *entry) {
	k.mu.Lock()
	e.refs--
	if e.refs == 0 {
		delete(k.entries, key)
	}
	k.mu.Unlock()
}