package keyspacemigration

import (
	"sync"
	"sync/atomic"
	"testing"
)

func TestTryAcquireExclusive(t *testing.T) {
	l := New()

	release, ok := l.TryAcquire("ks-a")
	if !ok {
		t.Fatal("first acquire should succeed")
	}
	if !l.IsLocked("ks-a") {
		t.Fatal("ks-a should be locked after acquire")
	}
	if l.IsLocked("ks-b") {
		t.Fatal("ks-b must not be locked")
	}
	if _, ok := l.TryAcquire("ks-a"); ok {
		t.Fatal("second acquire of same key should fail")
	}

	// A different keyspace can be acquired independently.
	releaseB, ok := l.TryAcquire("ks-b")
	if !ok {
		t.Fatal("acquire of distinct key should succeed")
	}

	release()
	if l.IsLocked("ks-a") {
		t.Fatal("ks-a should be released")
	}
	if !l.IsLocked("ks-b") {
		t.Fatal("ks-b should still be held")
	}

	// Re-acquire after release works.
	release2, ok := l.TryAcquire("ks-a")
	if !ok {
		t.Fatal("re-acquire after release should succeed")
	}
	release2()
	releaseB()
}

func TestReleaseIsIdempotent(t *testing.T) {
	l := New()
	release, ok := l.TryAcquire("ks")
	if !ok {
		t.Fatal("acquire should succeed")
	}

	release()
	release() // must not panic or affect a future holder

	release2, ok := l.TryAcquire("ks")
	if !ok {
		t.Fatal("second acquire should succeed")
	}
	if !l.IsLocked("ks") {
		t.Fatal("ks should be held by second acquire")
	}
	// First release closure must not free the second holder's lock.
	release()
	if !l.IsLocked("ks") {
		t.Fatal("stale release closure must not affect second holder")
	}
	release2()
	if l.IsLocked("ks") {
		t.Fatal("second release should free the lock")
	}
}

func TestTryAcquireConcurrent(t *testing.T) {
	l := New()
	const goroutines = 64

	var wg sync.WaitGroup
	var winners int32
	wg.Add(goroutines)
	for i := 0; i < goroutines; i++ {
		go func() {
			defer wg.Done()
			if release, ok := l.TryAcquire("ks"); ok {
				atomic.AddInt32(&winners, 1)
				release()
			}
		}()
	}
	wg.Wait()

	// At least one goroutine wins; we don't require an exact count
	// because the lock is only held briefly so others may also see it free.
	if winners == 0 {
		t.Fatal("expected at least one goroutine to acquire the lock")
	}
	if l.IsLocked("ks") {
		t.Fatal("lock must be free once all goroutines released")
	}
}