package keyedmutex

import (
	"sync"
)

type keyedLock struct {
	lock  sync.Mutex
	count int
}

type KeyedMutex[K comparable] struct {
	mu    sync.Mutex
	byKey map[K]*keyedLock
}

func New[K comparable]() *KeyedMutex[K] {
	return &KeyedMutex[K]{
		byKey: map[K]*keyedLock{},
	}
}

func (m *KeyedMutex[K]) getLock(key K) *keyedLock {
	m.mu.Lock()
	defer m.mu.Unlock()

	item, ok := m.byKey[key]
	if !ok {
		item = &keyedLock{}
		m.byKey[key] = item
	}
	item.count++
	return item
}

func (m *KeyedMutex[K]) release(key K, unlock bool) {
	m.mu.Lock()
	defer m.mu.Unlock()

	item, ok := m.byKey[key]
	if !ok {
		panic("unlock of unlocked mutex")
	}

	item.count--
	if unlock {
		item.lock.Unlock()
	}

	if item.count == 0 {
		delete(m.byKey, key)
	}
}

func (m *KeyedMutex[K]) Lock(key K) {
	m.getLock(key).lock.Lock()
}

func (m *KeyedMutex[K]) TryLock(key K) bool {
	if ok := m.getLock(key).lock.TryLock(); !ok {
		m.release(key, false)
		return false
	}
	return true
}

func (m *KeyedMutex[K]) Unlock(key K) {
	m.release(key, true)
}