package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"
	"sync"
	"testing"

	gomock "github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/module/modagent"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/module/starboard_vulnerability"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/tool/syncz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/tool/testing/mock_modagent"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/tool/testing/testhelpers"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/pkg/agentcfg"
	"go.uber.org/zap/zaptest"
)

var (
	_ modagent.LeaderModule = &module{}
	_ modagent.Factory      = &Factory{}

	// Init common data
	agentConfigOCSSetting = agentcfg.ContainerScanningCF{
		Cadence: "1 * * * *",
		VulnerabilityReport: &agentcfg.VulnerabilityReport{
			Namespaces: []string{"default"},
		},
	}

	policyOCSSetting = agentcfg.ContainerScanningCF{
		Cadence: "2 * * * *",
		VulnerabilityReport: &agentcfg.VulnerabilityReport{
			Namespaces: []string{"kube-system"},
		},
	}

	defaultResourceRequirements = agentcfg.ResourceRequirements{
		Limits: &agentcfg.Resource{
			Cpu:    defaultTrivyResourceLimitsCpu,
			Memory: defaultTrivyResourceLimitsMemory,
		},
		Requests: &agentcfg.Resource{
			Cpu:    defaultTrivyResourceRequestsCpu,
			Memory: defaultTrivyResourceRequestsMemory,
		},
	}

	customResourceRequirements = agentcfg.ResourceRequirements{
		Limits: &agentcfg.Resource{
			Cpu:    "100m",
			Memory: "100Mi",
		},
		Requests: &agentcfg.Resource{
			Cpu:    "100m",
			Memory: "100Mi",
		},
	}
)

func TestEmptyConfig(t *testing.T) {
	m := new(module)
	cfg := &agentcfg.AgentConfiguration{}

	assert.NoError(t, m.DefaultAndValidateConfiguration(cfg))
	finalResourceRequirements := cfg.ContainerScanning.ResourceRequirements

	// m.DefaultAndValidateConfiguration would default resource requirements even if config is empty since security policies could still enforce vulnerability scan
	assert.Equal(t, defaultTrivyResourceLimitsCpu, finalResourceRequirements.Limits.Cpu)
	assert.Equal(t, defaultTrivyResourceLimitsMemory, finalResourceRequirements.Limits.Memory)
	assert.Equal(t, defaultTrivyResourceRequestsCpu, finalResourceRequirements.Requests.Cpu)
	assert.Equal(t, defaultTrivyResourceRequestsMemory, finalResourceRequirements.Requests.Memory)
}

func TestConfigWithResourceRequirements(t *testing.T) {
	m := new(module)
	cfg := &agentcfg.AgentConfiguration{
		ContainerScanning: &agentcfg.ContainerScanningCF{
			ResourceRequirements: &customResourceRequirements,
		},
	}

	assert.NoError(t, m.DefaultAndValidateConfiguration(cfg))
	finalResourceRequirements := cfg.ContainerScanning.ResourceRequirements

	assert.Equal(t, customResourceRequirements.Limits.Cpu, finalResourceRequirements.Limits.Cpu)
	assert.Equal(t, customResourceRequirements.Limits.Memory, finalResourceRequirements.Limits.Memory)
	assert.Equal(t, customResourceRequirements.Requests.Cpu, finalResourceRequirements.Requests.Cpu)
	assert.Equal(t, customResourceRequirements.Requests.Memory, finalResourceRequirements.Requests.Memory)
}

