// Copyright 2024 ByteDance Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package span

import (
	"fmt"
	"sync"
	"testing"

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

func TestSpanClass(t *testing.T) {
	assert.Equal(t, spanClass(minSpanObject), minSpanClass)

	assert.Equal(t, spanClass(1), 1)

	assert.Equal(t, spanClass(2), 2)
	assert.Equal(t, spanClass(3), 2)

	assert.Equal(t, spanClass(4), 3)
	assert.Equal(t, spanClass(5), 3)
	assert.Equal(t, spanClass(6), 3)
	assert.Equal(t, spanClass(7), 3)

	assert.Equal(t, spanClass(8), 4)
	assert.Equal(t, spanClass(9), 4)

	assert.Equal(t, spanClass(32), 6)
	assert.Equal(t, spanClass(33), 6)
	assert.Equal(t, spanClass(63), 6)

	assert.Equal(t, spanClass(64), 7)
	assert.Equal(t, spanClass(65), 7)
	assert.Equal(t, spanClass(127), 7)
	assert.Equal(t, spanClass(128), 8)
	assert.Equal(t, spanClass(129), 8)
}

func TestSpan(t *testing.T) {
	bc := NewSpan(1024 * 32)

	var wg sync.WaitGroup
	for c := 0; c < 12; c++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := 0; i < 1024; i++ {
				buf := []byte("123")
				assert.Equal(t, bc.Copy(buf), buf)

				buf = make([]byte, 32)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 63)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 64)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 1024)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 32*1024)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 64*1024)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 128*1024)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
			}
		}()
	}
	wg.Wait()
}

func TestSpanCache(t *testing.T) {
	bc := NewSpanCache(1024 * 32)

	var wg sync.WaitGroup
	for c := 0; c < 12; c++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := 0; i < 1024; i++ {
				buf := []byte("123")
				assert.Equal(t, bc.Copy(buf), buf)

				buf = make([]byte, 32)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 63)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 64)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 1024)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 32*1024)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 64*1024)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 128*1024)
				assert.Equal(t, len(bc.Copy(buf)), len(buf))
			}
		}()
	}
	wg.Wait()
}

var benchStringSizes = []int{
	15, 31, 127, // not cached
	128, 256, 1024 - 1, 4*1024 - 1,
	32*1024 - 1, 128*1024 - 1,
	512*1024 - 1, // not cached
}

func BenchmarkSpanCacheCopy(b *testing.B) {
	for _, sz := range benchStringSizes {
		b.Run(fmt.Sprintf("size=%d", sz), func(b *testing.B) {
			from := make([]byte, sz)
			sc := NewSpanCache(maxSpanObject * 8)
			b.RunParallel(func(pb *testing.PB) {
				b.ReportAllocs()
				b.ResetTimer()
				var buffer []byte
				for pb.Next() {
					buffer = sc.Copy(from)
					buffer[0] = 'a'
					buffer[sz-1] = 'z'
				}
				_ = buffer
			})
		})
	}
}

func BenchmarkMakeAndCopy(b *testing.B) {
	for _, sz := range benchStringSizes {
		b.Run(fmt.Sprintf("size=%d", sz), func(b *testing.B) {
			from := make([]byte, sz)
			b.RunParallel(func(pb *testing.PB) {
				b.ReportAllocs()
				b.ResetTimer()
				var buffer []byte
				for pb.Next() {
					buffer = make([]byte, sz)
					copy(buffer, from)
					buffer[0] = 'a'
					buffer[sz-1] = 'z'
				}
				_ = buffer
			})
		})
	}
}