package main

import (
	"context"
	"crypto/sha256"
	"encoding/base32"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"reflect"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/dynamic"
	"k8s.io/client-go/dynamic/dynamicinformer"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/retry"
	"k8s.io/client-go/util/workqueue"
)

const (
	prepareStage   = "prepare"
	promotionStage = "promotion"
	cleanupStage   = "cleanup"
	scanAccepted   = "Accepted"
	scanPending    = "Pending"
	scanRejected   = "Rejected"
)

var (
	imageGVR         = schema.GroupVersionResource{Group: "mirror.quorate.dev", Version: "v1alpha1", Resource: "mirroredimages"}
	workflowGVR      = schema.GroupVersionResource{Group: "argoproj.io", Version: "v1alpha1", Resource: "workflows"}
	digestRule       = regexp.MustCompile(`^sha256:[0-9a-f]{64}$`)
	tagRule          = regexp.MustCompile(`^[A-Za-z0-9_][A-Za-z0-9_.-]{0,127}$`)
	retentionTagRule = regexp.MustCompile(`^sha256-[0-9a-f]{64}$`)
)

type controllerConfig struct {
	namespace         string
	pollInterval      time.Duration
	prepareTemplate   string
	promotionTemplate string
	cleanupTemplate   string
}

type controller struct {
	client   dynamic.Interface
	registry registryReader
	config   controllerConfig
}

type registryReader interface {
	scan(context.Context, string, string, string) (string, error)
	exists(context.Context, string, string) (bool, error)
	physicallyPresent(context.Context, string, string) (bool, error)
}

type retryRef struct {
	digest    int
	stage     string
	successor bool
}

type imageKey struct {
	namespace string
	name      string
}

func (c *controller) run(ctx context.Context) error {
	factory := dynamicinformer.NewFilteredDynamicSharedInformerFactory(c.client, 0, c.config.namespace, nil)
	informer := factory.ForResource(imageGVR).Informer()
	queue := workqueue.NewTypedRateLimitingQueueWithConfig(
		workqueue.DefaultTypedControllerRateLimiter[imageKey](),
		workqueue.TypedRateLimitingQueueConfig[imageKey]{Name: "mirrored-image"},
	)
	defer queue.ShutDown()

	enqueue := func(object any) {
		image, ok := object.(*unstructured.Unstructured)
		if ok {
			queue.Add(imageKey{namespace: image.GetNamespace(), name: image.GetName()})
		}
	}
	_, err := informer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc:    enqueue,
		UpdateFunc: func(_, newObject any) { enqueue(newObject) },
	})
	if err != nil {
		return fmt.Errorf("register MirroredImage informer: %w", err)
	}

	factory.Start(ctx.Done())
	if !cache.WaitForCacheSync(ctx.Done(), informer.HasSynced) {
		if err := ctx.Err(); err != nil {
			return err
		}
		return errors.New("MirroredImage informer did not synchronize")
	}
	go func() {
		<-ctx.Done()
		queue.ShutDown()
	}()

	for c.processNext(ctx, queue) {
	}
	return ctx.Err()
}

func (c *controller) processNext(ctx context.Context, queue workqueue.TypedRateLimitingInterface[imageKey]) bool {
	key, shutdown := queue.Get()
	if shutdown {
		return false
	}
	defer queue.Done(key)

	found, err := c.reconcileKey(ctx, key.namespace, key.name)
	if err != nil {
		slog.Error("reconcile MirroredImage", "namespace", key.namespace, "name", key.name, "error", err)
		c.recordError(ctx, key.namespace, key.name, err)
		queue.AddRateLimited(key)
		return true
	}
	queue.Forget(key)
	if found {
		queue.AddAfter(key, c.config.pollInterval)
	}
	return true
}

func (c *controller) reconcileKey(ctx context.Context, namespace, name string) (bool, error) {
	found := false
	err := retry.RetryOnConflict(retry.DefaultRetry, func() error {
		object, err := c.client.Resource(imageGVR).Namespace(namespace).Get(ctx, name, metav1.GetOptions{})
		if apierrors.IsNotFound(err) {
			return nil
		}
		if err != nil {
			return err
		}
		found = true
		return c.reconcile(ctx, object)
	})
	return found, err
}

func (c *controller) reconcile(ctx context.Context, object *unstructured.Unstructured) error {
	image, err := decodeImage(object)
	if err != nil {
		return err
	}
	if image.Metadata.DeletionTimestamp != nil {
		return nil
	}
	if err := validateImage(&image); err != nil {
		return err
	}
	retried, err := c.reconcileRetry(ctx, &image)
	if err != nil {
		return err
	}
	if !retried {
		if err := c.reconcileDesired(ctx, &image); err != nil {
			return err
		}
	}
	image.Status.ObservedGeneration = image.Metadata.Generation
	return c.updateStatus(ctx, object, &image)
}

