package server

import (
	"errors"
	"net/http"
	"strings"
	"testing"

	"buf.build/go/protovalidate"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	gapi "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/gitlab/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/job_router/rpc"
	runner_controller_rpc "gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/runner_controller/rpc"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/server_api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/httpz"
	"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_gitlab"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_modserver"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_otel"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testhelpers"
	"go.uber.org/mock/gomock"
)

const (
	jobResponse = `{
		"id": 123,
		"token": "job-token-123",
		"job_info": {
			"project_id": 100,
			"project_name": "test-project",
			"project_full_path": "group/test-project",
			"namespace_id": 10,
			"root_namespace_id": 1,
			"organization_id": 1,
			"user_id": 42
		}
	}`
)

func TestRequestID(t *testing.T) {
	assert.Equal(t, requestIDMetadataKey, strings.ToLower(httpz.RequestIDHeader))
}

func TestGetJob_NoContent(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockRPCAPI := mock_modserver.NewMockAgentRPCAPI(ctrl)
	mockAdmissionControlCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControlRequestsCounter := mock_otel.NewMockInt64Counter(ctrl)

	mockRPCAPI.EXPECT().AgentInfo(gomock.Any(), gomock.Any()).Return(&server_api.RunnerInfo{
		Key: testhelpers.RunnerKey1,
	}, nil)

	gitlabClient := mock_gitlab.SetupClient(t, "", gapi.JobsRequestAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPost, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.WriteHeader(http.StatusNoContent)
	})
	v, err := protovalidate.New()
	require.NoError(t, err)

	s := &server{
		gitlabClient:                         gitlabClient,
		validator:                            v,
		admissionControlTotalCounter:         mockAdmissionControlCounter,
		admissionControlRequestsTotalCounter: mockAdmissionControlRequestsCounter,
	}

	ctx := mock_modserver.IncomingAgentCtx(t, mockRPCAPI)

	req := &rpc.GetJobRequest{
		JobRequest: []byte(`{}`),
	}
	resp, err := s.GetJob(ctx, req)

	require.NoError(t, err)
	assert.Empty(t, resp.JobResponse)
}

func TestGetJob_WithJob_RunnerControllerAdmits(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockRPCAPI := mock_modserver.NewMockAgentRPCAPI(ctrl)
	mockAdmissionControlCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControlRequestsCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControllerClient := NewMockadmissionControllerClient(ctrl)

	mockRPCAPI.EXPECT().AgentInfo(gomock.Any(), gomock.Any()).Return(&server_api.RunnerInfo{
		Key: testhelpers.RunnerKey1,
	}, nil)

	mux := http.NewServeMux()
	mux.HandleFunc(gapi.JobsRequestAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPost, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusCreated)
		_, err := w.Write([]byte(jobResponse))
		assert.NoError(t, err)
	})
	mux.HandleFunc(gapi.GetRunnerControllersForJobAdmissionAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodGet, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, err := w.Write([]byte(`{"runner_controllers":[{"id":1111,"state":"enabled"}]}`))
		assert.NoError(t, err)
	})

	gitlabClient := mock_gitlab.SetupClient(t, "", "/", mux.ServeHTTP)
	v, err := protovalidate.New()
	require.NoError(t, err)

	// Mock the runner controller client stream to admit the job
	gomock.InOrder(
		mockAdmissionControllerClient.EXPECT().AdmitJob(
			gomock.Any(),
			api.AgentKey{ID: 1111, Type: api.AgentTypeRunnerController},
			&runner_controller_rpc.AdmitJobRequest{
				ProjectId:       100,
				ProjectName:     "test-project",
				ProjectFullPath: "group/test-project",
				NamespaceId:     10,
				RootNamespaceId: 1,
				OrganizationId:  1,
				UserId:          42,
				Image:           &runner_controller_rpc.Image{},
				RunnerId:        testhelpers.RunnerKey1.ID,
			}).
			Return(&runner_controller_rpc.AdmitJobResponse{
				AdmissionResponse: &runner_controller_rpc.AdmitJobResponse_Admitted{Admitted: &runner_controller_rpc.Admitted{}},
			}, nil),

		mockAdmissionControlRequestsCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "admitted"})),
		mockAdmissionControlCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "admitted"})),
	)

	s := &server{
		gitlabClient:                         gitlabClient,
		validator:                            v,
		newAdmissionControllerClient:         func() admissionControllerClient { return mockAdmissionControllerClient },
		admissionControlTotalCounter:         mockAdmissionControlCounter,
		admissionControlRequestsTotalCounter: mockAdmissionControlRequestsCounter,
	}

	ctx := mock_modserver.IncomingAgentCtx(t, mockRPCAPI)
	req := &rpc.GetJobRequest{
		JobRequest: []byte(`{}`),
	}

	resp, err := s.GetJob(ctx, req)

	require.NoError(t, err)
	assert.JSONEq(t, jobResponse, string(resp.JobResponse))
}

