package testutil

import (
	"encoding/binary"

	cid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	"git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object"
	oid "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/object/id"
	"go.uber.org/atomic"
	"golang.org/x/exp/rand"
)

// 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
}

var _ AddressGenerator = &SeqAddrGenerator{}

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
}

// 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
}

// ObjectGenerator is the interface of types that generate object entries.
type ObjectGenerator interface {
	Next() *object.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
}

var _ ObjectGenerator = &SeqObjGenerator{}

func generateObjectWithOIDWithCIDWithSize(oid oid.ID, cid cid.ID, sz uint64) *object.Object {
	data := make([]byte, sz)
	_, _ = rand.Read(data)
	obj := GenerateObjectWithCIDWithPayload(cid, data)
	obj.SetID(oid)
	return obj
}

func (g *SeqObjGenerator) Next() *object.Object {
	var id oid.ID
	binary.LittleEndian.PutUint64(id[:], g.cnt.Inc())
	return generateObjectWithOIDWithCIDWithSize(id, cid.ID{}, g.ObjSize)
}

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

var _ ObjectGenerator = &RandObjGenerator{}

func (g *RandObjGenerator) Next() *object.Object {
	return generateObjectWithOIDWithCIDWithSize(oid.ID{}, cid.ID{}, g.ObjSize)
}

// 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() *object.Object {
	var id oid.ID
	binary.LittleEndian.PutUint64(id[:], uint64(1+rand.Int63n(int64(g.MaxObjects))))
	return generateObjectWithOIDWithCIDWithSize(id, cid.ID{}, g.ObjSize)
}

func AddressFromObject(obj *object.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
}