/*
 * Copyright 2024 CloudWeGo Authors
 *
 * 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 mem

import (
	"fmt"
	"sync"
	"sync/atomic"
	"testing"

	"github.com/cloudwego/kitex/internal/test"
)

func TestSpanClass(t *testing.T) {
	test.DeepEqual(t, spanClass(minSpanObject), minSpanClass)

	test.DeepEqual(t, spanClass(1), 1)

	test.DeepEqual(t, spanClass(2), 2)
	test.DeepEqual(t, spanClass(3), 2)

	test.DeepEqual(t, spanClass(4), 3)
	test.DeepEqual(t, spanClass(5), 3)
	test.DeepEqual(t, spanClass(6), 3)
	test.DeepEqual(t, spanClass(7), 3)

	test.DeepEqual(t, spanClass(8), 4)
	test.DeepEqual(t, spanClass(9), 4)

	test.DeepEqual(t, spanClass(32), 6)
	test.DeepEqual(t, spanClass(33), 6)
	test.DeepEqual(t, spanClass(63), 6)

	test.DeepEqual(t, spanClass(64), 7)
	test.DeepEqual(t, spanClass(65), 7)
	test.DeepEqual(t, spanClass(127), 7)
	test.DeepEqual(t, spanClass(128), 8)
	test.DeepEqual(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 < 128; i++ {
				buf := []byte("123")
				test.DeepEqual(t, bc.Copy(buf), buf)

				buf = make([]byte, 32)
				test.DeepEqual(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 63)
				test.DeepEqual(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 64)
				test.DeepEqual(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 1024)
				test.DeepEqual(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 32*1024)
				test.DeepEqual(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 64*1024)
				test.DeepEqual(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 128*1024)
				test.DeepEqual(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 < 128; i++ {
				buf := []byte("123")
				test.DeepEqual(t, bc.Copy(buf), buf)

				buf = make([]byte, 32)
				test.DeepEqual(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 63)
				test.DeepEqual(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 64)
				test.DeepEqual(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 1024)
				test.DeepEqual(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 32*1024)
				test.DeepEqual(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 64*1024)
				test.DeepEqual(t, len(bc.Copy(buf)), len(buf))
				buf = make([]byte, 128*1024)
				test.DeepEqual(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)
			buffers := make([][]byte, 1024)
			var bidx int32
			b.ReportAllocs()
			b.ResetTimer()
			b.RunParallel(func(pb *testing.PB) {
				idx := atomic.AddInt32(&bidx, 1)
				for pb.Next() {
					buffer := sc.Copy(from)
					buffer[0] = 'a'
					buffer[sz-1] = 'z'
					buffers[idx] = buffer
				}
			})
			_ = buffers
		})
	}
}

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