package engine

import (
	"context"
	"sync"
	"testing"

	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/core/object"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/blobstor/common"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/blobstor/teststore"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/internal/testutil"
	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/shard"
	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"
	"github.com/stretchr/testify/require"
)

func TestRebalance(t *testing.T) {
	t.Parallel()

	te := newEngineWithErrorThreshold(t, "", 0)
	defer func() {
		require.NoError(t, te.ng.Close(context.Background()))
	}()

	const (
		objCount  = 20
		copyCount = (objCount + 2) / 3
	)

	type objectWithShard struct {
		bestShard  shard.ID
		worstShard shard.ID
		object     *objectSDK.Object
	}

	objects := make([]objectWithShard, objCount)
	for i := range objects {
		obj := testutil.GenerateObjectWithCID(cidtest.ID())
		obj.SetPayload(make([]byte, errSmallSize))
		objects[i].object = obj

		shards := te.ng.sortShardsByWeight(object.AddressOf(obj))
		objects[i].bestShard = *shards[0].Shard.ID()
		objects[i].worstShard = *shards[1].Shard.ID()
	}

	for i := range objects {
		var prm shard.PutPrm
		prm.SetObject(objects[i].object)

		var err1, err2 error
		te.ng.mtx.RLock()
		// Every 3rd object (i%3 == 0) is put to both shards, others are distributed.
		if i%3 != 1 {
			_, err1 = te.ng.shards[te.shards[0].id.String()].Shard.Put(context.Background(), prm)
		}
		if i%3 != 2 {
			_, err2 = te.ng.shards[te.shards[1].id.String()].Shard.Put(context.Background(), prm)
		}
		te.ng.mtx.RUnlock()

		require.NoError(t, err1)
		require.NoError(t, err2)
	}

	var removedMtx sync.Mutex
	var removed []deleteEvent
	for _, shard := range te.shards {
		id := *shard.id
		shard.largeFileStorage.SetOption(teststore.WithDelete(func(prm common.DeletePrm) (common.DeleteRes, error) {
			removedMtx.Lock()
			removed = append(removed, deleteEvent{shardID: id, addr: prm.Address})
			removedMtx.Unlock()
			return common.DeleteRes{}, nil
		}))
	}

	err := te.ng.RemoveDuplicates(context.Background(), RemoveDuplicatesPrm{})
	require.NoError(t, err)

	require.Equal(t, copyCount, len(removed))

	removedMask := make([]bool, len(objects))
loop:
	for i := range removed {
		for j := range objects {
			if removed[i].addr == object.AddressOf(objects[j].object) {
				require.Equal(t, objects[j].worstShard, removed[i].shardID,
					"object %d was expected to be removed from another shard", j)
				removedMask[j] = true
				continue loop
			}
		}
		require.FailNow(t, "unexpected object was removed", removed[i].addr)
	}

	for i := 0; i < copyCount; i++ {
		if i%3 == 0 {
			require.True(t, removedMask[i], "object %d was expected to be removed", i)
		} else {
			require.False(t, removedMask[i], "object %d was not expected to be removed", i)
		}
	}
}

func TestRebalanceSingleThread(t *testing.T) {
	t.Parallel()

	te := newEngineWithErrorThreshold(t, "", 0)
	defer func() {
		require.NoError(t, te.ng.Close(context.Background()))
	}()

	obj := testutil.GenerateObjectWithCID(cidtest.ID())
	obj.SetPayload(make([]byte, errSmallSize))

	var prm shard.PutPrm
	prm.SetObject(obj)
	te.ng.mtx.RLock()
	_, err1 := te.ng.shards[te.shards[0].id.String()].Shard.Put(context.Background(), prm)
	_, err2 := te.ng.shards[te.shards[1].id.String()].Shard.Put(context.Background(), prm)
	te.ng.mtx.RUnlock()
	require.NoError(t, err1)
	require.NoError(t, err2)

	signal := make(chan struct{})  // unblock rebalance
	started := make(chan struct{}) // make sure rebalance is started
	for _, shard := range te.shards {
		shard.largeFileStorage.SetOption(teststore.WithDelete(func(common.DeletePrm) (common.DeleteRes, error) {
			close(started)
			<-signal
			return common.DeleteRes{}, nil
		}))
	}

	var firstErr error
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		firstErr = te.ng.RemoveDuplicates(context.Background(), RemoveDuplicatesPrm{})
	}()

	<-started
	secondErr := te.ng.RemoveDuplicates(context.Background(), RemoveDuplicatesPrm{})
	require.ErrorIs(t, secondErr, errRemoveDuplicatesInProgress)

	close(signal)
	wg.Wait()
	require.NoError(t, firstErr)
}

type deleteEvent struct {
	shardID shard.ID
	addr    oid.Address
}

func TestRebalanceExitByContext(t *testing.T) {
	te := newEngineWithErrorThreshold(t, "", 0)
	defer func() {
		require.NoError(t, te.ng.Close(context.Background()))
	}()

	objects := make([]*objectSDK.Object, 4)
	for i := range objects {
		obj := testutil.GenerateObjectWithCID(cidtest.ID())
		obj.SetPayload(make([]byte, errSmallSize))
		objects[i] = obj
	}

	for i := range objects {
		var prm shard.PutPrm
		prm.SetObject(objects[i])

		te.ng.mtx.RLock()
		_, err1 := te.ng.shards[te.shards[0].id.String()].Shard.Put(context.Background(), prm)
		_, err2 := te.ng.shards[te.shards[1].id.String()].Shard.Put(context.Background(), prm)
		te.ng.mtx.RUnlock()

		require.NoError(t, err1)
		require.NoError(t, err2)
	}

	var removed []deleteEvent
	deleteCh := make(chan struct{})
	signal := make(chan struct{})
	for _, shard := range te.shards {
		id := *shard.id
		shard.largeFileStorage.SetOption(teststore.WithDelete(func(prm common.DeletePrm) (common.DeleteRes, error) {
			deleteCh <- struct{}{}
			<-signal
			removed = append(removed, deleteEvent{shardID: id, addr: prm.Address})
			return common.DeleteRes{}, nil
		}))
	}

	ctx, cancel := context.WithCancel(context.Background())

	var rebalanceErr error
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		rebalanceErr = te.ng.RemoveDuplicates(ctx, RemoveDuplicatesPrm{Concurrency: 1})
	}()

	const removeCount = 3
	for i := 0; i < removeCount-1; i++ {
		<-deleteCh
		signal <- struct{}{}
	}
	<-deleteCh
	cancel()
	close(signal)

	wg.Wait()
	require.ErrorIs(t, rebalanceErr, context.Canceled)
	require.Equal(t, removeCount, len(removed))
}