package main

import (
	"bytes"
	"encoding/binary"
	"math/rand"
	"sort"
	"testing"

	"github.com/stretchr/testify/require"
)

// Search returns the smallest index i in [0, n) at which f(i) is true, assuming that on the range [0, n),
func Search(n int, f func(i int) bool) int {
	var i int
	for n > 2 {
		middle := i + (n >> 1)
		n = (n + 1) >> 1
		if !f(middle) {
			i = middle
		}
	}
	if n > 1 && !f(i) {
		i++
	}
	if n > 0 && !f(i) {
		i++
	}
	return i
}

func int64ToBz(n uint64) []byte {
	var key [8]byte
	binary.BigEndian.PutUint64(key[:], n)
	return key[:]
}

func genRandItems(n int) [][]byte {
	r := rand.New(rand.NewSource(0))
	items := make([][]byte, n)
	itemsM := make(map[uint64]bool)
	for i := 0; i < n; i++ {
		for {
			key := uint64(r.Int63n(10000000000000000))
			if !itemsM[key] {
				itemsM[key] = true
				items[i] = int64ToBz(key)
				break
			}
		}
	}
	return items
}

func generateRandomData(n int) [][]byte {
	data := make([][]byte, n)
	for i := 0; i < n; i++ {
		buf := make([]byte, 8)
		binary.BigEndian.PutUint64(buf, uint64(i))
		data[i] = buf
	}
	return data
}

func BenchmarkBSearch(b *testing.B) {
	items := genRandItems(1000000)
	target := items[500]
	sort.Slice(items, func(i, j int) bool {
		return bytes.Compare(items[i], items[j]) < 0
	})

	cmp := func(i int) bool {
		return bytes.Compare(items[i], target) >= 0
	}

	b.Run("std", func(b *testing.B) {
		i := sort.Search(len(items), cmp)
		require.Equal(b, target, items[i])

		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			i := sort.Search(len(items), cmp)
			_ = items[i]
		}
	})
	b.Run("new", func(b *testing.B) {
		i := Search(len(items), cmp)
		require.Equal(b, target, items[i])

		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			i := Search(len(items), cmp)
			_ = items[i]
		}
	})
}