package keyedmutex

import (
	"container/list"
	"sync"
)

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

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

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

func (m KeyedMutex) lock(key string) 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) TryLock(key string) 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) Unlock(key string) {
	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{}{}
	}
}