package agentk

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"sync/atomic"
	"time"

	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/api"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/module/starboard_vulnerability/agentk/resources"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/internal/tool/logz"
	"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v18/pkg/agentcfg"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
)

const (
	maxParallel int = 10 // Trivy scanner Pods batch size

	// Fatal image pull errors which indicate that an OCS job should not be retried.
	// Source: https://github.com/kubernetes/kubernetes/blob/605e94f6dfa9e98831e7ed0932dd18cd75d99e46/pkg/kubelet/images/types.go#L27-L42

	// ErrImagePullBackOff - Container image pull failed, kubelet is backing off image pull
	ErrImagePullBackOff = "ImagePullBackOff"
	// ErrImageInspect - Unable to inspect imag
	ErrImageInspect = "ImageInspectError"
	// ErrImagePull - General image pull error
	ErrImagePull = "ErrImagePull"
	// ErrImageNeverPull - Required Image is absent on host and PullPolicy is NeverPullImage
	ErrImageNeverPull = "ErrImageNeverPull"
	// ErrInvalidImageName - Unable to parse the image name.
	ErrInvalidImageName = "InvalidImageName"
	scanStatusRunning   = "Running"
	scanStatusSuccess   = "Success"
	scanStatusFailed    = "Failed"
)

type Scanner interface {
	scan(ctx context.Context, namespaceScanner NamespaceScanner, reporter Reporter) error
}

type scanner struct {
	log                  *slog.Logger
	kubeClientset        kubernetes.Interface
	gitlabAgentNamespace string
	agentID              api.AgentKey
	resourceRequirements *agentcfg.ResourceRequirements
	targetNamespaces     []string
	persistOcsStatus     bool
	scannerTimeout       time.Duration
	reportMaxSizeBytes   uint64
	resourceTypes        []string
}

func newScanner(
	log *slog.Logger,
	kubeClientset kubernetes.Interface,
	gitlabAgentNamespace string,
	agentID api.AgentKey,
	resourceRequirements *agentcfg.ResourceRequirements,
	targetNamespaces []string,
	persistOcsStatus bool,
	scannerTimeout time.Duration,
	reportMaxSizeBytes uint64,
	resourceTypes []string,
) *scanner {
	return &scanner{
		log:                  log,
		kubeClientset:        kubeClientset,
		gitlabAgentNamespace: gitlabAgentNamespace,
		agentID:              agentID,
		resourceRequirements: resourceRequirements,
		targetNamespaces:     targetNamespaces,
		persistOcsStatus:     persistOcsStatus,
		scannerTimeout:       scannerTimeout,
		reportMaxSizeBytes:   reportMaxSizeBytes,
		resourceTypes:        resourceTypes,
	}
}

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 *scanner) scan(ctx context.Context, namespaceScanner NamespaceScanner, reporter Reporter) error {
	s.log.Info("Start Trivy k8s scan")

	// Default vuln resolution to true
	var shouldResolveVulns atomic.Bool
	shouldResolveVulns.Store(true)

	statusRecorder := statusRecorder{
		configMaps:           s.kubeClientset.CoreV1().ConfigMaps(s.gitlabAgentNamespace),
		gitlabAgentNamespace: s.gitlabAgentNamespace,
		agentKey:             s.agentID,
		persistOcsStatus:     s.persistOcsStatus,
	}
	err := statusRecorder.initScanStatus(ctx)
	if err != nil {
		return fmt.Errorf("error initializing OCS status config map: %w", err)
	}

	var allUUIDs uuidCollection

	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))
			statusRecorder.recordScanStart(ctx, scanLogger, targetNamespace)

			scanningManager := scanningManager{
				kubeClientset: s.kubeClientset,
				resourcesManager: &resources.Manager{
					Requirements: s.resourceRequirements,
				},
				agentKey:             s.agentID,
				gitlabAgentNamespace: s.gitlabAgentNamespace,
				scanLogger:           scanLogger,
				scannerTimeout:       s.scannerTimeout,
				reportMaxSizeBytes:   s.reportMaxSizeBytes,
				resourceTypes:        s.resourceTypes,
			}
			uuids, err := namespaceScanner.scan(ctx, scanLogger, targetNamespace, reporter, &scanningManager)
			if err != nil {
				// Always disable resolve vulns if an error is encountered
				shouldResolveVulns.Store(false)
				statusRecorder.recordScanEnd(ctx, scanLogger, targetNamespace, 0, scanStatusFailed)
				// Not logging errors for context canceled 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
			}
			statusRecorder.recordScanEnd(ctx, scanLogger, targetNamespace, len(uuids), scanStatusSuccess)
			allUUIDs.Append(uuids)
		})
	}

	wg.Wait()

	if !shouldResolveVulns.Load() {
		s.log.Warn("Skipping vulnerability resolution due to errors detected in one or more namespace scans")
		return nil
	}

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

	return nil
}