func TestConfigWithPartialResourceRequirements(t *testing.T) {
	testCases := []struct {
		description            string
		inputLimitsCpu         string
		inputLimitsMemory      string
		inputRequestsCpu       string
		inputRequestsMemory    string
		expectedLimitsCpu      string
		expectedLimitsMemory   string
		expectedRequestsCpu    string
		expectedRequestsMemory string
	}{
		{
			description:            "only has limits cpu",
			inputLimitsCpu:         customResourceRequirements.Limits.Cpu,
			inputLimitsMemory:      "",
			inputRequestsCpu:       "",
			inputRequestsMemory:    "",
			expectedLimitsCpu:      customResourceRequirements.Limits.Cpu,
			expectedLimitsMemory:   defaultResourceRequirements.Limits.Memory,
			expectedRequestsCpu:    defaultResourceRequirements.Requests.Cpu,
			expectedRequestsMemory: defaultResourceRequirements.Requests.Memory,
		},
		{
			description:            "only has limits memory",
			inputLimitsCpu:         "",
			inputLimitsMemory:      customResourceRequirements.Limits.Memory,
			inputRequestsCpu:       "",
			inputRequestsMemory:    "",
			expectedLimitsCpu:      defaultResourceRequirements.Limits.Cpu,
			expectedLimitsMemory:   customResourceRequirements.Limits.Memory,
			expectedRequestsCpu:    defaultResourceRequirements.Requests.Cpu,
			expectedRequestsMemory: defaultResourceRequirements.Requests.Memory,
		},
		{
			description:            "only has requests cpu",
			inputLimitsCpu:         "",
			inputLimitsMemory:      "",
			inputRequestsCpu:       customResourceRequirements.Requests.Cpu,
			inputRequestsMemory:    "",
			expectedLimitsCpu:      defaultResourceRequirements.Limits.Cpu,
			expectedLimitsMemory:   defaultResourceRequirements.Limits.Memory,
			expectedRequestsCpu:    customResourceRequirements.Requests.Cpu,
			expectedRequestsMemory: defaultResourceRequirements.Requests.Memory,
		},
		{
			description:            "only has requests memory",
			inputLimitsCpu:         "",
			inputLimitsMemory:      "",
			inputRequestsCpu:       "",
			inputRequestsMemory:    customResourceRequirements.Requests.Memory,
			expectedLimitsCpu:      defaultResourceRequirements.Limits.Cpu,
			expectedLimitsMemory:   defaultResourceRequirements.Limits.Memory,
			expectedRequestsCpu:    defaultResourceRequirements.Requests.Cpu,
			expectedRequestsMemory: customResourceRequirements.Requests.Memory,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.description, func(t *testing.T) {
			m := new(module)
			cfg := &agentcfg.AgentConfiguration{
				ContainerScanning: &agentcfg.ContainerScanningCF{
					ResourceRequirements: &agentcfg.ResourceRequirements{
						Limits: &agentcfg.Resource{
							Cpu:    tc.inputLimitsCpu,
							Memory: tc.inputLimitsMemory,
						},
						Requests: &agentcfg.Resource{
							Cpu:    tc.inputRequestsCpu,
							Memory: tc.inputRequestsMemory,
						},
					},
				},
			}

			assert.NoError(t, m.DefaultAndValidateConfiguration(cfg))
			newResourceRequirements := cfg.ContainerScanning.ResourceRequirements

			assert.Equal(t, tc.expectedLimitsCpu, newResourceRequirements.Limits.Cpu)
			assert.Equal(t, tc.expectedLimitsMemory, newResourceRequirements.Limits.Memory)
			assert.Equal(t, tc.expectedRequestsCpu, newResourceRequirements.Requests.Cpu)
			assert.Equal(t, tc.expectedRequestsMemory, newResourceRequirements.Requests.Memory)
		})
	}
}

func TestName(t *testing.T) {
	m := new(module)

	assert.Equal(t, starboard_vulnerability.ModuleName, m.Name())
}

func TestCadenceValidation(t *testing.T) {
	// Possible values documented at https://pkg.go.dev/github.com/robfig/cron#hdr-CRON_Expression_Format
	testCases := []struct {
		description string
		cadence     string
		shouldError bool
	}{
		{
			description: "daily schedule",
			cadence:     "0 0 * * *",
			shouldError: false,
		},
		{
			description: "using all fields",
			cadence:     "0 0 1 1 1",
			shouldError: false,
		},
		{
			description: "range of seconds",
			cadence:     "0-30 0 * * *",
			shouldError: false,
		},
		{
			description: "range of minutes",
			cadence:     "0 0-23 * * *",
			shouldError: false,
		},
		{
			description: "range of hours",
			cadence:     "0 0 1-12 * *",
			shouldError: false,
		},
		{
			description: "range of days",
			cadence:     "0 0 * 1-12 *",
			shouldError: false,
		},
		{
			description: "range of months (numeric)",
			cadence:     "0 0 * * 1-6",
			shouldError: false,
		},
		{
			description: "using question mark",
			cadence:     "0 0 * ? *",
			shouldError: false,
		},
		{
			description: "seconds out of range",
			cadence:     "60 0 * * *",
			shouldError: true,
		},
		{
			description: "minutes out of range",
			cadence:     "60 0 * * *",
			shouldError: true,
		},
		{
			description: "hours out of range",
			cadence:     "0 24 1 * *",
			shouldError: true,
		},
		{
			description: "days out of range",
			cadence:     "0 0 * 32 *",
			shouldError: true,
		},
		{
			description: "months out of range",
			cadence:     "0 0 * * 13",
			shouldError: true,
		},
		{
			description: "zero day",
			cadence:     "0 0 * 0 *",
			shouldError: true,
		},
		{
			description: "zero month",
			cadence:     "0 0 * 0 0",
			shouldError: true,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.description, func(t *testing.T) {
			m := new(module)
			cfg := &agentcfg.AgentConfiguration{
				ContainerScanning: &agentcfg.ContainerScanningCF{
					Cadence: tc.cadence,
				},
			}

			err := m.DefaultAndValidateConfiguration(cfg)
			if tc.shouldError {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
			}
		})
	}
}

