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/module/runner_controller/rpc"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/tool/grpcz"
	"gitlab.com/gitlab-org/cluster-integration/tunnel/tool/retry"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type server struct {
	rpc.UnsafeRunnerControllerServiceServer

	tunnelHandler          router.Handler
	getAgentInfoPollConfig retry.PollConfig
}

func (s *server) AdmitJob(server grpc.BidiStreamingServer[rpc.AdmitJobResponse, rpc.AdmitJobRequest]) error {
	ctx := server.Context()
	ageCtx := grpcz.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(ctx, ageCtx, agentInfo.AgentKey(), func(tunnelRetErr chan<- error) (router.Forwarder, error) {
			return router.NewSimpleRPCForwarder(server, tunnelRetErr), nil
		}), retry.Done
	})
}