func (c *controller) reconcileRetry(ctx context.Context, image *mirroredImage) (bool, error) {
	request := image.Metadata.Annotations[retryAnnotation]
	if request == "" || request == image.Status.LastHandledRetry {
		return false, nil
	}
	target, found := retryTarget(image, request)
	if !found {
		return false, fmt.Errorf("retry Workflow %q is not a recorded attempt", request)
	}
	attempt := attemptAt(image, target)
	if target.successor {
		if otherAttemptActive(image, attempt.WorkflowName) {
			return false, nil
		}
		status := &image.Status.Digests[target.digest]
		if err := c.advanceAttempt(ctx, image, status, target.stage, attempt); err != nil {
			return false, err
		}
		image.Status.LastHandledRetry = request
		image.Status.Message = "retry " + target.stage + ": " + phaseOrPending(attempt.Phase)
		return true, nil
	}
	if otherAttemptActive(image, request) || !terminalPhase(attempt.Phase) {
		return false, nil
	}
	status := &image.Status.Digests[target.digest]
	successor := c.newAttempt(image, status, target.stage, request)
	switch target.stage {
	case prepareStage:
		status.Prepare = successor
		status.Scan = ""
		status.Promotion = nil
	case promotionStage:
		status.Promotion = successor
	case cleanupStage:
		status.Cleanup = successor
	default:
		return false, fmt.Errorf("unknown retry stage %q", target.stage)
	}
	if target.stage != cleanupStage && status.Cleanup != nil {
		status.Cleanup = c.newAttempt(image, status, cleanupStage, "")
	}
	image.Status.Message = "retry " + target.stage + ": Pending"
	return true, nil
}

func (c *controller) reconcileDesired(ctx context.Context, image *mirroredImage) error {
	desired := make(map[string]digestSpec, len(image.Spec.Digests))
	failedMessage := ""
	for _, item := range image.Spec.Digests {
		desired[item.Digest] = item
		index := digestIndex(image.Status.Digests, item.Digest)
		if index < 0 {
			image.Status.Digests = append(image.Status.Digests, digestStatus{Digest: item.Digest, Tag: item.Tag})
		}
	}

	for index := 0; index < len(image.Status.Digests); index++ {
		status := &image.Status.Digests[index]
		item, wanted := desired[status.Digest]
		retiring := status.Cleanup != nil || !wanted
		if retiring {
			active := activeAttempt(status)
			if active != nil {
				stage := activeStage(status)
				if err := c.advanceAttempt(ctx, image, status, stage, active); err != nil {
					return err
				}
				image.Status.Message = "awaiting " + stage + " before cleanup: " + phaseOrPending(active.Phase)
				return nil
			}
			switch {
			case status.Cleanup == nil:
				status.Cleanup = c.newAttempt(image, status, cleanupStage, "")
				image.Status.Message = "cleanup: Pending"
				return nil
			case !terminalPhase(status.Cleanup.Phase):
				if err := c.advanceAttempt(ctx, image, status, cleanupStage, status.Cleanup); err != nil {
					return err
				}
				image.Status.Message = "cleanup: " + phaseOrPending(status.Cleanup.Phase)
				return nil
			case status.Cleanup.Phase != "Succeeded":
				if failedMessage == "" {
					failedMessage = "cleanup: " + status.Cleanup.Phase
				}
				continue
			default:
				complete, err := c.cleanupComplete(ctx, image.Metadata.Name, status.Digest)
				if err != nil {
					return err
				}
				if !complete {
					image.Status.Message = "cleanup: awaiting repository GC"
					return nil
				}
				if wanted {
					*status = digestStatus{Digest: item.Digest, Tag: item.Tag}
					image.Status.Message = "cleanup completed before remirroring"
					return nil
				}
				image.Status.Digests = append(image.Status.Digests[:index], image.Status.Digests[index+1:]...)
				image.Status.Message = "cleanup completed"
				return nil
			}
		}

		switch {
		case status.Cleanup != nil:
			return fmt.Errorf("digest %q has cleanup state while desired", status.Digest)
		case status.Prepare == nil:
			status.Prepare = c.newAttempt(image, status, prepareStage, "")
			image.Status.Message = "prepare: Pending"
			return nil
		case !terminalPhase(status.Prepare.Phase):
			if err := c.advanceAttempt(ctx, image, status, prepareStage, status.Prepare); err != nil {
				return err
			}
			image.Status.Message = "prepare: " + phaseOrPending(status.Prepare.Phase)
			return nil
		case status.Prepare.Phase != "Succeeded":
			if failedMessage == "" {
				failedMessage = "prepare: " + status.Prepare.Phase
			}
			continue
		case status.Scan == scanRejected:
			if failedMessage == "" {
				failedMessage = "scan: " + status.Scan
			}
			continue
		case status.Scan != scanAccepted:
			phase, err := c.registry.scan(ctx, "candidates/"+image.Metadata.Name, status.Digest, image.Spec.Policy.RejectAtOrAbove)
			if err != nil {
				return err
			}
			status.Scan = phase
			if phase == scanRejected {
				if failedMessage == "" {
					failedMessage = "scan: " + phase
				}
				continue
			}
			image.Status.Message = "scan: " + phase
			return nil
		case status.Promotion == nil:
			status.Promotion = c.newAttempt(image, status, promotionStage, "")
			image.Status.Message = "promotion: Pending"
			return nil
		case !terminalPhase(status.Promotion.Phase):
			if err := c.advanceAttempt(ctx, image, status, promotionStage, status.Promotion); err != nil {
				return err
			}
			image.Status.Message = "promotion: " + phaseOrPending(status.Promotion.Phase)
			return nil
		case status.Promotion.Phase != "Succeeded":
			if failedMessage == "" {
				failedMessage = "promotion: " + status.Promotion.Phase
			}
		default:
			present, err := c.registry.exists(ctx, "mirror/"+image.Metadata.Name, status.Digest)
			if err != nil {
				return err
			}
			if !present {
				image.Status.Message = "promotion: awaiting mirror"
				return nil
			}
		}
	}

	if failedMessage == "" {
		image.Status.Message = "desired digests reconciled"
	} else {
		image.Status.Message = failedMessage
	}
	return nil
}

