package blobstor

import (
	"encoding/binary"
	"fmt"
	"os"
	"testing"

	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/blobstor/blobovniczatree"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/blobstor/common"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/blobstor/fstree"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/blobstor/memstore"
	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	cidtest "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id/test"
	objectSDK "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
	oid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/id"
	oidtest "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/id/test"
	"github.com/stretchr/testify/require"
	"go.uber.org/atomic"
	"golang.org/x/exp/rand"
	"golang.org/x/exp/slices"
)

// The storages to benchmark. Each storage has a description and a function which returns the actual
// storage along with a cleanup function.
var storages = []struct {
	desc   string
	create func(*testing.B) (common.Storage, func())
}{
	{
		desc: "memstore",
		create: func(*testing.B) (common.Storage, func()) {
			return memstore.New(), func() {}
		},
	},
	{
		desc: "fstree_nosync",
		create: func(b *testing.B) (common.Storage, func()) {
			dir, err := os.MkdirTemp(os.TempDir(), "fstree_nosync")
			if err != nil {
				b.Fatalf("creating fstree_nosync root path: %v", err)
			}
			cleanup := func() { os.RemoveAll(dir) }
			return fstree.New(
				fstree.WithPath(dir),
				fstree.WithDepth(2),
				fstree.WithDirNameLen(2),
				fstree.WithNoSync(true),
			), cleanup
		},
	},
	{
		desc: "fstree",
		create: func(b *testing.B) (common.Storage, func()) {
			dir, err := os.MkdirTemp(os.TempDir(), "fstree")
			if err != nil {
				b.Fatalf("creating fstree root path: %v", err)
			}
			cleanup := func() { os.RemoveAll(dir) }
			return fstree.New(
				fstree.WithPath(dir),
				fstree.WithDepth(2),
				fstree.WithDirNameLen(2),
			), cleanup
		},
	},
	{
		desc: "blobovniczatree",
		create: func(b *testing.B) (common.Storage, func()) {
			dir, err := os.MkdirTemp(os.TempDir(), "blobovniczatree")
			if err != nil {
				b.Fatalf("creating blobovniczatree root path: %v", err)
			}
			cleanup := func() { os.RemoveAll(dir) }
			return blobovniczatree.NewBlobovniczaTree(
				blobovniczatree.WithRootPath(dir),
			), cleanup
		},
	},
}

func BenchmarkSubstorageReadPerf(b *testing.B) {
	readTests := []struct {
		desc    string
		size    int
		objGen  func() objectGenerator
		addrGen func() addressGenerator
	}{
		{
			desc:    "seq100",
			size:    10000,
			objGen:  func() objectGenerator { return &seqObjGenerator{objSize: 100} },
			addrGen: func() addressGenerator { return &seqAddrGenerator{maxID: 100} },
		},
		{
			desc:    "rand100",
			size:    10000,
			objGen:  func() objectGenerator { return &seqObjGenerator{objSize: 100} },
			addrGen: func() addressGenerator { return randAddrGenerator(10000) },
		},
	}
	for _, tt := range readTests {
		for _, stEntry := range storages {
			b.Run(fmt.Sprintf("%s-%s", stEntry.desc, tt.desc), func(b *testing.B) {
				objGen := tt.objGen()
				st, cleanup := stEntry.create(b)

				require.NoError(b, st.Open(false))
				require.NoError(b, st.Init())

				// Fill database
				for i := 0; i < tt.size; i++ {
					obj := objGen.Next()
					addr := addressFromObject(obj)
					raw, err := obj.Marshal()
					require.NoError(b, err)
					if _, err := st.Put(common.PutPrm{
						Address: addr,
						RawData: raw,
					}); err != nil {
						b.Fatalf("writing entry: %v", err)
					}
				}

				// Benchmark reading
				addrGen := tt.addrGen()
				b.ResetTimer()
				b.RunParallel(func(pb *testing.PB) {
					for pb.Next() {
						_, err := st.Get(common.GetPrm{Address: addrGen.Next()})
						require.NoError(b, err)
					}
				})

				require.NoError(b, st.Close())
				cleanup()
			})
		}
	}
}

