package server

import (
	"context"
	"errors"
	"testing"

	"github.com/stretchr/testify/assert"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent_registrar/agentk_rpc"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/agent_tracker"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/modserver"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_agent_tracker"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/mock_event_tracker"
	"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/testhelpers"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/testing/testlogger"
	"go.uber.org/mock/gomock"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func TestRegister_Agentk(t *testing.T) {
	mockRPCAPI, mockAgentTracker, mockEventTracker, s, ctx := setupAgentkServer(t)
	req := agentkRegisterRequest()

	mockRPCAPI.EXPECT().
		Log().
		Return(testlogger.New(t))
	mockRPCAPI.EXPECT().
		AgentInfo(gomock.Any(), gomock.Any()).
		Return(testhelpers.AgentkInfoObj(), nil)
	mockAgentTracker.EXPECT().
		RegisterAgentkExpiring(gomock.Any(), gomock.Any()).
		Do(func(ctx context.Context, connectedAgentInfo *agent_tracker.ConnectedAgentkInfo) error {
			assert.Equal(t, testhelpers.AgentkKey1.ID, connectedAgentInfo.AgentId)
			assert.Equal(t, testhelpers.ProjectID, connectedAgentInfo.ProjectId)
			assert.EqualValues(t, 123456789, connectedAgentInfo.ConnectionId)
			return nil
		})
	expectedEvent := &RegisterAgentkEvent{
		ProjectID:          testhelpers.ProjectID,
		AgentVersion:       "v17.1.0",
		Architecture:       "amd64",
		ConnectionID:       123456789,
		AgentID:            testhelpers.AgentkKey1.ID,
		KubernetesVersion:  "1.30",
		ExtraTelemetryData: map[string]string{"some-key": "some-value"},
	}
	mockEventTracker.EXPECT().EmitEvent(expectedEvent)

	resp, err := s.Register(ctx, req)
	assert.NotNil(t, resp)
	assert.NoError(t, err)
}

func TestRegister_Agentk_AgentInfo_Error(t *testing.T) {
	mockRPCAPI, _, _, s, ctx := setupAgentkServer(t)
	req := agentkRegisterRequest()

	mockRPCAPI.EXPECT().
		Log().
		Return(testlogger.New(t))
	mockRPCAPI.EXPECT().
		AgentInfo(gomock.Any(), gomock.Any()).
		Return(nil, status.Error(codes.Unavailable, "failed to register agent"))

	resp, err := s.Register(ctx, req)
	assert.Nil(t, resp)
	assert.Equal(t, codes.Unavailable, status.Code(err))
}

func TestRegister_Agentk_registerAgent_Error(t *testing.T) {
	mockRPCAPI, mockAgentTracker, _, s, ctx := setupAgentkServer(t)
	req := agentkRegisterRequest()

	expectedErr := errors.New("expected error")

	mockRPCAPI.EXPECT().
		Log().
		Return(testlogger.New(t)).
		AnyTimes()
	mockRPCAPI.EXPECT().
		AgentInfo(gomock.Any(), gomock.Any()).
		Return(testhelpers.AgentkInfoObj(), nil)
	mockAgentTracker.EXPECT().
		RegisterAgentkExpiring(gomock.Any(), gomock.Any()).
		Return(expectedErr)
	mockRPCAPI.EXPECT().
		HandleProcessingError(gomock.Any(), gomock.Any(), expectedErr, logz.AgentKey(testhelpers.AgentkKey1))

	resp, err := s.Register(ctx, req)
	assert.Nil(t, resp)
	assert.Equal(t, codes.Unavailable, status.Code(err))
}

func setupAgentkServer(t *testing.T) (*mock_modserver.MockAgentRPCAPI,
	*mock_agent_tracker.MockTracker, *mock_event_tracker.MockEventsInterface, *agentkServer, context.Context) {
	ctrl := gomock.NewController(t)

	mockRPCAPI := mock_modserver.NewMockAgentRPCAPI(ctrl)
	mockAgentTracker := mock_agent_tracker.NewMockTracker(ctrl)
	mockEventTracker := mock_event_tracker.NewMockEventsInterface(ctrl)

	s := &agentkServer{
		agentRegisterer:           mockAgentTracker,
		registerAgentEventTracker: mockEventTracker,
	}

	ctx := modserver.InjectAgentRPCAPI(context.Background(), mockRPCAPI)

	return mockRPCAPI, mockAgentTracker, mockEventTracker, s, ctx
}

func agentkRegisterRequest() *agentk_rpc.RegisterRequest {
	return &agentk_rpc.RegisterRequest{
		Meta:       mock_modserver.AgentkMeta(),
		InstanceId: 123456789,
	}
}