func TestGetJob_WithJob_MultipleRunnerControllersAdmit(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockRPCAPI := mock_modserver.NewMockAgentRPCAPI(ctrl)
	mockAdmissionControlCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControlRequestsCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControllerClient := NewMockadmissionControllerClient(ctrl)

	mockRPCAPI.EXPECT().AgentInfo(gomock.Any(), gomock.Any()).Return(&server_api.RunnerInfo{
		Key: testhelpers.RunnerKey1,
	}, nil)

	mux := http.NewServeMux()
	mux.HandleFunc(gapi.JobsRequestAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPost, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusCreated)
		_, err := w.Write([]byte(jobResponse))
		assert.NoError(t, err)
	})
	mux.HandleFunc(gapi.GetRunnerControllersForJobAdmissionAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodGet, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, err := w.Write([]byte(`{"runner_controllers":[{"id":2001,"state":"enabled"},{"id":2002,"state":"enabled"},{"id":2003,"state":"enabled"}]}`))
		assert.NoError(t, err)
	})

	gitlabClient := mock_gitlab.SetupClient(t, "", "/", mux.ServeHTTP)
	v, err := protovalidate.New()
	require.NoError(t, err)

	gomock.InOrder(
		mockAdmissionControllerClient.EXPECT().AdmitJob(gomock.Any(), api.AgentKey{ID: 2001, Type: api.AgentTypeRunnerController}, gomock.Any()).
			Return(&runner_controller_rpc.AdmitJobResponse{
				AdmissionResponse: &runner_controller_rpc.AdmitJobResponse_Admitted{Admitted: &runner_controller_rpc.Admitted{}},
			}, nil),
		mockAdmissionControlRequestsCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "admitted"})),

		mockAdmissionControllerClient.EXPECT().AdmitJob(gomock.Any(), api.AgentKey{ID: 2002, Type: api.AgentTypeRunnerController}, gomock.Any()).
			Return(&runner_controller_rpc.AdmitJobResponse{
				AdmissionResponse: &runner_controller_rpc.AdmitJobResponse_Admitted{Admitted: &runner_controller_rpc.Admitted{}},
			}, nil),
		mockAdmissionControlRequestsCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "admitted"})),

		mockAdmissionControllerClient.EXPECT().AdmitJob(gomock.Any(), api.AgentKey{ID: 2003, Type: api.AgentTypeRunnerController}, gomock.Any()).
			Return(&runner_controller_rpc.AdmitJobResponse{
				AdmissionResponse: &runner_controller_rpc.AdmitJobResponse_Admitted{Admitted: &runner_controller_rpc.Admitted{}},
			}, nil),
		mockAdmissionControlRequestsCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "admitted"})),

		mockAdmissionControlCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "admitted"})),
	)

	s := &server{
		gitlabClient:                         gitlabClient,
		validator:                            v,
		newAdmissionControllerClient:         func() admissionControllerClient { return mockAdmissionControllerClient },
		admissionControlTotalCounter:         mockAdmissionControlCounter,
		admissionControlRequestsTotalCounter: mockAdmissionControlRequestsCounter,
	}

	ctx := mock_modserver.IncomingAgentCtx(t, mockRPCAPI)
	req := &rpc.GetJobRequest{
		JobRequest: []byte(`{}`),
	}

	resp, err := s.GetJob(ctx, req)

	require.NoError(t, err)
	assert.JSONEq(t, jobResponse, string(resp.JobResponse))
}

