package tree

import (
	"context"
	"testing"
	"time"

	"git.frostfs.info/TrueCloudLab/frostfs-node/pkg/local_object_storage/pilorama"
	cidSDK "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id"
	cidtest "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/container/id/test"
	"github.com/google/uuid"
	"google.golang.org/grpc"
)

func BenchmarkGetSubTree(b *testing.B) {
	const count = simpleHeapFallbackThreshold

	d := pilorama.CIDDescriptor{CID: cidtest.ID(), Size: 1}
	treeID := "sometree"

	body := &GetSubTreeRequest_Body{
		TreeId: treeID,
		RootId: pilorama.RootID,
		Depth:  2,
		OrderBy: &GetSubTreeRequest_Body_Order{
			Direction: GetSubTreeRequest_Body_Order_Asc,
		},
	}
	p := newBenchPilorama(count)

	b.Run("latency", func(b *testing.B) {
		b.ReportAllocs()
		for i := 0; i < b.N; i++ {
			acc := benchServer{start: time.Now(), errIndex: count + 1}
			err := getSubTree(context.Background(), &acc, d.CID, body, p)
			if err != nil {
				b.Fatalf("Error: %v, expected: %d, got %d", err, count, acc.seen)
			}
			b.ReportMetric(float64(uint64(acc.first)/uint64(b.N)), "ns/op")
		}
	})
	b.Run("total time", func(b *testing.B) {
		b.ReportAllocs()
		for i := 0; i < b.N; i++ {
			acc := benchServer{start: time.Now(), errIndex: count + 1}
			err := getSubTree(context.Background(), &acc, d.CID, body, p)
			if err != nil {
				b.Fatalf("Error: %v, expected: %d, got %d", err, count, acc.seen)
			}
		}
	})
}

// benchPilorama represents flat pilorama, i.e. tree of height 2 with many child nodes.
type benchPilorama struct {
	pilorama.Forest // Satisfy interface.
	nodes           []pilorama.NodeInfo
}

func newBenchPilorama(size int) *benchPilorama {
	nodes := make([]pilorama.NodeInfo, 0, size)

	for i := 1; i <= size; i++ { // Start with 1 to avoid intersecting with RootID = 0.
		nodes = append(nodes, pilorama.NodeInfo{
			ParentID: pilorama.RootID,
			ID:       pilorama.Node(i),
			Meta: pilorama.Meta{
				Items: []pilorama.KeyValue{{
					Key:   pilorama.AttributeFilename,
					Value: []byte(uuid.New().String()),
				}},
			},
		})
	}
	return &benchPilorama{nodes: nodes}
}

func (p *benchPilorama) TreeGetChildren(ctx context.Context, cid cidSDK.ID, treeID string, nodeID pilorama.Node) ([]pilorama.NodeInfo, error) {
	switch nodeID {
	case pilorama.RootID:
		result := make([]pilorama.NodeInfo, len(p.nodes))
		copy(result, p.nodes)
		return result, nil
	default:
		panic("unexpected")
	}
}

func (p *benchPilorama) TreeGetMeta(ctx context.Context, cid cidSDK.ID, treeID string, root pilorama.Node) (pilorama.Meta, pilorama.Node, error) {
	if root == pilorama.RootID {
		return pilorama.Meta{}, pilorama.RootID, nil
	}
	return p.nodes[root-1].Meta, p.nodes[root-1].ParentID, nil
}

type benchServer struct {
	grpc.ServerStream // to satisfy the interface
	start             time.Time
	first             time.Duration
	seen              int
	errIndex          int
}

var _ TreeService_GetSubTreeServer = (*benchServer)(nil)

func (s *benchServer) Send(r *GetSubTreeResponse) error {
	if s.seen == 1 {
		s.first = time.Since(s.start)
	}
	s.seen++

	if s.errIndex >= 0 {
		if s.seen == s.errIndex+1 {
			return errSubTreeSend
		}
		if s.errIndex >= 0 && s.seen > s.errIndex {
			return errSubTreeSendAfterError
		}
	}
	return nil
}