package grpctool

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/grpctool/test"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_otel"
	"go.opentelemetry.io/otel/attribute"
	otelmetric "go.opentelemetry.io/otel/metric"
	metricnoop "go.opentelemetry.io/otel/metric/noop"
	"go.uber.org/mock/gomock"
	"google.golang.org/grpc"
	"google.golang.org/grpc/stats"
)

func TestRequestsInFlightStatsHandler(t *testing.T) {
	ctrl := gomock.NewController(t)
	meter := mock_otel.NewMockMeter(ctrl)
	sc := mock_otel.NewMockInt64UpDownCounter(ctrl)
	cc := mock_otel.NewMockInt64UpDownCounter(ctrl)
	meter.EXPECT().
		Int64UpDownCounter("grpc_server_requests_in_flight", gomock.Any()).
		Return(sc, nil)
	meter.EXPECT().
		Int64UpDownCounter("grpc_client_requests_in_flight", gomock.Any()).
		Return(cc, nil)

	ssh, err := NewServerRequestsInFlightStatsHandler(meter)
	require.NoError(t, err)
	csh, err := NewClientRequestsInFlightStatsHandler(meter)
	require.NoError(t, err)

	expectedAttrs := attribute.NewSet(
		grpcServiceAttr.String("gitlab.agent.grpctool.test.Testing"),
		grpcMethodAttr.String("RequestResponse"),
	)
	assertAttrs := func(ctx context.Context, _ int64, opts ...otelmetric.AddOption) {
		ocfg := otelmetric.NewAddConfig(opts)
		attrs := ocfg.Attributes()
		assert.True(t, attrs.Equals(&expectedAttrs))
	}
	ats := &test.GRPCTestingServer{
		UnaryFunc: func(ctx context.Context, request *test.Request) (*test.Response, error) {
			return &test.Response{}, nil
		},
	}
	conn := setupWithOpts(t, ats,
		[]grpc.ServerOption{grpc.StatsHandler(ssh)},
		[]grpc.DialOption{grpc.WithStatsHandler(csh)},
	)
	client := test.NewTestingClient(conn)

	gomock.InOrder(
		cc.EXPECT().
			Add(gomock.Any(), int64(1), gomock.Any()).
			Do(assertAttrs),
		cc.EXPECT().
			Add(gomock.Any(), int64(-1), gomock.Any()).
			Do(assertAttrs),
	)

	gomock.InOrder(
		sc.EXPECT().
			Add(gomock.Any(), int64(1), gomock.Any()).
			Do(assertAttrs),
		sc.EXPECT().
			Add(gomock.Any(), int64(-1), gomock.Any()).
			Do(assertAttrs),
	)

	_, err = client.RequestResponse(context.Background(), &test.Request{})
	require.NoError(t, err)
}

func Benchmark_RequestsInFlightStatsHandler_HandleRPC(b *testing.B) {
	h, err := NewClientRequestsInFlightStatsHandler(metricnoop.NewMeterProvider().Meter("test"))
	b.ReportAllocs()
	require.NoError(b, err)
	begin := &stats.Begin{}
	info := &stats.RPCTagInfo{
		FullMethodName: "/service1/rpc1",
	}
	bg := context.Background()

	for b.Loop() {
		ctx := h.TagRPC(bg, info)
		h.HandleRPC(ctx, begin)
	}
}