func TestIsConfigEqual(t *testing.T) {
	oldCfg := configurationToUpdateData{containerScanningConfig: &agentcfg.ContainerScanningCF{}}
	newCfg := configurationToUpdateData{containerScanningConfig: &agentcfg.ContainerScanningCF{}}

	assert.True(t, isConfigEqual(oldCfg, newCfg))

	newCfg.containerScanningConfig = &agentcfg.ContainerScanningCF{
		Cadence: "0 0 * * *",
	}

	assert.False(t, isConfigEqual(oldCfg, newCfg))
}

const (
	AGENT_CONFIG = "AGENT CONFIG"
	POLICY       = "POLICY"
)

type customWorkerFactory struct {
	cfg      configurationToUpdateData
	cfgMutex sync.Mutex
	wg       sync.WaitGroup
}

func (f *customWorkerFactory) New(cfg configurationToUpdateData) syncz.Worker {
	// decrement wg
	defer f.wg.Done()

	// Save cfg
	f.cfgMutex.Lock()
	defer f.cfgMutex.Unlock()
	f.cfg = cfg

	return syncz.WorkerFunc(func(ctx context.Context) {}) // do nothing
}

func TestModuleConfigScenarios(t *testing.T) {
	testCases := []struct {
		description                        string
		policyHasOCSSetting                bool
		agentConfigHasOCSSetting           bool
		agentConfigHasResourceRequirements bool
		expectCustomResourceSet            bool
		expectOCSSettingFrom               string
	}{
		{
			description:                        "OCS setting from policy, resource not set",
			policyHasOCSSetting:                true,
			agentConfigHasOCSSetting:           false,
			agentConfigHasResourceRequirements: false,
			expectCustomResourceSet:            false,
			expectOCSSettingFrom:               POLICY,
		},
		{
			description:                        "OCS setting from policy and agentConfig, resource not set",
			policyHasOCSSetting:                true,
			agentConfigHasOCSSetting:           true,
			agentConfigHasResourceRequirements: false,
			expectCustomResourceSet:            false,
			expectOCSSettingFrom:               POLICY,
		},
		{
			description:                        "OCS setting from policy and agentConfig, resource set",
			policyHasOCSSetting:                true,
			agentConfigHasOCSSetting:           true,
			agentConfigHasResourceRequirements: true,
			expectCustomResourceSet:            true,
			expectOCSSettingFrom:               POLICY,
		},
		{
			description:                        "OCS setting from policy, resource set",
			policyHasOCSSetting:                true,
			agentConfigHasOCSSetting:           false,
			agentConfigHasResourceRequirements: true,
			expectCustomResourceSet:            true,
			expectOCSSettingFrom:               POLICY,
		},
		{
			description:                        "OCS setting from agentConfig, resource not set",
			policyHasOCSSetting:                false,
			agentConfigHasOCSSetting:           true,
			agentConfigHasResourceRequirements: false,
			expectCustomResourceSet:            false,
			expectOCSSettingFrom:               AGENT_CONFIG,
		},
		{
			description:                        "OCS setting from agentConfig, resource set",
			policyHasOCSSetting:                false,
			agentConfigHasOCSSetting:           true,
			agentConfigHasResourceRequirements: true,
			expectCustomResourceSet:            true,
			expectOCSSettingFrom:               AGENT_CONFIG,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.description, func(t *testing.T) {
			// Setup module
			policiesUpdateDataChannel := make(chan configurationToUpdateData)
			ctrl := gomock.NewController(t)
			api := mock_modagent.NewMockApi(ctrl)
			cnwh := customWorkerFactory{}
			m := module{
				log:                       zaptest.NewLogger(t),
				api:                       api,
				policiesUpdateDataChannel: policiesUpdateDataChannel,
				workerFactory:             cnwh.New,
			}

			// Stub policies_configuration request
			api.EXPECT().
				MakeGitLabRequest(gomock.Any(), "/policies_configuration", gomock.Any()).
				MinTimes(0).
				DoAndReturn(func(ctx context.Context, path string, opts ...modagent.GitLabRequestOption) (*modagent.GitLabResponse, error) {
					body, err := json.Marshal(map[string]interface{}{})
					if err != nil {
						return nil, err
					}
					return &modagent.GitLabResponse{
						StatusCode: http.StatusOK,
						Body:       io.NopCloser(bytes.NewReader(body)),
					}, nil
				})
			// Stub GetAgentId request
			api.EXPECT().
				GetAgentId(gomock.Any()).
				MinTimes(0).
				Return(testhelpers.AgentId, nil)

			ctx := context.Background()
			var wg sync.WaitGroup
			wg.Add(1)
			defer wg.Wait()

			repositoryConfigChan := make(chan *agentcfg.AgentConfiguration)
			defer close(repositoryConfigChan)

			go func() {
				defer wg.Done()
				err := m.Run(ctx, repositoryConfigChan)
				assert.NoError(t, err)
			}()

			// Configure agentConfig to be sent to module
			agentConfig := agentcfg.AgentConfiguration{
				ContainerScanning: &agentcfg.ContainerScanningCF{},
			}
			if tc.agentConfigHasOCSSetting {
				agentConfig.ContainerScanning.Cadence = agentConfigOCSSetting.Cadence
				agentConfig.ContainerScanning.VulnerabilityReport = agentConfigOCSSetting.VulnerabilityReport
				// Increment wg since change in the agentConfig OCS setting would cause module.Run to update the new
				// worker holder and we want to wait for that process to complete before continuing with the test.
				cnwh.wg.Add(1)
			}
			if tc.agentConfigHasResourceRequirements {
				agentConfig.ContainerScanning.ResourceRequirements = &customResourceRequirements
			}

			// Send agentConfig to module
			repositoryConfigChan <- &agentConfig

			// Configure policy to be sent to module
			policy := configurationToUpdateData{
				agentId: testhelpers.AgentId,
			}
			if tc.policyHasOCSSetting {
				policy.containerScanningConfig = &policyOCSSetting
				// Increment wg since change in the policy would cause module.Run to update the new worker holder
				// and we want to wait for that process to complete before continuing with the test.
				cnwh.wg.Add(1)
			}

			// Send policy to module
			policiesUpdateDataChannel <- policy

			// Wait for module to complete processing of new CS config before asserting
			cnwh.wg.Wait()
			cnwh.cfgMutex.Lock()
			defer cnwh.cfgMutex.Unlock()
			finalContainerScanningConfig := cnwh.cfg.containerScanningConfig

			if tc.expectOCSSettingFrom == AGENT_CONFIG {
				assert.Equal(t, agentConfigOCSSetting.Cadence, finalContainerScanningConfig.Cadence)
				assert.Equal(t, agentConfigOCSSetting.VulnerabilityReport, finalContainerScanningConfig.VulnerabilityReport)
			}
			if tc.expectOCSSettingFrom == POLICY {
				assert.Equal(t, policyOCSSetting.Cadence, finalContainerScanningConfig.Cadence)
				assert.Equal(t, policyOCSSetting.VulnerabilityReport, finalContainerScanningConfig.VulnerabilityReport)
			}
			if tc.expectCustomResourceSet {
				assert.Equal(t, &customResourceRequirements, finalContainerScanningConfig.ResourceRequirements)
			}
		})
	}
}