func BenchmarkSubstorageWritePerf(b *testing.B) {
	generators := []struct {
		desc   string
		create func() objectGenerator
	}{
		{desc: "rand10", create: func() objectGenerator { return &randObjGenerator{objSize: 10} }},
		{desc: "rand100", create: func() objectGenerator { return &randObjGenerator{objSize: 100} }},
		{desc: "rand1000", create: func() objectGenerator { return &randObjGenerator{objSize: 1000} }},
		{desc: "overwrite10", create: func() objectGenerator { return &overwriteObjGenerator{objSize: 10, maxObjects: 100} }},
		{desc: "overwrite100", create: func() objectGenerator { return &overwriteObjGenerator{objSize: 100, maxObjects: 100} }},
		{desc: "overwrite1000", create: func() objectGenerator { return &overwriteObjGenerator{objSize: 1000, maxObjects: 100} }},
	}

	for _, genEntry := range generators {
		for _, stEntry := range storages {
			b.Run(fmt.Sprintf("%s-%s", stEntry.desc, genEntry.desc), func(b *testing.B) {
				gen := genEntry.create()
				st, cleanup := stEntry.create(b)

				require.NoError(b, st.Open(false))
				require.NoError(b, st.Init())

				b.ResetTimer()
				b.RunParallel(func(pb *testing.PB) {
					for pb.Next() {
						obj := gen.Next()
						addr := addressFromObject(obj)
						raw, err := obj.Marshal()
						require.NoError(b, err)
						if _, err := st.Put(common.PutPrm{
							Address: addr,
							RawData: raw,
						}); err != nil {
							b.Fatalf("writing entry: %v", err)
						}
					}
				})

				require.NoError(b, st.Close())
				cleanup()
			})
		}
	}
}

func BenchmarkSubstorageIteratePerf(b *testing.B) {
	iterateTests := []struct {
		desc   string
		size   int
		objGen func() objectGenerator
	}{
		{
			desc:   "rand100",
			size:   10000,
			objGen: func() objectGenerator { return &randObjGenerator{objSize: 100} },
		},
	}
	for _, tt := range iterateTests {
		for _, stEntry := range storages {
			b.Run(fmt.Sprintf("%s-%s", stEntry.desc, tt.desc), func(b *testing.B) {
				objGen := tt.objGen()
				st, cleanup := stEntry.create(b)

				require.NoError(b, st.Open(false))
				require.NoError(b, st.Init())

				// Fill database
				for i := 0; i < tt.size; i++ {
					obj := objGen.Next()
					addr := addressFromObject(obj)
					raw, err := obj.Marshal()
					require.NoError(b, err)
					if _, err := st.Put(common.PutPrm{
						Address: addr,
						RawData: raw,
					}); err != nil {
						b.Fatalf("writing entry: %v", err)
					}
				}

				// Benchmark iterate
				cnt := 0
				b.ResetTimer()
				_, err := st.Iterate(common.IteratePrm{
					Handler: func(elem common.IterationElement) error {
						cnt++
						return nil
					},
				})
				require.NoError(b, err)
				require.Equal(b, tt.size, cnt)
				b.StopTimer()

				require.NoError(b, st.Close())
				cleanup()
			})
		}
	}
}

func addressFromObject(obj *objectSDK.Object) oid.Address {
	var addr oid.Address
	if id, isSet := obj.ID(); isSet {
		addr.SetObject(id)
	} else {
		panic("object ID is not set")
	}
	if cid, isSet := obj.ContainerID(); isSet {
		addr.SetContainer(cid)
	} else {
		panic("container ID is not set")
	}
	return addr
}

// addressGenerator is the interface of types that generate object addresses.
type addressGenerator interface {
	Next() oid.Address
}

// seqAddrGenerator is an addressGenerator that generates addresses sequentially and wraps around the given max ID.
type seqAddrGenerator struct {
	cnt   atomic.Uint64
	maxID uint64
}

func (g *seqAddrGenerator) Next() oid.Address {
	var id oid.ID
	binary.LittleEndian.PutUint64(id[:], ((g.cnt.Inc()-1)%g.maxID)+1)
	var addr oid.Address
	addr.SetContainer(cid.ID{})
	addr.SetObject(id)
	return addr
}