func TestGetJob_WithJob_MultipleRunnerControllersSecondRejects(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockRPCAPI := mock_modserver.NewMockAgentRPCAPI(ctrl)
	mockAdmissionControlCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControlRequestsCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControllerClient := NewMockadmissionControllerClient(ctrl)

	mockRPCAPI.EXPECT().AgentInfo(gomock.Any(), gomock.Any()).Return(&server_api.RunnerInfo{
		Key: testhelpers.RunnerKey1,
	}, nil)

	mux := http.NewServeMux()
	mux.HandleFunc(gapi.JobsRequestAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPost, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusCreated)
		_, err := w.Write([]byte(jobResponse))
		assert.NoError(t, err)
	})
	mux.HandleFunc(gapi.GetRunnerControllersForJobAdmissionAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodGet, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, err := w.Write([]byte(`{"runner_controllers":[{"id":3001,"state":"enabled"},{"id":3002,"state":"enabled"},{"id":3003,"state":"enabled"}]}`))
		assert.NoError(t, err)
	})
	mux.HandleFunc(gapi.UpdateCIJobAPIPath(123), func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPut, r.Method)
		w.WriteHeader(http.StatusOK)
	})

	gitlabClient := mock_gitlab.SetupClient(t, "", "/", mux.ServeHTTP)
	v, err := protovalidate.New()
	require.NoError(t, err)

	gomock.InOrder(
		mockAdmissionControllerClient.EXPECT().AdmitJob(gomock.Any(), api.AgentKey{ID: 3001, Type: api.AgentTypeRunnerController}, gomock.Any()).
			Return(&runner_controller_rpc.AdmitJobResponse{
				AdmissionResponse: &runner_controller_rpc.AdmitJobResponse_Admitted{Admitted: &runner_controller_rpc.Admitted{}},
			}, nil),
		mockAdmissionControlRequestsCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "admitted"})),

		mockAdmissionControllerClient.EXPECT().AdmitJob(gomock.Any(), api.AgentKey{ID: 3002, Type: api.AgentTypeRunnerController}, gomock.Any()).
			Return(&runner_controller_rpc.AdmitJobResponse{
				AdmissionResponse: &runner_controller_rpc.AdmitJobResponse_Rejected{Rejected: &runner_controller_rpc.Rejected{Reason: "not admitted"}},
			}, nil),
		mockAdmissionControlRequestsCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "rejected"})),
	)

	mockAdmissionControlCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "rejected"}))

	// Note: No mock for the third runner controller - it should not be called

	s := &server{
		gitlabClient:                         gitlabClient,
		validator:                            v,
		newAdmissionControllerClient:         func() admissionControllerClient { return mockAdmissionControllerClient },
		admissionControlTotalCounter:         mockAdmissionControlCounter,
		admissionControlRequestsTotalCounter: mockAdmissionControlRequestsCounter,
	}

	ctx := mock_modserver.IncomingAgentCtx(t, mockRPCAPI)
	req := &rpc.GetJobRequest{
		JobRequest: []byte(`{}`),
	}

	resp, err := s.GetJob(ctx, req)

	require.NoError(t, err)
	assert.Empty(t, resp.JobResponse)
}

func TestGetJob_WithJob_AdmissionControlNotAvailable_NonEnterpriseInstances(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockRPCAPI := mock_modserver.NewMockAgentRPCAPI(ctrl)
	mockAdmissionControlCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControlRequestsCounter := mock_otel.NewMockInt64Counter(ctrl)

	mockRPCAPI.EXPECT().AgentInfo(gomock.Any(), gomock.Any()).Return(&server_api.RunnerInfo{
		Key: testhelpers.RunnerKey1,
	}, nil)

	mux := http.NewServeMux()
	mux.HandleFunc(gapi.JobsRequestAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPost, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusCreated)
		_, err := w.Write([]byte(jobResponse))
		assert.NoError(t, err)
	})
	mux.HandleFunc(gapi.GetRunnerControllersForJobAdmissionAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodGet, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.WriteHeader(http.StatusNotFound)
	})

	gitlabClient := mock_gitlab.SetupClient(t, "", "/", mux.ServeHTTP)
	v, err := protovalidate.New()
	require.NoError(t, err)

	mockAdmissionControlCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "disabled"}))

	s := &server{
		gitlabClient:                         gitlabClient,
		validator:                            v,
		admissionControlTotalCounter:         mockAdmissionControlCounter,
		admissionControlRequestsTotalCounter: mockAdmissionControlRequestsCounter,
	}

	ctx := mock_modserver.IncomingAgentCtx(t, mockRPCAPI)
	req := &rpc.GetJobRequest{
		JobRequest: []byte(`{}`),
	}

	resp, err := s.GetJob(ctx, req)

	require.NoError(t, err)
	assert.JSONEq(t, jobResponse, string(resp.JobResponse))
}