func (c *controller) newAttempt(image *mirroredImage, status *digestStatus, stage, retryOf string) *attemptStatus {
	attempt := &attemptStatus{RetryOf: retryOf, Phase: "Pending", Generation: image.Metadata.Generation}
	lineage := []string{retryOf}
	if stage == promotionStage && status.Prepare != nil {
		lineage = append(lineage, status.Prepare.WorkflowName)
	}
	if stage == cleanupStage {
		for _, copy := range []*attemptStatus{status.Prepare, status.Promotion} {
			if copy != nil {
				lineage = append(lineage, copy.WorkflowName)
			}
		}
	}
	attempt.WorkflowName = attemptName(image, status.Digest, stage, attempt.Generation, lineage...)
	return attempt
}

func (c *controller) advanceAttempt(
	ctx context.Context,
	image *mirroredImage,
	status *digestStatus,
	stage string,
	attempt *attemptStatus,
) error {
	if attempt.WorkflowName == "" {
		return errors.New("workflow attempt has no name")
	}
	template, parameters, err := c.workflowInput(image, status, stage)
	if err != nil {
		return err
	}
	if err := c.ensureWorkflow(ctx, image, stage, attempt.WorkflowName, template, parameters); err != nil {
		return err
	}
	phase, err := c.workflowPhase(ctx, image.Metadata.Namespace, attempt.WorkflowName)
	if err != nil {
		return err
	}
	attempt.Phase = phaseOrPending(phase)
	return nil
}

func (c *controller) cleanupComplete(ctx context.Context, name, digest string) (bool, error) {
	candidate, err := c.registry.physicallyPresent(ctx, "candidates/"+name, digest)
	if err != nil {
		return false, err
	}
	mirror, err := c.registry.physicallyPresent(ctx, "mirror/"+name, digest)
	if err != nil {
		return false, err
	}
	return !candidate && !mirror, nil
}

func (c *controller) workflowInput(
	image *mirroredImage,
	status *digestStatus,
	stage string,
) (string, map[string]string, error) {
	parameters := map[string]string{
		"digest": status.Digest,
	}
	switch stage {
	case prepareStage:
		parameters["tag"] = status.Tag
		parameters["source-repository"] = image.Spec.Upstream
		parameters["target-repository"] = "candidates/" + image.Metadata.Name
		return c.config.prepareTemplate, parameters, nil
	case promotionStage:
		parameters["tag"] = status.Tag
		parameters["source-repository"] = "candidates/" + image.Metadata.Name
		parameters["target-repository"] = "mirror/" + image.Metadata.Name
		return c.config.promotionTemplate, parameters, nil
	case cleanupStage:
		parameters["candidate-repository"] = "candidates/" + image.Metadata.Name
		parameters["mirror-repository"] = "mirror/" + image.Metadata.Name
		return c.config.cleanupTemplate, parameters, nil
	default:
		return "", nil, fmt.Errorf("unknown Workflow stage %q", stage)
	}
}

