package agent

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"regexp"
	"strings"
	"sync"
	"time"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/module/modagent"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/tool/logz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/pkg/agentcfg"
	"go.uber.org/zap"
	corev1 "k8s.io/api/core/v1"
	k8errors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
)

const (
	maxParallel                      int = 10 // Trivy scanner Pods batch size
	TrivyScannerVersion                  = "0.38.3"
	trivyImageRef                        = "aquasec/trivy:" + TrivyScannerVersion
	kindPod                              = "Pod"
	kindReplicaSet                       = "ReplicaSet"
	kindReplicationController            = "ReplicationController"
	kindStatefulSet                      = "StatefulSet"
	kindDaemonSet                        = "DaemonSet"
	kindCronJob                          = "CronJob"
	kindJob                              = "Job"
	envVarTrivyResourceRequestCpu        = "TRIVY_RESOURCE_REQUEST_CPU"
	envVarTrivyResourceRequestMemory     = "TRIVY_RESOURCE_REQUEST_MEMORY"
	envVarTrivyResourceLimitCpu          = "TRIVY_RESOURCE_LIMIT_CPU"
	envVarTrivyResourceLimitMemory       = "TRIVY_RESOURCE_LIMIT_MEMORY"
)

var (
	// This regex matches a JSON object with first key ClusterName.
	// Example input: `some Error\n  \n  {\n\t\"ClusterName\": \"\"\n\t}`
	// Example output: `{\n\t\"ClusterName\": \"\"\n\t}`
	// first part `\{\n\s*"ClusterName`: Matches for the opening brace `{` as well as the first key ClusterName
	// second part `(?s).*`: matches all characters after ClusterName. setting the s flag to allow . to match newline characters
	jsonLogRegexp = regexp.MustCompile(`\{\n\s*"ClusterName(?s).*$`)
)

type scanJob struct {
	log                       *zap.Logger
	api                       modagent.Api
	kubeClientset             kubernetes.Interface
	gitlabAgentNamespace      string
	gitlabAgentServiceAccount string
	agentID                   int64
	targetNamespaces          []string
	logParser                 LogParser
	resourceRequirements      *agentcfg.ResourceRequirements
}

type resourcesQuantity struct {
	requestsCpu    resource.Quantity
	requestsMemory resource.Quantity
	limitsCpu      resource.Quantity
	limitsMemory   resource.Quantity
}

type resourceQuantity struct {
	Cpu    resource.Quantity
	Memory resource.Quantity
}

type LogParser interface {
	ParsePodLogsToReport(logs []byte) (ConsolidatedReport, error)
}

type logParserImpl struct{}

func (lp *logParserImpl) ParsePodLogsToReport(logs []byte) (ConsolidatedReport, error) {
	// logs could contain errors that is printed before the JSON object. This regex retrieves the JSON report from the logs. See TestParsePodLogsToReportWithError for example of log output with error.
	match := jsonLogRegexp.FindSubmatch(logs)

	if match == nil {
		return ConsolidatedReport{}, errors.New("no JSON report found")
	}

	var consolidatedReport ConsolidatedReport
	err := json.Unmarshal(match[0], &consolidatedReport)
	if err != nil {
		return ConsolidatedReport{}, fmt.Errorf("error parsing JSON in logs for Trivy Pod: %w", err)
	}
	return consolidatedReport, nil
}

func (s *scanJob) Run(ctx context.Context) {
	if err := s.scan(ctx); err != nil {
		s.log.Error("Error running vulnerability scan", logz.Error(err))
	}
}

type uuidCollection struct {
	uuids []string
	mux   sync.Mutex
}

func (u *uuidCollection) Append(uuids []string) {
	u.mux.Lock()
	u.uuids = append(u.uuids, uuids...)
	u.mux.Unlock()
}

func (u *uuidCollection) Items() []string {
	u.mux.Lock()
	defer u.mux.Unlock()
	return u.uuids
}

