safemap

func TestMap_Basic(t *testing.T)

{
	m := New[string, int]()

	m.Set("a", 1)
	m.Set("b", 2)

	if v, ok := m.Get("a"); !ok || v != 1 {
		t.Errorf("Get a: got %d, want 1", v)
	}

	if !m.Has("b") {
		t.Error("Has b: should be true")
	}

	if m.Has("c") {
		t.Error("Has c: should be false")
	}

	if m.Len() != 2 {
		t.Errorf("Len: got %d, want 2", m.Len())
	}
}

func TestMap_Range(t *testing.T)

{
	m := New[string, int]()
	m.Set("a", 1)
	m.Set("b", 2)
	m.Set("c", 3)

	count := 0
	m.Range(func(k string, v int) bool {
		count++
		return count < 2
	})

	if count != 2 {
		t.Errorf("Range should stop early, got %d iterations", count)
	}
}

func TestMap_GetOrSet(t *testing.T)

{
	m := New[string, int]()

	v1 := m.GetOrSet("key", 42)
	if v1 != 42 {
		t.Errorf("GetOrSet first call: got %d, want 42", v1)
	}

	v2 := m.GetOrSet("key", 100)
	if v2 != 42 {
		t.Errorf("GetOrSet second call: got %d, want 42", v2)
	}
}

func TestMap_Compute(t *testing.T)

{
	m := New[string, int]()

	m.Compute("counter", func(v int, exists bool) int {
		return v + 1
	})

	if v, _ := m.Get("counter"); v != 1 {
		t.Errorf("Compute first: got %d, want 1", v)
	}

	m.Compute("counter", func(v int, exists bool) int {
		return v + 1
	})

	if v, _ := m.Get("counter"); v != 2 {
		t.Errorf("Compute second: got %d, want 2", v)
	}
}

func TestMap_Concurrent(t *testing.T)

{
	m := New[int, int]()
	var wg sync.WaitGroup

	for i := 0; i < 100; i++ {
		wg.Add(1)
		go func(n int) {
			defer wg.Done()
			m.Set(n, n*10)
			m.Get(n)
			m.Has(n)
		}(i)
	}

	wg.Wait()

	if m.Len() != 100 {
		t.Errorf("Concurrent: got %d items, want 100", m.Len())
	}
}

func NewSharded[K comparable, V any](hasher Hasher[K], shardCount int) *ShardedMap[K, V]

{
	if shardCount <= 0 {
		shardCount = 32
	}
	// Ensure power of 2
	if bits.OnesCount(uint(shardCount)) != 1 {
		shardCount = 1 << bits.Len(uint(shardCount))
	}

	shards := make([]*Map[K, V], shardCount)
	for i := 0; i < shardCount; i++ {
		shards[i] = New[K, V]()
	}

	return &ShardedMap[K, V]{
		shards: shards,
		mask:   uint64(shardCount - 1),
		hasher: hasher,
	}
}

func TestShardedMap_Basic(t *testing.T)

{
	m := NewSharded[string, int](StringHasher, 4)

	m.Set("foo", 1)
	m.Set("bar", 2)

	if v, ok := m.Get("foo"); !ok || v != 1 {
		t.Errorf("Get(foo) = %v, %v; want 1, true", v, ok)
	}

	if v, ok := m.Get("bar"); !ok || v != 2 {
		t.Errorf("Get(bar) = %v, %v; want 2, true", v, ok)
	}

	if _, ok := m.Get("baz"); ok {
		t.Errorf("Get(baz) = _, true; want false")
	}

	if !m.Has("foo") {
		t.Errorf("Has(foo) = false; want true")
	}

	m.Delete("foo")
	if m.Has("foo") {
		t.Errorf("Has(foo) after delete = true; want false")
	}

	if m.Len() != 1 {
		t.Errorf("Len() = %d; want 1", m.Len())
	}

	m.Clear()
	if m.Len() != 0 {
		t.Errorf("Len() after Clear = %d; want 0", m.Len())
	}
}

func TestShardedMap_Range(t *testing.T)

{
	m := NewSharded[string, int](StringHasher, 4)
	count := 100
	for i := 0; i < count; i++ {
		m.Set(fmt.Sprintf("key-%d", i), i)
	}

	seen := 0
	m.Range(func(k string, v int) bool {
		seen++
		return true
	})

	if seen != count {
		t.Errorf("Range visited %d items; want %d", seen, count)
	}
}

func TestShardedMap_Concurrent(t *testing.T)

{
	m := NewSharded[string, int](StringHasher, 16)
	var wg sync.WaitGroup
	workers := 10
	ops := 1000

	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func(id int) {
			defer wg.Done()
			for j := 0; j < ops; j++ {
				key := fmt.Sprintf("key-%d-%d", id, j)
				m.Set(key, j) // Write

				// Read same key or random key
				rKey := fmt.Sprintf("key-%d-%d", rand.Intn(workers), rand.Intn(ops))
				m.Get(rKey)

				if j%10 == 0 {
					m.Delete(key)
				}
			}
		}(i)
	}

	wg.Wait()
}

func TestShardedMap_GetOrSet(t *testing.T)

{
	m := NewSharded[string, int](StringHasher, 4)
	v := m.GetOrSet("foo", 10)
	if v != 10 {
		t.Errorf("GetOrSet(foo) = %d; want 10", v)
	}

	v = m.GetOrSet("foo", 20)
	if v != 10 {
		t.Errorf("GetOrSet(foo) existing = %d; want 10", v)
	}
}

func TestShardedMap_Compute(t *testing.T)

{
	m := NewSharded[string, int](StringHasher, 4)
	m.Set("foo", 1)

	v := m.Compute("foo", func(existing int, exists bool) int {
		return existing + 1
	})

	if v != 2 {
		t.Errorf("Compute(foo) = %d; want 2", v)
	}

	v = m.Compute("bar", func(existing int, exists bool) int {
		if !exists {
			return 5
		}
		return existing
	})

	if v != 5 {
		t.Errorf("Compute(bar) = %d; want 5", v)
	}
}

func TestShardedMap_KeysDistribution(t *testing.T)

{
	// Verify that keys are actually distributed across shards
	// This uses internal knowledge check
	m := NewSharded[string, int](StringHasher, 4)
	count := 1000
	for i := 0; i < count; i++ {
		m.Set(strconv.Itoa(i), i)
	}

	// Each shard should have roughly count/4 items
	// Not strict check due to potential collision, but shouldn't be 0
	for i, shard := range m.shards {
		if shard.Len() == 0 {
			t.Errorf("Shard %d is empty, bad distribution or small sample", i)
		}
	}
}