package dataext

import (
	"context"
	"sync"
	"time"
)

// MultiMutex is a simple map[key -> mutex]
type MultiMutex[TKey comparable] struct {
	mutextMap *SyncMap[TKey, *CASMutex]
}

func NewMultiMutex[TKey comparable]() *MultiMutex[TKey] {
	return &MultiMutex[TKey]{
		mutextMap: NewSyncMap[TKey, *CASMutex](),
	}
}

// TryLockWithContext attempts to acquire the lock, blocking until resources
// are available or ctx is done (timeout or cancellation).
func (mm *MultiMutex[TKey]) TryLockWithContext(ctx context.Context, key TKey) bool {
	lck := mm.mutextMap.GetAndSetIfNotContainsFunc(key, NewCASMutex)
	return lck.TryLockWithContext(ctx)
}

// Lock acquires the lock.
// If it is currently held by others, Lock will wait until it has a chance to acquire it.
func (mm *MultiMutex[TKey]) Lock(key TKey) {
	lck := mm.mutextMap.GetAndSetIfNotContainsFunc(key, NewCASMutex)
	lck.Lock()
}

// TryLock attempts to acquire the lock without blocking.
// Return false if someone is holding it now.
func (mm *MultiMutex[TKey]) TryLock(key TKey) bool {
	lck := mm.mutextMap.GetAndSetIfNotContainsFunc(key, NewCASMutex)
	return lck.TryLock()
}

// TryLockWithTimeout attempts to acquire the lock within a period of time.
// Return false if spending time is more than duration and no chance to acquire it.
func (mm *MultiMutex[TKey]) TryLockWithTimeout(key TKey, duration time.Duration) bool {
	lck := mm.mutextMap.GetAndSetIfNotContainsFunc(key, NewCASMutex)
	return lck.TryLockWithTimeout(duration)
}

// Unlock releases the lock.
func (mm *MultiMutex[TKey]) Unlock(key TKey) {
	lck := mm.mutextMap.GetAndSetIfNotContainsFunc(key, NewCASMutex)
	lck.Unlock()
}

// RTryLockWithContext attempts to acquire the read lock, blocking until resources
// are available or ctx is done (timeout or cancellation).
func (mm *MultiMutex[TKey]) RTryLockWithContext(ctx context.Context, key TKey) bool {
	lck := mm.mutextMap.GetAndSetIfNotContainsFunc(key, NewCASMutex)
	return lck.RTryLockWithContext(ctx)
}

// RLock acquires the read lock.
// If it is currently held by others writing, RLock will wait until it has a chance to acquire it.
func (mm *MultiMutex[TKey]) RLock(key TKey) {
	lck := mm.mutextMap.GetAndSetIfNotContainsFunc(key, NewCASMutex)
	lck.RLock()
}

// RTryLock attempts to acquire the read lock without blocking.
// Return false if someone is writing it now.
func (mm *MultiMutex[TKey]) RTryLock(key TKey) bool {
	lck := mm.mutextMap.GetAndSetIfNotContainsFunc(key, NewCASMutex)
	return lck.RTryLock()
}

// RTryLockWithTimeout attempts to acquire the read lock within a period of time.
// Return false if spending time is more than duration and no chance to acquire it.
func (mm *MultiMutex[TKey]) RTryLockWithTimeout(duration time.Duration, key TKey) bool {
	lck := mm.mutextMap.GetAndSetIfNotContainsFunc(key, NewCASMutex)
	return lck.RTryLockWithTimeout(duration)
}

// RUnlock releases the read lock.
func (mm *MultiMutex[TKey]) RUnlock(key TKey) {
	lck := mm.mutextMap.GetAndSetIfNotContainsFunc(key, NewCASMutex)
	lck.RUnlock()
}

// RLocker returns a Locker interface that implements the Lock and Unlock methods
// by calling CASMutex.RLock and CASMutex.RUnlock.
func (mm *MultiMutex[TKey]) RLocker(key TKey) sync.Locker {
	lck := mm.mutextMap.GetAndSetIfNotContainsFunc(key, NewCASMutex)
	return lck.RLocker()
}

// Get returns a Locker interface
func (mm *MultiMutex[TKey]) Get(key TKey) sync.Locker {
	lck := mm.mutextMap.GetAndSetIfNotContainsFunc(key, NewCASMutex)
	return lck
}