func (s *scanJob) scan(ctx context.Context) error {
	s.log.Info("Start Trivy k8s scan")

	var allUuids uuidCollection
	reporter := NewReporter(s.log, s.api)

	var wg wait.Group
	limit := make(chan struct{}, maxParallel)

	for i := range s.targetNamespaces {
		targetNamespace := s.targetNamespaces[i]
		wg.Start(func() {
			limit <- struct{}{}

			defer func() { <-limit }()

			scanLogger := s.log.With(logz.TargetNamespace(targetNamespace))
			uuids, err := s.startPodScanForNamespace(ctx, scanLogger, targetNamespace, reporter)
			if err != nil {
				// Not logging errors for context cancelled since this is part of normal operation that can be triggered when agent configuration changes.
				if errors.Is(err, context.Canceled) {
					return
				}
				if errors.Is(err, context.DeadlineExceeded) {
					s.log.Error("Error running Trivy scan due to context timeout")
					return
				}

				s.log.Error("Error running Trivy scan", logz.Error(err))
				return
			}
			allUuids.Append(uuids)
		})
	}

	wg.Wait()

	if len(allUuids.Items()) != 0 {
		s.log.Info("Resolving no longer detected vulnerabilities in GitLab")
		err := reporter.ResolveVulnerabilities(ctx, allUuids.Items())
		s.log.Info("Resolved no longer detected vulnerabilities in GitLab")
		if err != nil {
			return fmt.Errorf("error resolving vulnerabilities: %w", err)
		}
	}

	return nil
}

func (s *scanJob) startPodScanForNamespace(ctx context.Context, scanLogger *zap.Logger, targetNamespace string, reporter *Reporter) ([]string, error) {
	// The default timeout for Trivy scan is 5 minutes, using a context.WithTimout to ensure that this function doesn't get stuck perpetually. Giving additional 5 minutes after Trivy scan is complete for parsing and transmitting vulnerability report.
	ctx, cancel := context.WithTimeout(ctx, 10*time.Minute)
	defer cancel()

	var uuids []string

	podName := fmt.Sprintf("trivy-scan-%s", targetNamespace)
	podSpec, err := s.getPodSpecForTrivyScanner(podName, targetNamespace)
	if err != nil {
		return nil, fmt.Errorf("error getting spec for Trivy Pod: %w", err)
	}

	scanLogger = scanLogger.With(logz.PodName(podName))
	scanLogger.Debug("Creating Trivy Pod")
	_, err = s.kubeClientset.CoreV1().Pods(s.gitlabAgentNamespace).Create(ctx, podSpec, metav1.CreateOptions{})
	if err != nil {
		// There could be a scenario where The previous Trivy Pod was not deleted. Delete the Pod to ensure that the next scan would be successful.
		if k8errors.IsAlreadyExists(err) {
			scanLogger.Debug("Trivy Pod already exists, deleting")
			err = s.kubeClientset.CoreV1().Pods(s.gitlabAgentNamespace).Delete(ctx, podName, metav1.DeleteOptions{})
			if err != nil {
				scanLogger.Error("Error deleting Pod", logz.Error(err))
			}
			return nil, errors.New("error creating Trivy Pod as Pod exists. Deleted Pod")
		}
		return nil, fmt.Errorf("error creating Trivy Pod: %w", err)
	}
	defer func() {
		scanLogger.Debug("Deleting Trivy Pod")
		// Using a separate context in the event that the context was cancelled.
		deleteCtx, deleteCtxCancel := context.WithTimeout(ctx, 10*time.Second)
		defer deleteCtxCancel()

		err = s.kubeClientset.CoreV1().Pods(s.gitlabAgentNamespace).Delete(deleteCtx, podName, metav1.DeleteOptions{})
		if err != nil {
			scanLogger.Error("Error deleting Pod", logz.Error(err))
		}
	}()

	scanLogger.Debug("Start watcher for Trivy Pod")
	watcher, err := s.kubeClientset.CoreV1().Pods(s.gitlabAgentNamespace).Watch(ctx, metav1.ListOptions{
		FieldSelector: fmt.Sprintf("metadata.name=%s", podName),
	})
	if err != nil {
		return nil, fmt.Errorf("error starting watcher for Trivy Pod: %w", err)
	}
	defer watcher.Stop()

	for {
		select {
		case event, ok := <-watcher.ResultChan():
			if !ok {
				return nil, errors.New("channel closed unexpectedly")
			}

			pod, ok := event.Object.(*corev1.Pod)
			if !ok {
				return nil, errors.New("watcher received unexpected object that is not a Pod")
			}

			scanLogger.Debug("pod status", logz.PodStatus(string(pod.Status.Phase)))

			if pod.Status.Phase == corev1.PodFailed {
				logs, err := s.getLogsFromPod(ctx, scanLogger, podName)
				if err != nil {
					scanLogger.Error("error getting logs from Pod Log", logz.Error(err))
				}
				// Only convert logs to string if debug level is enabled
				if scanLogger.Core().Enabled(zap.DebugLevel) {
					scanLogger.Debug("Logs from failed Pod", logz.PodLog(string(logs)))
				}
				// Log Pod failed reason
				for _, cs := range pod.Status.ContainerStatuses {
					if cs.State.Terminated != nil {
						return nil, fmt.Errorf("pod failed. Container terminated reason: %+v", cs.State.Terminated)
					}
				}

				return nil, errors.New("pod failed")
			}

			if pod.Status.Phase == corev1.PodSucceeded {

				logs, err := s.getLogsFromPod(ctx, scanLogger, podName)
				if err != nil {
					return nil, fmt.Errorf("error getting logs from Pod Log: %w", err)
				}
				consolidatedReport, err := s.logParser.ParsePodLogsToReport(logs)
				if err != nil {
					return nil, fmt.Errorf("error getting report from PodLog: %w", err)
				}

				findings := s.excludeControlledPodWorkloads(ctx, scanLogger, &consolidatedReport, targetNamespace)
				payloads, err := Convert(findings, s.agentID)
				if err != nil {
					return nil, fmt.Errorf("error processing vulnerability report: %w", err)
				}

				uuids, err = reporter.Transmit(ctx, payloads)
				if err != nil {
					return nil, fmt.Errorf("error transmitting vulnerability reports: %w", err)
				}
				return uuids, nil
			}
		case <-ctx.Done():
			scanLogger.Debug("Stopping watcher as context cancelled")
			return nil, ctx.Err()
		}
	}
}

