package keyedmutex

import (
	"container/list"
	"sync"
)

type KeyedMutex[K comparable] struct {
	mu       *sync.Mutex
	waitList map[K]*list.List
}

func New[K comparable]() KeyedMutex[K] {
	return KeyedMutex[K]{
		mu:       new(sync.Mutex),
		waitList: map[K]*list.List{},
	}
}

func (m KeyedMutex[K]) Lock(key K) {
	if ch := m.lock(key); ch != nil {
		<-ch
	}
}

func (m KeyedMutex[K]) lock(key K) chan struct{} {
	m.mu.Lock()
	defer m.mu.Unlock()

	if waitList, ok := m.waitList[key]; ok {
		ch := make(chan struct{})
		waitList.PushBack(ch)
		return ch
	}

	m.waitList[key] = list.New()
	return nil
}

func (m KeyedMutex[K]) TryLock(key K) bool {
	m.mu.Lock()
	defer m.mu.Unlock()

	if _, ok := m.waitList[key]; ok {
		return false
	}

	m.waitList[key] = list.New()
	return true
}

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

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

	if waitList.Len() == 0 {
		delete(m.waitList, key)
	} else {
		ch := waitList.Remove(waitList.Front()).(chan struct{})
		ch <- struct{}{}
	}
}