func TestGetJob_WithJob_AdmissionControlNotAvailable_JobRouterOrFeatureFlagDisabled(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockRPCAPI := mock_modserver.NewMockAgentRPCAPI(ctrl)
	mockAdmissionControlCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControlRequestsCounter := mock_otel.NewMockInt64Counter(ctrl)

	mockRPCAPI.EXPECT().AgentInfo(gomock.Any(), gomock.Any()).Return(&server_api.RunnerInfo{
		Key: testhelpers.RunnerKey1,
	}, nil)

	mux := http.NewServeMux()
	mux.HandleFunc(gapi.JobsRequestAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPost, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusCreated)
		_, err := w.Write([]byte(jobResponse))
		assert.NoError(t, err)
	})
	mux.HandleFunc(gapi.GetRunnerControllersForJobAdmissionAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodGet, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.WriteHeader(http.StatusNotImplemented)
	})

	gitlabClient := mock_gitlab.SetupClient(t, "", "/", mux.ServeHTTP)
	v, err := protovalidate.New()
	require.NoError(t, err)

	mockAdmissionControlCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "disabled"}))

	s := &server{
		gitlabClient:                         gitlabClient,
		validator:                            v,
		admissionControlTotalCounter:         mockAdmissionControlCounter,
		admissionControlRequestsTotalCounter: mockAdmissionControlRequestsCounter,
	}

	ctx := mock_modserver.IncomingAgentCtx(t, mockRPCAPI)
	req := &rpc.GetJobRequest{
		JobRequest: []byte(`{}`),
	}

	resp, err := s.GetJob(ctx, req)

	require.NoError(t, err)
	assert.JSONEq(t, jobResponse, string(resp.JobResponse))
}

func TestGetJob_WithJob_NoRunnerControllers(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockRPCAPI := mock_modserver.NewMockAgentRPCAPI(ctrl)
	mockAdmissionControlCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControlRequestsCounter := mock_otel.NewMockInt64Counter(ctrl)

	mockRPCAPI.EXPECT().AgentInfo(gomock.Any(), gomock.Any()).Return(&server_api.RunnerInfo{
		Key: testhelpers.RunnerKey1,
	}, nil)

	mux := http.NewServeMux()
	mux.HandleFunc(gapi.JobsRequestAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPost, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusCreated)
		_, err := w.Write([]byte(jobResponse))
		assert.NoError(t, err)
	})
	mux.HandleFunc(gapi.GetRunnerControllersForJobAdmissionAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodGet, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, err := w.Write([]byte(`{"runner_controllers":[]}`))
		assert.NoError(t, err)
	})

	gitlabClient := mock_gitlab.SetupClient(t, "", "/", mux.ServeHTTP)
	v, err := protovalidate.New()
	require.NoError(t, err)

	mockAdmissionControlCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "admitted"}))

	s := &server{
		gitlabClient:                         gitlabClient,
		validator:                            v,
		admissionControlTotalCounter:         mockAdmissionControlCounter,
		admissionControlRequestsTotalCounter: mockAdmissionControlRequestsCounter,
	}

	ctx := mock_modserver.IncomingAgentCtx(t, mockRPCAPI)
	req := &rpc.GetJobRequest{
		JobRequest: []byte(`{}`),
	}

	resp, err := s.GetJob(ctx, req)
	require.NoError(t, err)

	assert.JSONEq(t, jobResponse, string(resp.JobResponse))
}