func (c *controller) ensureWorkflow(
	ctx context.Context,
	image *mirroredImage,
	stage, name, template string,
	parameters map[string]string,
) error {
	if stage == promotionStage {
		_, err := c.client.Resource(workflowGVR).Namespace(image.Metadata.Namespace).Get(ctx, name, metav1.GetOptions{})
		if err == nil {
			return nil
		}
		if !apierrors.IsNotFound(err) {
			return err
		}
		phase, err := c.registry.scan(ctx, "candidates/"+image.Metadata.Name, parameters["digest"], image.Spec.Policy.RejectAtOrAbove)
		if err != nil {
			return err
		}
		if phase != scanAccepted {
			return fmt.Errorf("promotion scan: %s", phase)
		}
	}
	names := make([]string, 0, len(parameters))
	for name := range parameters {
		names = append(names, name)
	}
	sort.Strings(names)
	arguments := make([]any, 0, len(names))
	for _, name := range names {
		arguments = append(arguments, map[string]any{"name": name, "value": parameters[name]})
	}
	workflow := &unstructured.Unstructured{Object: map[string]any{
		"apiVersion": "argoproj.io/v1alpha1",
		"kind":       "Workflow",
		"metadata": map[string]any{
			"name":      name,
			"namespace": image.Metadata.Namespace,
			"labels": map[string]any{
				"app.kubernetes.io/name":     "mirrored-image",
				"app.kubernetes.io/instance": image.Metadata.Name,
				"mirror.quorate.dev/purpose": stage,
			},
			"ownerReferences": []any{map[string]any{
				"apiVersion":         "mirror.quorate.dev/v1alpha1",
				"kind":               "MirroredImage",
				"name":               image.Metadata.Name,
				"uid":                image.Metadata.UID,
				"controller":         true,
				"blockOwnerDeletion": false,
			}},
		},
		"spec": map[string]any{
			"workflowTemplateRef": map[string]any{"name": template},
			"arguments":           map[string]any{"parameters": arguments},
			"synchronization": map[string]any{
				"mutexes": []any{map[string]any{"name": workflowName("lock", image.Metadata.UID)}},
			},
		},
	}}
	_, err := c.client.Resource(workflowGVR).Namespace(image.Metadata.Namespace).Create(ctx, workflow, metav1.CreateOptions{})
	if apierrors.IsAlreadyExists(err) {
		return nil
	}
	return err
}

func (c *controller) workflowPhase(ctx context.Context, namespace, name string) (string, error) {
	object, err := c.client.Resource(workflowGVR).Namespace(namespace).Get(ctx, name, metav1.GetOptions{})
	if err != nil {
		return "", err
	}
	phase, _, _ := unstructured.NestedString(object.Object, "status", "phase")
	return phase, nil
}

func (c *controller) updateStatus(ctx context.Context, object *unstructured.Unstructured, image *mirroredImage) error {
	status, err := toMap(image.Status)
	if err != nil {
		return err
	}
	current, _, err := unstructured.NestedMap(object.Object, "status")
	if err != nil {
		return err
	}
	if reflect.DeepEqual(current, status) {
		return nil
	}
	updated := object.DeepCopy()
	updated.Object["status"] = status
	_, err = c.client.Resource(imageGVR).Namespace(updated.GetNamespace()).UpdateStatus(ctx, updated, metav1.UpdateOptions{})
	if apierrors.IsNotFound(err) {
		return nil
	}
	return err
}

func (c *controller) recordError(ctx context.Context, namespace, name string, reconcileError error) {
	object, err := c.client.Resource(imageGVR).Namespace(namespace).Get(ctx, name, metav1.GetOptions{})
	if err != nil {
		return
	}
	image, err := decodeImage(object)
	if err != nil {
		return
	}
	image.Status.Message = reconcileError.Error()
	image.Status.ObservedGeneration = image.Metadata.Generation
	_ = c.updateStatus(ctx, object, &image)
}

func decodeImage(object *unstructured.Unstructured) (mirroredImage, error) {
	content, err := json.Marshal(object.Object)
	if err != nil {
		return mirroredImage{}, err
	}
	var image mirroredImage
	if err := json.Unmarshal(content, &image); err != nil {
		return mirroredImage{}, err
	}
	return image, nil
}

