package core

import (
	"sync"
	"sync/atomic"
	"time"
)

// UpgradeableMutex wraps a RWMutex and provides safe upgrading / downgrading
// by informing you whether a write lock was achieved in the brief swap time
type UpgradeableMutex struct {
	wLast    int64
	internal sync.RWMutex
}

// RLock exactly wraps the internal RWMutex
func (mu *UpgradeableMutex) RLock() {
	mu.internal.RLock()
}

// RUnlock exactly wraps the internal RWMutex
func (mu *UpgradeableMutex) RUnlock() {
	mu.internal.RUnlock()
}

// Lock wraps the internal RWMutex, atomically storing the last write-lock time
func (mu *UpgradeableMutex) Lock() {
	mu.internal.Lock()
	atomic.StoreInt64(&mu.wLast, time.Now().UnixNano())
}

// Unlock exactly wraps the internal RWMutex
func (mu *UpgradeableMutex) Unlock() {
	mu.internal.Unlock()
}

// safeSwap stores the current time, performs the swap function and checks if
// any write locks were achieved during the swap function
func (mu *UpgradeableMutex) safeSwap(swapFn func()) bool {
	// Get the 'now' time
	now := time.Now().UnixNano()

	// Store now time
	atomic.StoreInt64(&mu.wLast, now)

	// Perform the swap
	swapFn()

	// Successful swap determined by if last write-lock
	// is still equal to 'now'
	return atomic.LoadInt64(&mu.wLast) == now
}

// UpgradeLock upgrades a read to a write lock, returning success state as a bool
func (mu *UpgradeableMutex) UpgradeLock() bool {
	return mu.safeSwap(func() {
		mu.internal.RUnlock()
		mu.internal.Lock()
	})
}

// DowngradeLock downgrades a write to a read lock, returning success state as a bool
func (mu *UpgradeableMutex) DowngradeLock() bool {
	return mu.safeSwap(func() {
		mu.internal.Unlock()
		mu.internal.RLock()
	})
}