func TestGetJob_WithJob_RunnerControllerError(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockRPCAPI := mock_modserver.NewMockAgentRPCAPI(ctrl)
	mockAdmissionControlCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControlRequestsCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControllerClient := NewMockadmissionControllerClient(ctrl)

	mockRPCAPI.EXPECT().AgentInfo(gomock.Any(), gomock.Any()).Return(&server_api.RunnerInfo{
		Key: testhelpers.RunnerKey1,
	}, nil)

	mux := http.NewServeMux()
	mux.HandleFunc(gapi.JobsRequestAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPost, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusCreated)
		_, err := w.Write([]byte(jobResponse))
		assert.NoError(t, err)
	})
	mux.HandleFunc(gapi.GetRunnerControllersForJobAdmissionAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodGet, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, err := w.Write([]byte(`{"runner_controllers":[{"id":789,"state":"enabled"}]}`))
		assert.NoError(t, err)
	})
	mux.HandleFunc(gapi.UpdateCIJobAPIPath(123), func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPut, r.Method)
		w.WriteHeader(http.StatusOK)
	})

	gitlabClient := mock_gitlab.SetupClient(t, "", "/", mux.ServeHTTP)
	v, err := protovalidate.New()
	require.NoError(t, err)

	gomock.InOrder(
		mockAdmissionControllerClient.EXPECT().AdmitJob(gomock.Any(), api.AgentKey{ID: 789, Type: api.AgentTypeRunnerController}, gomock.Any()).
			Return(nil, errors.New("some error")),

		mockAdmissionControlRequestsCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "failed"})),

		mockAdmissionControlCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "failed"})),

		mockRPCAPI.EXPECT().
			HandleProcessingError(gomock.Any(), "Unable to perform admission control", gomock.Any(), gomock.Any()),
	)

	s := &server{
		gitlabClient:                         gitlabClient,
		validator:                            v,
		newAdmissionControllerClient:         func() admissionControllerClient { return mockAdmissionControllerClient },
		admissionControlTotalCounter:         mockAdmissionControlCounter,
		admissionControlRequestsTotalCounter: mockAdmissionControlRequestsCounter,
	}

	ctx := mock_modserver.IncomingAgentCtx(t, mockRPCAPI)
	req := &rpc.GetJobRequest{
		JobRequest: []byte(`{}`),
	}

	resp, err := s.GetJob(ctx, req)

	require.NoError(t, err)
	assert.Empty(t, resp.JobResponse)
}

func TestGetJob_WithJob_RunnerControllerRejects(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockRPCAPI := mock_modserver.NewMockAgentRPCAPI(ctrl)
	mockAdmissionControlCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControlRequestsCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControllerClient := NewMockadmissionControllerClient(ctrl)

	mockRPCAPI.EXPECT().AgentInfo(gomock.Any(), gomock.Any()).Return(&server_api.RunnerInfo{
		Key: testhelpers.RunnerKey1,
	}, nil)

	mux := http.NewServeMux()
	mux.HandleFunc(gapi.JobsRequestAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPost, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusCreated)
		_, err := w.Write([]byte(jobResponse))
		assert.NoError(t, err)
	})
	mux.HandleFunc(gapi.GetRunnerControllersForJobAdmissionAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodGet, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, err := w.Write([]byte(`{"runner_controllers":[{"id":999,"state":"enabled"}]}`))
		assert.NoError(t, err)
	})
	mux.HandleFunc(gapi.UpdateCIJobAPIPath(123), func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPut, r.Method)
		w.WriteHeader(http.StatusOK)
	})

	gitlabClient := mock_gitlab.SetupClient(t, "", "/", mux.ServeHTTP)
	v, err := protovalidate.New()
	require.NoError(t, err)

	gomock.InOrder(
		mockAdmissionControllerClient.EXPECT().AdmitJob(gomock.Any(), api.AgentKey{ID: 999, Type: api.AgentTypeRunnerController}, gomock.Any()).
			Return(&runner_controller_rpc.AdmitJobResponse{
				AdmissionResponse: &runner_controller_rpc.AdmitJobResponse_Rejected{Rejected: &runner_controller_rpc.Rejected{Reason: "not admitted"}},
			}, nil),
		mockAdmissionControlRequestsCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "rejected"})),

		mockAdmissionControlCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "rejected"})),
	)

	s := &server{
		gitlabClient:                         gitlabClient,
		validator:                            v,
		newAdmissionControllerClient:         func() admissionControllerClient { return mockAdmissionControllerClient },
		admissionControlTotalCounter:         mockAdmissionControlCounter,
		admissionControlRequestsTotalCounter: mockAdmissionControlRequestsCounter,
	}

	ctx := mock_modserver.IncomingAgentCtx(t, mockRPCAPI)
	req := &rpc.GetJobRequest{
		JobRequest: []byte(`{}`),
	}

	resp, err := s.GetJob(ctx, req)

	require.NoError(t, err)
	assert.Empty(t, resp.JobResponse)
}