func TestModuleConfigChangeScenarios(t *testing.T) {
	testCases := []struct {
		description          string
		lastOCSSettingFrom   string
		expectOCSSettingFrom string
	}{
		{
			description:          "policy configured after agentConfig",
			lastOCSSettingFrom:   POLICY,
			expectOCSSettingFrom: POLICY,
		},
		{
			description:          "agentConfig configured after policy",
			lastOCSSettingFrom:   POLICY,
			expectOCSSettingFrom: POLICY,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.description, func(t *testing.T) {
			// Setup module
			policiesUpdateDataChannel := make(chan configurationToUpdateData)
			ctrl := gomock.NewController(t)
			api := mock_modagent.NewMockApi(ctrl)
			cnwh := customWorkerFactory{}
			m := module{
				log:                       zaptest.NewLogger(t),
				api:                       api,
				policiesUpdateDataChannel: policiesUpdateDataChannel,
				workerFactory:             cnwh.New,
			}

			// Stub policies_configuration request
			api.EXPECT().
				MakeGitLabRequest(gomock.Any(), "/policies_configuration", gomock.Any()).
				MinTimes(0).
				DoAndReturn(func(ctx context.Context, path string, opts ...modagent.GitLabRequestOption) (*modagent.GitLabResponse, error) {
					body, err := json.Marshal(map[string]interface{}{})
					if err != nil {
						return nil, err
					}
					return &modagent.GitLabResponse{
						StatusCode: http.StatusOK,
						Body:       io.NopCloser(bytes.NewReader(body)),
					}, nil
				})
			// Stub GetAgentId request
			api.EXPECT().
				GetAgentId(gomock.Any()).
				MinTimes(0).
				Return(testhelpers.AgentId, nil)

			ctx := context.Background()
			var wg sync.WaitGroup
			wg.Add(1)
			defer wg.Wait()

			repositoryConfigChan := make(chan *agentcfg.AgentConfiguration)
			defer close(repositoryConfigChan)

			go func() {
				defer wg.Done()
				err := m.Run(ctx, repositoryConfigChan)
				assert.NoError(t, err)
			}()

			// Configure agentConfig to be sent to module
			agentConfig := agentcfg.AgentConfiguration{
				ContainerScanning: &agentcfg.ContainerScanningCF{
					Cadence:              agentConfigOCSSetting.Cadence,
					VulnerabilityReport:  agentConfigOCSSetting.VulnerabilityReport,
					ResourceRequirements: &customResourceRequirements,
				},
			}

			// Configure policy to be sent to module
			policy := configurationToUpdateData{
				agentId:                 testhelpers.AgentId,
				containerScanningConfig: &policyOCSSetting,
			}

			// Send agentConfig and policy to module. Increment wg since these config changes would cause module.Run
			// to update the new worker holder and we want to wait for that process to complete before continuing with the test.
			if tc.lastOCSSettingFrom == POLICY {
				cnwh.wg.Add(1)
				repositoryConfigChan <- &agentConfig
				cnwh.wg.Add(1)
				policiesUpdateDataChannel <- policy
			}
			if tc.lastOCSSettingFrom == AGENT_CONFIG {
				cnwh.wg.Add(1)
				policiesUpdateDataChannel <- policy
				cnwh.wg.Add(1)
				repositoryConfigChan <- &agentConfig
			}

			// Wait for module to complete processing of new CS config before asserting
			cnwh.wg.Wait()
			cnwh.cfgMutex.Lock()
			defer cnwh.cfgMutex.Unlock()
			finalContainerScanningConfig := cnwh.cfg.containerScanningConfig

			if tc.expectOCSSettingFrom == AGENT_CONFIG {
				assert.Equal(t, agentConfigOCSSetting.Cadence, finalContainerScanningConfig.Cadence)
				assert.Equal(t, agentConfigOCSSetting.VulnerabilityReport, finalContainerScanningConfig.VulnerabilityReport)
			}
			if tc.expectOCSSettingFrom == POLICY {
				assert.Equal(t, policyOCSSetting.Cadence, finalContainerScanningConfig.Cadence)
				assert.Equal(t, policyOCSSetting.VulnerabilityReport, finalContainerScanningConfig.VulnerabilityReport)
			}

			assert.Equal(t, &customResourceRequirements, finalContainerScanningConfig.ResourceRequirements)
		})
	}
}