func TestSeqAddrGenerator(t *testing.T) {
	gen := &seqAddrGenerator{maxID: 10}
	for i := 1; i <= 20; i++ {
		addr := gen.Next()
		id := addr.Object()

		require.Equal(t, uint64((i-1)%int(gen.maxID)+1), binary.LittleEndian.Uint64(id[:]))
	}
}

// randAddrGenerator is an addressGenerator that generates random addresses in the given range.
type randAddrGenerator uint64

func (g randAddrGenerator) Next() oid.Address {
	var id oid.ID
	binary.LittleEndian.PutUint64(id[:], uint64(1+int(rand.Int63n(int64(g)))))
	var addr oid.Address
	addr.SetContainer(cid.ID{})
	addr.SetObject(id)
	return addr
}

func TestRandAddrGenerator(t *testing.T) {
	gen := randAddrGenerator(5)
	for i := 0; i < 50; i++ {
		addr := gen.Next()
		id := addr.Object()
		k := binary.LittleEndian.Uint64(id[:])

		require.True(t, 1 <= k && k <= uint64(gen))
	}
}

// objectGenerator is the interface of types that generate object entries.
type objectGenerator interface {
	Next() *objectSDK.Object
}

// seqObjGenerator is an objectGenerator that generates entries with random payloads of size objSize and sequential IDs.
type seqObjGenerator struct {
	cnt     atomic.Uint64
	objSize uint64
}

func (g *seqObjGenerator) Next() *objectSDK.Object {
	var id oid.ID
	binary.LittleEndian.PutUint64(id[:], g.cnt.Inc())
	return genObject(id, cid.ID{}, g.objSize)
}

func TestSeqObjGenerator(t *testing.T) {
	gen := &seqObjGenerator{objSize: 10}
	var addrs []string
	for i := 1; i <= 10; i++ {
		obj := gen.Next()
		id, isSet := obj.ID()
		addrs = append(addrs, addressFromObject(obj).EncodeToString())

		require.True(t, isSet)
		require.Equal(t, gen.objSize, uint64(len(obj.Payload())))
		require.Equal(t, uint64(i), binary.LittleEndian.Uint64(id[:]))
	}
	require.True(t, slices.IsSorted(addrs))
}

// randObjGenerator is an objectGenerator that generates entries with random IDs and payloads of size objSize.
type randObjGenerator struct {
	objSize uint64
}

func (g *randObjGenerator) Next() *objectSDK.Object {
	return genObject(oidtest.ID(), cidtest.ID(), g.objSize)
}

func TestRandObjGenerator(t *testing.T) {
	gen := &randObjGenerator{objSize: 10}
	for i := 0; i < 10; i++ {
		obj := gen.Next()

		require.Equal(t, gen.objSize, uint64(len(obj.Payload())))
	}
}

// overwriteObjGenerator is an objectGenerator that generates entries with random payloads of size objSize and at most maxObjects distinct IDs.
type overwriteObjGenerator struct {
	objSize    uint64
	maxObjects uint64
}

func (g *overwriteObjGenerator) Next() *objectSDK.Object {
	var id oid.ID
	binary.LittleEndian.PutUint64(id[:], uint64(1+rand.Int63n(int64(g.maxObjects))))
	return genObject(id, cid.ID{}, g.objSize)
}

func TestOverwriteObjGenerator(t *testing.T) {
	gen := &overwriteObjGenerator{
		objSize:    10,
		maxObjects: 4,
	}
	for i := 0; i < 40; i++ {
		obj := gen.Next()
		id, isSet := obj.ID()
		i := binary.LittleEndian.Uint64(id[:])

		require.True(t, isSet)
		require.Equal(t, gen.objSize, uint64(len(obj.Payload())))
		require.True(t, 1 <= i && i <= gen.maxObjects)
	}
}

// Generates an object with random payload and the given address and size.
// TODO(#86): there's some testing-related dupes in many places. Probably worth
// spending some time cleaning up a bit.
func genObject(id oid.ID, cid cid.ID, sz uint64) *objectSDK.Object {
	raw := objectSDK.New()

	raw.SetID(id)
	raw.SetContainerID(cid)

	payload := make([]byte, sz)
	rand.Read(payload)
	raw.SetPayload(payload)

	return raw
}