func TestGetJob_WithJob_MultipleRunnerControllersSecondRejectsInDryRun(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockRPCAPI := mock_modserver.NewMockAgentRPCAPI(ctrl)
	mockAdmissionControlCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControlRequestsCounter := mock_otel.NewMockInt64Counter(ctrl)
	mockAdmissionControllerClient := NewMockadmissionControllerClient(ctrl)

	mockRPCAPI.EXPECT().AgentInfo(gomock.Any(), gomock.Any()).Return(&server_api.RunnerInfo{
		Key: testhelpers.RunnerKey1,
	}, nil)

	mux := http.NewServeMux()
	mux.HandleFunc(gapi.JobsRequestAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPost, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusCreated)
		_, err := w.Write([]byte(jobResponse))
		assert.NoError(t, err)
	})
	mux.HandleFunc(gapi.GetRunnerControllersForJobAdmissionAPIPath, func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodGet, r.Method)
		assert.Equal(t, string(testhelpers.AgentToken), r.Header.Get(httpz.RunnerTokenHeader))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, err := w.Write([]byte(`{"runner_controllers":[{"id":3001,"state":"enabled"},{"id":3002,"state":"dry_run"},{"id":3003,"state":"enabled"}]}`))
		assert.NoError(t, err)
	})
	mux.HandleFunc(gapi.UpdateCIJobAPIPath(123), func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPut, r.Method)
		w.WriteHeader(http.StatusOK)
	})

	gitlabClient := mock_gitlab.SetupClient(t, "", "/", mux.ServeHTTP)
	v, err := protovalidate.New()
	require.NoError(t, err)

	gomock.InOrder(
		mockAdmissionControllerClient.EXPECT().AdmitJob(gomock.Any(), api.AgentKey{ID: 3001, Type: api.AgentTypeRunnerController}, gomock.Any()).
			Return(&runner_controller_rpc.AdmitJobResponse{
				AdmissionResponse: &runner_controller_rpc.AdmitJobResponse_Admitted{Admitted: &runner_controller_rpc.Admitted{}},
			}, nil),
		mockAdmissionControlRequestsCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "admitted"})),

		mockAdmissionControllerClient.EXPECT().AdmitJob(gomock.Any(), api.AgentKey{ID: 3002, Type: api.AgentTypeRunnerController}, gomock.Any()).
			Return(&runner_controller_rpc.AdmitJobResponse{
				AdmissionResponse: &runner_controller_rpc.AdmitJobResponse_Rejected{Rejected: &runner_controller_rpc.Rejected{Reason: "not admitted"}},
			}, nil),
		mockAdmissionControlRequestsCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "rejected", "dry_run": true})),

		mockAdmissionControllerClient.EXPECT().AdmitJob(gomock.Any(), api.AgentKey{ID: 3003, Type: api.AgentTypeRunnerController}, gomock.Any()).
			Return(&runner_controller_rpc.AdmitJobResponse{
				AdmissionResponse: &runner_controller_rpc.AdmitJobResponse_Admitted{Admitted: &runner_controller_rpc.Admitted{}},
			}, nil),
		mockAdmissionControlRequestsCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "admitted"})),
	)

	mockAdmissionControlCounter.EXPECT().Add(gomock.Any(), int64(1), matcher.MatchOtelAttributes(map[string]any{"result": "admitted"}))

	s := &server{
		gitlabClient:                         gitlabClient,
		validator:                            v,
		newAdmissionControllerClient:         func() admissionControllerClient { return mockAdmissionControllerClient },
		admissionControlTotalCounter:         mockAdmissionControlCounter,
		admissionControlRequestsTotalCounter: mockAdmissionControlRequestsCounter,
	}

	ctx := mock_modserver.IncomingAgentCtx(t, mockRPCAPI)
	req := &rpc.GetJobRequest{
		JobRequest: []byte(`{}`),
	}

	resp, err := s.GetJob(ctx, req)

	require.NoError(t, err)
	assert.JSONEq(t, jobResponse, string(resp.JobResponse))
}
