package tunserver

import (
	"context"
	"io"
	"log/slog"
	"net"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/grpctool"
	"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/matcher"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_modshared"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_tool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testhelpers"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testlogger"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/tlstool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tunnel/rpc"
	"go.uber.org/mock/gomock"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"
	"k8s.io/apimachinery/pkg/util/wait"
)

var (
	_ grpc.StreamHandler = (*Router[grpctool.URLTarget])(nil).routeToTunclient
	_ DataCallback       = (*wrappingCallback)(nil)
)

func TestRouter_UnaryHappyPath(t *testing.T) {
	ctrl := gomock.NewController(t)
	unaryResponse := &test.Response{Message: &test.Response_Scalar{Scalar: 123}}
	payloadMD, responseMD, trailersMD := meta()
	payloadReq := &test.Request{S1: "123"}
	var (
		headerResp  metadata.MD
		trailerResp metadata.MD
	)
	tun := NewMockTunnel(ctrl)
	tun.EXPECT().
		ForwardStream(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
		Do(forwardStream(t, payloadMD, payloadReq, unaryResponse, responseMD, trailersMD))
	runRouterTest(t, tun, func(client test.TestingClient) {
		ctx := metadata.NewOutgoingContext(context.Background(), payloadMD)
		// grpc.Header() and grpc.Trailer are ok here because it's a unary RPC.
		response, err := client.RequestResponse(ctx, payloadReq, grpc.Header(&headerResp), grpc.Trailer(&trailerResp)) //nolint: forbidigo
		require.NoError(t, err)
		matcher.AssertProtoEqual(t, response, unaryResponse)
		mdContains(t, responseMD, headerResp)
		mdContains(t, trailersMD, trailerResp)
	})
}

func TestRouter_UnaryImmediateError(t *testing.T) {
	ctrl := gomock.NewController(t)
	statusWithDetails, err := status.New(codes.InvalidArgument, "some expected error").
		WithDetails(&test.Request{S1: "some details of the error"})
	require.NoError(t, err)
	tun := NewMockTunnel(ctrl)
	tun.EXPECT().
		ForwardStream(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
		Return(statusWithDetails.Err())
	runRouterTest(t, tun, func(client test.TestingClient) {
		_, err = client.RequestResponse(context.Background(), &test.Request{S1: "123"})
		require.Error(t, err)
		receivedStatus := status.Convert(err).Proto()
		matcher.AssertProtoEqual(t, receivedStatus, statusWithDetails.Proto())
	})
}

func TestRouter_UnaryErrorAfterHeader(t *testing.T) {
	ctrl := gomock.NewController(t)
	payloadMD, responseMD, trailersMD := meta()
	statusWithDetails, err := status.New(codes.InvalidArgument, "some expected error").
		WithDetails(&test.Request{S1: "some details of the error"})
	require.NoError(t, err)
	var (
		headerResp  metadata.MD
		trailerResp metadata.MD
	)
	tun := NewMockTunnel(ctrl)
	tun.EXPECT().
		ForwardStream(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
		DoAndReturn(func(log *slog.Logger, rpcAPI modshared.RPCAPI, incomingStream grpc.ServerStream, cb DataCallback) error {
			verifyMeta(t, incomingStream, payloadMD)
			assert.NoError(t, cb.Header(rpc.MetaToMetadataKV(responseMD)))
			assert.NoError(t, cb.Trailer(rpc.MetaToMetadataKV(trailersMD)))
			return statusWithDetails.Err()
		})
	runRouterTest(t, tun, func(client test.TestingClient) {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		ctx = metadata.NewOutgoingContext(ctx, payloadMD)
		// grpc.Header() and grpc.Trailer are ok here because it's a unary RPC.
		_, err := client.RequestResponse(ctx, &test.Request{S1: "123"}, grpc.Header(&headerResp), grpc.Trailer(&trailerResp)) //nolint: forbidigo
		require.Error(t, err)
		receivedStatus := status.Convert(err).Proto()
		matcher.AssertProtoEqual(t, receivedStatus, statusWithDetails.Proto())
		mdContains(t, responseMD, headerResp)
		mdContains(t, trailersMD, trailerResp)
	})
}

func TestRouter_StreamHappyPath(t *testing.T) {
	ctrl := gomock.NewController(t)
	streamResponse := &test.Response{Message: &test.Response_Scalar{Scalar: 123}}
	payloadMD, responseMD, trailersMD := meta()
	payloadReq := &test.Request{S1: "123"}
	tun := NewMockTunnel(ctrl)
	tun.EXPECT().
		ForwardStream(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
		Do(forwardStream(t, payloadMD, payloadReq, streamResponse, responseMD, trailersMD))
	runRouterTest(t, tun, func(client test.TestingClient) {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		ctx = metadata.NewOutgoingContext(ctx, payloadMD)
		stream, err := client.StreamingRequestResponse(ctx)
		require.NoError(t, err)
		err = stream.Send(payloadReq)
		require.NoError(t, err)
		err = stream.CloseSend()
		require.NoError(t, err)
		response, err := stream.Recv()
		require.NoError(t, err)
		matcher.AssertProtoEqual(t, response, streamResponse)
		_, err = stream.Recv()
		assert.Equal(t, io.EOF, err)
		verifyHeaderAndTrailer(t, stream, responseMD, trailersMD)
	})
}

func TestRouter_StreamImmediateError(t *testing.T) {
	ctrl := gomock.NewController(t)
	statusWithDetails, err := status.New(codes.InvalidArgument, "some expected error").
		WithDetails(&test.Request{S1: "some details of the error"})
	require.NoError(t, err)
	tun := NewMockTunnel(ctrl)
	tun.EXPECT().
		ForwardStream(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
		Return(statusWithDetails.Err())
	runRouterTest(t, tun, func(client test.TestingClient) {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		stream, err := client.StreamingRequestResponse(ctx)
		require.NoError(t, err)
		_, err = stream.Recv()
		require.Error(t, err)
		receivedStatus := status.Convert(err).Proto()
		matcher.AssertProtoEqual(t, receivedStatus, statusWithDetails.Proto())
	})
}

func TestRouter_StreamErrorAfterHeader(t *testing.T) {
	ctrl := gomock.NewController(t)
	payloadMD, responseMD, trailersMD := meta()
	statusWithDetails, err := status.New(codes.InvalidArgument, "some expected error").
		WithDetails(&test.Request{S1: "some details of the error"})
	require.NoError(t, err)
	tun := NewMockTunnel(ctrl)
	tun.EXPECT().
		ForwardStream(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
		DoAndReturn(func(log *slog.Logger, rpcAPI modshared.RPCAPI, incomingStream grpc.ServerStream, cb DataCallback) error {
			verifyMeta(t, incomingStream, payloadMD)
			assert.NoError(t, cb.Header(rpc.MetaToMetadataKV(responseMD)))
			assert.NoError(t, cb.Trailer(rpc.MetaToMetadataKV(trailersMD)))
			return statusWithDetails.Err()
		})
	runRouterTest(t, tun, func(client test.TestingClient) {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		ctx = metadata.NewOutgoingContext(ctx, payloadMD)
		stream, err := client.StreamingRequestResponse(ctx)
		require.NoError(t, err)
		_, err = stream.Recv()
		require.Error(t, err)
		receivedStatus := status.Convert(err).Proto()
		matcher.AssertProtoEqual(t, receivedStatus, statusWithDetails.Proto())
		verifyHeaderAndTrailer(t, stream, responseMD, trailersMD)
	})
}

func TestRouter_StreamVisitorErrorAfterErrorMessage(t *testing.T) {
	ctrl := gomock.NewController(t)
	payloadMD, responseMD, trailersMD := meta()
	statusWithDetails, err := status.New(codes.InvalidArgument, "some expected error").
		WithDetails(&test.Request{S1: "some details of the error"})
	require.NoError(t, err)
	tun := NewMockTunnel(ctrl)
	tun.EXPECT().
		ForwardStream(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
		DoAndReturn(func(log *slog.Logger, rpcAPI modshared.RPCAPI, incomingStream grpc.ServerStream, cb DataCallback) error {
			verifyMeta(t, incomingStream, payloadMD)
			assert.NoError(t, cb.Header(rpc.MetaToMetadataKV(responseMD)))
			assert.NoError(t, cb.Trailer(rpc.MetaToMetadataKV(trailersMD)))
			assert.NoError(t, cb.Error(statusWithDetails.Proto()))
			return status.Error(codes.Unavailable, "expected return error")
		})
	runRouterTest(t, tun, func(client test.TestingClient) {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()
		ctx = metadata.NewOutgoingContext(ctx, payloadMD)
		stream, err := client.StreamingRequestResponse(ctx)
		require.NoError(t, err)
		_, err = stream.Recv()
		require.EqualError(t, err, "rpc error: code = Unavailable desc = expected return error")
		verifyHeaderAndTrailer(t, stream, responseMD, trailersMD)
	})
}

func meta() (metadata.MD, metadata.MD, metadata.MD) {
	payloadMD := metadata.Pairs("key1", "value1")
	responseMD := metadata.Pairs("key2", "value2")
	trailersMD := metadata.Pairs("key3", "value3")
	return payloadMD, responseMD, trailersMD
}

func verifyHeaderAndTrailer(t *testing.T, stream grpc.ClientStream, responseMD, trailersMD metadata.MD) {
	headerResp, err := stream.Header()
	require.NoError(t, err)
	mdContains(t, responseMD, headerResp)
	mdContains(t, trailersMD, stream.Trailer())
}

func forwardStream(t *testing.T, payloadMD metadata.MD, payloadReq *test.Request, response *test.Response, responseMD, trailersMD metadata.MD) func(*slog.Logger, modshared.RPCAPI, grpc.ServerStream, DataCallback) error {
	return func(log *slog.Logger, rpcAPI modshared.RPCAPI, incomingStream grpc.ServerStream, cb DataCallback) error {
		verifyMeta(t, incomingStream, payloadMD)
		var req test.Request
		err := incomingStream.RecvMsg(&req)
		assert.NoError(t, err)
		matcher.AssertProtoEqual(t, payloadReq, &req)
		data, err := proto.Marshal(response)
		assert.NoError(t, err)
		assert.NoError(t, cb.Header(rpc.MetaToMetadataKV(responseMD)))
		assert.NoError(t, cb.Message(data))
		assert.NoError(t, cb.Trailer(rpc.MetaToMetadataKV(trailersMD)))
		return nil
	}
}

func verifyMeta(t *testing.T, incomingStream grpc.ServerStream, payloadMD metadata.MD) {
	md, _ := metadata.FromIncomingContext(incomingStream.Context())
	mdContains(t, payloadMD, md)
}

func mdContains(t *testing.T, expectedMd, actualMd metadata.MD) {
	for k, v := range expectedMd {
		assert.Equalf(t, v, actualMd[k], "key: %s", k)
	}
}

// test:client(default codec) --> kas:internal server(raw codec) --> router_kas handler -->
// client from kas_pool(raw with fallback codec) --> kas:private server(raw with fallback codec) -->
// router_agent handler --> tunnel finder --> tunnel.ForwardStream()
func runRouterTest(t *testing.T, tunnel *MockTunnel, runTest func(client test.TestingClient)) {
	ctrl := gomock.NewController(t)
	rep := mock_tool.NewMockErrReporter(ctrl)
	log := testlogger.New(t)
	querier := NewMockPollingGatewayURLQuerier[grpctool.URLTarget](ctrl)
	modAPI := mock_modshared.NewMockAPI(ctrl)
	fh := NewMockFindHandle(ctrl)
	privateAPIServerListener, err := net.Listen("tcp", "localhost:0")
	require.NoError(t, err)
	defer privateAPIServerListener.Close()

	querier.EXPECT().
		CachedGatewayURLs(testhelpers.AgentkKey1)
	factory := func(ctx context.Context, fullMethodName string) modshared.RPCAPI {
		rpcAPI := mock_modshared.NewMockRPCAPI(ctrl)
		return rpcAPI
	}

	privateAPIServer := grpc.NewServer(
		grpc.StatsHandler(grpctool.NewServerMaxConnAgeStatsHandler(context.Background(), 0)),
		grpc.ChainStreamInterceptor(
			modshared.StreamRPCAPIInterceptor(factory),
		),
		grpc.ChainUnaryInterceptor(
			modshared.UnaryRPCAPIInterceptor(factory),
		),
		grpc.ForceServerCodecV2(grpctool.RawCodecWithProtoFallback{}),
	)

	kasPool := grpctool.NewPool(log, rep, grpctool.DefaultNewConnection(credentials.NewTLS(tlstool.ClientConfig())))
	defer func() {
		kasPool.Shutdown(time.Minute)
	}()
	rp := NewMockRouterPlugin[grpctool.URLTarget](ctrl)
	rp.EXPECT().
		FindReadyGateway(gomock.Any(), gomock.Any(), gomock.Any()).
		DoAndReturn(func(ctx context.Context, log *slog.Logger, method string) (ReadyGateway[grpctool.URLTarget], *slog.Logger, api.AgentKey, error) {
			gf := NewGatewayFinder(
				ctx,
				log,
				kasPool,
				querier,
				modAPI,
				method,
				grpctool.URLTarget("grpc://"+privateAPIServerListener.Addr().String()),
				testhelpers.AgentkKey1,
				testhelpers.NewPollConfig(time.Minute),
				// We don't want any nondeterministic polls to other KAS
				5*time.Second,
			)
			rg, err := gf.Find(ctx)
			return rg, log, testhelpers.AgentkKey1, status.FromContextError(err).Err() // must return status-compatible error
		})
	gomock.InOrder(
		rp.EXPECT().
			FindTunnel(gomock.Any(), gomock.Any()).
			Return(true, log, fh, nil),
		fh.EXPECT().
			Get(gomock.Any()).
			Return(tunnel, nil),
		rp.EXPECT().
			PrepareStreamForForwarding(gomock.Any()).
			DoAndReturn(func(stream grpc.ServerStream) (grpc.ServerStream, error) {
				return stream, nil
			}),
		tunnel.EXPECT().Done(gomock.Any()),
		fh.EXPECT().Done(gomock.Any()),
	)

	r := &Router[grpctool.URLTarget]{
		Plugin:           rp,
		PrivateAPIServer: privateAPIServer,
	}
	r.RegisterTunclientAPI(&test.Testing_ServiceDesc)
	var wg wait.Group
	defer wg.Wait()
	defer privateAPIServer.GracefulStop()
	wg.Start(func() {
		assert.NoError(t, privateAPIServer.Serve(privateAPIServerListener))
	})
	routingConn := &RoutingClientConn[grpctool.URLTarget]{
		Log:    log,
		API:    modAPI,
		Plugin: rp,
	}
	client := test.NewTestingClient(routingConn)
	runTest(client)
}