// Exclude controlled Pods to prevent duplicate findings from being presented.
// For example a Pod controlled by a replicaset would have the same underlying image and result in duplicate findings.
// Unfortunately, we can't filter out such Pods when runnings the Trivy k8s scan at point of writing.
func (s *scanJob) excludeControlledPodWorkloads(ctx context.Context, scanLogger *zap.Logger, consolidatedReport *ConsolidatedReport, targetNamespace string) []Resource {
	findings := consolidatedReport.Findings

	var qualifiedFindings []Resource
	for _, finding := range findings {
		isControlledPod, err := s.podControlledByBuiltInWorkload(ctx, scanLogger, targetNamespace, finding.Name, finding.Kind)
		if err != nil {
			scanLogger.Error("Error checking if Pod is controlled", logz.Error(err))
			continue
		}
		if !isControlledPod {
			qualifiedFindings = append(qualifiedFindings, finding)
		}
	}
	return qualifiedFindings
}

func (s *scanJob) podControlledByBuiltInWorkload(ctx context.Context, scanLogger *zap.Logger, targetNamespace string, name string, kind string) (bool, error) {
	if kind == kindPod {
		pod, err := s.kubeClientset.CoreV1().Pods(targetNamespace).Get(ctx, name, metav1.GetOptions{})
		if err != nil {
			return false, fmt.Errorf("error getting Pod: %s. %w", name, err)
		}
		accessor, err := meta.Accessor(pod)
		if err != nil {
			return false, fmt.Errorf("error getting Pod accessor: %s. %w", name, err)
		}
		controller := metav1.GetControllerOf(accessor)
		if controller == nil {
			return false, nil
		}

		_, ok := podOwnerBuiltInWorkloadAPIVersions[controller.APIVersion]
		if ok {
			scanLogger.Sugar().Debugf("Skipping Pod %q controlled by built-in workload: %s/%s", accessor.GetName(), controller.Kind, controller.Name)
			return true, nil
		}
		return false, nil
	}
	return false, nil
}