func validateImage(image *mirroredImage) error {
	if image.Spec.Upstream == "" {
		return errors.New("upstream repository is empty")
	}
	if image.Spec.Policy.RejectAtOrAbove == "" {
		return errors.New("severity threshold is empty")
	}
	seen := make(map[string]struct{}, len(image.Spec.Digests))
	for index, item := range image.Spec.Digests {
		if !digestRule.MatchString(item.Digest) {
			return fmt.Errorf("invalid digest at index %d", index)
		}
		if item.Tag != "" && !tagRule.MatchString(item.Tag) {
			return fmt.Errorf("invalid tag at index %d", index)
		}
		if retentionTagRule.MatchString(item.Tag) && item.Tag != "sha256-"+strings.TrimPrefix(item.Digest, "sha256:") {
			return fmt.Errorf("tag at index %d is owned by another digest", index)
		}
		if _, found := seen[item.Digest]; found {
			return fmt.Errorf("duplicate digest at index %d", index)
		}
		seen[item.Digest] = struct{}{}
	}
	return nil
}

func retryTarget(image *mirroredImage, workflow string) (retryRef, bool) {
	for index := range image.Status.Digests {
		status := &image.Status.Digests[index]
		if status.Prepare != nil {
			if status.Prepare.WorkflowName == workflow {
				return retryRef{digest: index, stage: prepareStage}, true
			}
			if status.Prepare.RetryOf == workflow {
				return retryRef{digest: index, stage: prepareStage, successor: true}, true
			}
		}
		if status.Promotion != nil {
			if status.Promotion.WorkflowName == workflow {
				return retryRef{digest: index, stage: promotionStage}, true
			}
			if status.Promotion.RetryOf == workflow {
				return retryRef{digest: index, stage: promotionStage, successor: true}, true
			}
		}
		if status.Cleanup != nil {
			if status.Cleanup.WorkflowName == workflow {
				return retryRef{digest: index, stage: cleanupStage}, true
			}
			if status.Cleanup.RetryOf == workflow {
				return retryRef{digest: index, stage: cleanupStage, successor: true}, true
			}
		}
	}
	return retryRef{}, false
}

func attemptAt(image *mirroredImage, target retryRef) *attemptStatus {
	status := &image.Status.Digests[target.digest]
	if target.stage == prepareStage {
		return status.Prepare
	}
	if target.stage == promotionStage {
		return status.Promotion
	}
	return status.Cleanup
}

func otherAttemptActive(image *mirroredImage, excluded string) bool {
	for index := range image.Status.Digests {
		status := &image.Status.Digests[index]
		for _, attempt := range []*attemptStatus{status.Prepare, status.Promotion, status.Cleanup} {
			if attempt == status.Cleanup && attempt != nil && attempt.Phase == "Pending" && activeAttempt(status) != nil {
				continue
			}
			if attempt != nil && attempt.WorkflowName != excluded && !terminalPhase(attempt.Phase) {
				return true
			}
		}
	}
	return false
}

func activeAttempt(status *digestStatus) *attemptStatus {
	for _, attempt := range []*attemptStatus{status.Prepare, status.Promotion} {
		if attempt != nil && !terminalPhase(attempt.Phase) {
			return attempt
		}
	}
	return nil
}

func activeStage(status *digestStatus) string {
	if status.Prepare != nil && !terminalPhase(status.Prepare.Phase) {
		return prepareStage
	}
	return promotionStage
}

func digestIndex(statuses []digestStatus, digest string) int {
	for index := range statuses {
		if statuses[index].Digest == digest {
			return index
		}
	}
	return -1
}

func toMap(value any) (map[string]any, error) {
	content, err := json.Marshal(value)
	if err != nil {
		return nil, err
	}
	var result map[string]any
	err = json.Unmarshal(content, &result)
	return result, err
}

func workflowName(purpose string, fields ...string) string {
	digest := sha256.Sum256([]byte("v1\x00" + purpose + "\x00" + strings.Join(fields, "\x00")))
	encoded := base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(digest[:15])
	return "mi-" + purpose + "-" + strings.ToLower(encoded)
}

func attemptName(image *mirroredImage, digest, stage string, generation int64, lineage ...string) string {
	fields := []string{image.Metadata.UID, digest}
	lineageFound := false
	for _, value := range lineage {
		if value != "" {
			fields = append(fields, value)
			lineageFound = true
		}
	}
	if !lineageFound {
		fields = append(fields, strconv.FormatInt(generation, 10))
	}
	return workflowName(stage, fields...)
}

func terminalPhase(phase string) bool {
	return phase == "Succeeded" || phase == "Failed" || phase == "Error"
}

func phaseOrPending(phase string) string {
	if phase == "" {
		return "Pending"
	}
	return phase
}