package locker

import "sync"

type Locker interface {
	AssignLocker(key string) sync.Locker
	ReleaseLocker(key string)
}

func New() Locker {
	return &mapLocker{m: make(map[string]*mutexCount)}
}

type mapLocker struct {
	sync.Mutex

	m map[string]*mutexCount
}

type mutexCount struct {
	sync.Mutex
	count int
}

func (ml *mapLocker) AssignLocker(key string) sync.Locker {
	ml.Lock()
	defer ml.Unlock()

	if l, ok := ml.m[key]; ok {
		l.count++
		return l
	}
	mu := &mutexCount{count: 1}
	ml.m[key] = mu
	return mu
}

func (ml *mapLocker) ReleaseLocker(key string) {
	ml.Lock()
	defer ml.Unlock()

	if l, ok := ml.m[key]; ok {
		l.count--
		if l.count == 0 {
			delete(ml.m, key)
		}
	}
}