// podOwnerBuiltInWorkloadAPIVersions contains the API versions of Pod owners that are built in workloads.
var podOwnerBuiltInWorkloadAPIVersions = map[string]struct{}{
	"v1":       {},
	"apps/v1":  {},
	"batch/v1": {},
}

func (s *scanJob) getLogsFromPod(ctx context.Context, scanLogger *zap.Logger, podName string) ([]byte, error) {
	scanLogger.Debug("Read Trivy Pod logs")
	maxBytesToRead := int64(10 * 1000 * 1000)
	req := s.kubeClientset.CoreV1().Pods(s.gitlabAgentNamespace).GetLogs(podName, &corev1.PodLogOptions{
		LimitBytes: &maxBytesToRead, //Limit logs to be read up to 10MB to prevent crashing agent
	})
	podLogs, err := req.Stream(ctx)
	if err != nil {
		return nil, fmt.Errorf("error getting logs for Trivy Pod: %w", err)
	}
	defer func() {
		err = podLogs.Close()
		if err != nil {
			scanLogger.Error("Error closing Pod logs stream", logz.Error(err))
		}
	}()

	logs, err := io.ReadAll(podLogs)
	if err != nil {
		return nil, fmt.Errorf("error reading logs for Trivy Pod: %w", err)
	}
	return logs, nil
}

func (s *scanJob) getPodSpecForTrivyScanner(podName string, targetNamespace string) (*corev1.Pod, error) {
	workloadsToScan := getWorkloadsToScanString()

	rq, err := s.parseAndValidateResourceRequirements()
	if err != nil {
		return nil, err
	}

	return &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      podName,
			Namespace: s.gitlabAgentNamespace,
		},
		Spec: corev1.PodSpec{
			ServiceAccountName: s.gitlabAgentServiceAccount,
			Containers: []corev1.Container{
				{
					Name: "trivy",
					Resources: corev1.ResourceRequirements{
						Limits: corev1.ResourceList{
							corev1.ResourceMemory: rq.limitsMemory,
							corev1.ResourceCPU:    rq.limitsCpu,
						},
						Requests: corev1.ResourceList{
							corev1.ResourceMemory: rq.requestsMemory,
							corev1.ResourceCPU:    rq.requestsCpu,
						},
					},
					Image: trivyImageRef,
					Command: []string{
						"trivy",
						"k8s",
						workloadsToScan,
						"-n",
						targetNamespace,
						"--no-progress",
						"--quiet",
						"--report=summary",
						"--scanners=vuln",
						"--db-repository",
						"registry.gitlab.com/gitlab-org/security-products/dependencies/trivy-db-glad",
						"--format",
						"json",
					},
				},
			},
			RestartPolicy: corev1.RestartPolicyNever,
		},
	}, nil
}

func (s *scanJob) parseAndValidateResourceRequirements() (resourcesQuantity, error) {
	requestsResource, err := parseAndValidateResource(s.resourceRequirements.Requests)
	if err != nil {
		return resourcesQuantity{}, err
	}

	limitsResource, err := parseAndValidateResource(s.resourceRequirements.Limits)
	if err != nil {
		return resourcesQuantity{}, err
	}
	return resourcesQuantity{
		requestsCpu:    requestsResource.Cpu,
		requestsMemory: requestsResource.Memory,
		limitsCpu:      limitsResource.Cpu,
		limitsMemory:   limitsResource.Memory,
	}, nil
}

func getWorkloadsToScanString() string {
	workloadKinds := []string{
		kindPod,
		kindReplicaSet,
		kindReplicationController,
		kindStatefulSet,
		kindDaemonSet,
		kindCronJob,
		kindJob,
	}
	return strings.Join(workloadKinds, ",")
}

func parseAndValidateResource(resources *agentcfg.Resource) (resourceQuantity, error) {
	cpu, err := resource.ParseQuantity(resources.Cpu)
	if err != nil {
		return resourceQuantity{}, err
	}

	memory, err := resource.ParseQuantity(resources.Memory)
	if err != nil {
		return resourceQuantity{}, err
	}
	return resourceQuantity{
		Cpu:    cpu,
		Memory: memory,
	}, nil
}