package server

import (
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent2kas_tunnel/router"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modserver"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/grpctool"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/retry"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tunnel/rpc"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type server struct {
	rpc.UnsafeReverseTunnelServer
	tunnelHandler          router.Handler
	getAgentInfoPollConfig retry.PollConfigFactory
}

func (s *server) Connect(server grpc.BidiStreamingServer[rpc.ConnectRequest, rpc.ConnectResponse]) error {
	ctx := server.Context()
	ageCtx := grpctool.MaxConnectionAgeContextFromStreamContext(ctx)
	rpcAPI := modserver.AgentRPCAPIFromContext(ctx)
	log := rpcAPI.Log()
	return rpcAPI.PollWithBackoff(s.getAgentInfoPollConfig(), func() (error, retry.AttemptResult) { //nolint:staticcheck
		agentInfo, err := rpcAPI.AgentInfo(ctx, log)
		if err != nil {
			if status.Code(err) == codes.Unavailable {
				return nil, retry.Backoff
			}
			return err, retry.Done // no wrap
		}
		return s.tunnelHandler.HandleTunnel(ageCtx, agentInfo, server), retry.Done
	})
}
