package main

import (
	"context"
	"errors"
	"fmt"
	"testing"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/client-go/dynamic/fake"
)

const testDigest = "sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"

func TestRetentionTagOwnership(t *testing.T) {
	t.Parallel()
	for _, test := range []struct {
		tag   string
		valid bool
	}{
		{"", true},
		{"stable", true},
		{"sha256-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", true},
		{"sha256-bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", false},
		{"sha256-release", true},
	} {
		t.Run(test.tag, func(t *testing.T) {
			image := testImage()
			image.Spec.Digests = []digestSpec{{Digest: testDigest, Tag: test.tag}}
			err := validateImage(&image)
			if (err == nil) != test.valid {
				t.Fatalf("RetentionTags.AliasAllowed: valid=%t err=%v", test.valid, err)
			}
		})
	}
}

type fakeRegistry struct {
	scanPhase string
	present   map[string]bool
	stored    map[string]bool
}

func (f *fakeRegistry) scan(context.Context, string, string, string) (string, error) {
	return f.scanPhase, nil
}

func (f *fakeRegistry) exists(_ context.Context, repository, digest string) (bool, error) {
	return f.present[repository+"@"+digest], nil
}

func (f *fakeRegistry) physicallyPresent(_ context.Context, repository, digest string) (bool, error) {
	if f.stored != nil {
		return f.stored[repository+"@"+digest], nil
	}
	return f.present[repository+"@"+digest], nil
}

func TestCleanupObservedMeansGone(t *testing.T) {
	t.Parallel()
	for _, repository := range []string{"candidates/pause", "mirror/pause"} {
		t.Run(repository, func(t *testing.T) {
			registry := &fakeRegistry{stored: map[string]bool{repository + "@" + testDigest: true}}
			instance := testController(newClient(), registry)
			image := testImage()
			image.Spec.Digests = nil
			image.Status.Digests = []digestStatus{{Digest: testDigest, Cleanup: &attemptStatus{WorkflowName: "cleanup", Phase: "Succeeded"}}}
			reconcileDesired(context.Background(), t, instance, &image)
			if len(image.Status.Digests) != 1 {
				t.Fatal("CleanupObservedMeansGone: physical manifest was forgotten")
			}
			delete(registry.stored, repository+"@"+testDigest)
			reconcileDesired(context.Background(), t, instance, &image)
			if len(image.Status.Digests) != 0 {
				t.Fatal("physically absent digest was retained")
			}
		})
	}
}

func TestCopyRetryRearmsCleanup(t *testing.T) {
	t.Parallel()
	for _, stage := range []string{prepareStage, promotionStage} {
		for _, readd := range []bool{false, true} {
			t.Run(fmt.Sprintf("%s/readd=%t", stage, readd), func(t *testing.T) {
				ctx := context.Background()
				client := newClient()
				registry := &fakeRegistry{scanPhase: scanAccepted, stored: map[string]bool{}}
				instance := testController(client, registry)
				image := testImage()
				image.Spec.Digests = nil
				image.Status.Digests = []digestStatus{{
					Digest: testDigest, Tag: "stable", Scan: scanAccepted,
					Prepare:   &attemptStatus{WorkflowName: "prepare", Phase: "Succeeded"},
					Promotion: &attemptStatus{WorkflowName: "promotion", Phase: "Succeeded"},
				}}
				status := &image.Status.Digests[0]
				status.Cleanup = instance.newAttempt(&image, status, cleanupStage, "")
				status.Cleanup.Phase = "Succeeded"
				oldCleanup := status.Cleanup.WorkflowName
				if _, err := client.Resource(workflowGVR).Namespace("mirrored-image").Create(ctx, workflowObject(oldCleanup, "Succeeded"), metav1.CreateOptions{}); err != nil {
					t.Fatal(err)
				}
				image.Metadata.Annotations = map[string]string{retryAnnotation: stage}
				for pass := 0; pass < 2; pass++ {
					handled, err := instance.reconcileRetry(ctx, &image)
					if err != nil || !handled {
						t.Fatalf("retry pass %d: handled=%t err=%v", pass, handled, err)
					}
					if pass == 0 {
						if image.Status.LastHandledRetry != "" {
							t.Fatal("RetryAcknowledgedAfterCreation")
						}
						if readd {
							image.Spec.Digests = []digestSpec{{Digest: testDigest, Tag: "edited"}}
						}
					}
				}
				if image.Status.LastHandledRetry != stage {
					t.Fatal("retry was not acknowledged")
				}
				copy := status.Prepare
				if stage == promotionStage {
					copy = status.Promotion
				}
				setPhase(ctx, t, client, copy.WorkflowName, "Succeeded")
				registry.stored["candidates/pause@"+testDigest] = true
				reconcileDesired(ctx, t, instance, &image)
				reconcileDesired(ctx, t, instance, &image)
				if status.Cleanup == nil || status.Cleanup.WorkflowName == oldCleanup || terminalPhase(status.Cleanup.Phase) {
					t.Fatal("copy retry reused completed cleanup")
				}
				fresh := status.Cleanup.WorkflowName
				if _, err := client.Resource(workflowGVR).Namespace("mirrored-image").Get(ctx, fresh, metav1.GetOptions{}); err != nil {
					t.Fatal(err)
				}
				if status.Tag != "stable" {
					t.Fatal("RetryPreservesTagIdentity")
				}
				setPhase(ctx, t, client, fresh, "Succeeded")
				reconcileDesired(ctx, t, instance, &image)
				reconcileDesired(ctx, t, instance, &image)
				if len(image.Status.Digests) != 1 {
					t.Fatal("CleanupObservedMeansGone")
				}
				delete(registry.stored, "candidates/pause@"+testDigest)
				reconcileDesired(ctx, t, instance, &image)
				if readd {
					if status.Cleanup != nil || status.Prepare != nil || status.Tag != "edited" {
						t.Fatal("re-addition did not wait for fresh cleanup")
					}
				} else if len(image.Status.Digests) != 0 {
					t.Fatal("retirement did not complete")
				}
			})
		}
	}
}

func TestDigestLifecycle(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	client := newClient()
	registry := &fakeRegistry{scanPhase: scanAccepted, present: map[string]bool{}}
	instance := testController(client, registry)
	image := testImage()

	reconcileDesired(ctx, t, instance, &image)
	status := &image.Status.Digests[0]
	if status.Prepare == nil || status.Prepare.WorkflowName == "" {
		t.Fatal("prepare workflow was not recorded")
	}
	reconcileDesired(ctx, t, instance, &image)
	setPhase(ctx, t, client, status.Prepare.WorkflowName, "Succeeded")
	reconcileDesired(ctx, t, instance, &image)
	reconcileDesired(ctx, t, instance, &image)
	if image.Status.Digests[0].Scan != scanAccepted {
		t.Fatal("candidate scan was not accepted")
	}
	reconcileDesired(ctx, t, instance, &image)
	status = &image.Status.Digests[0]
	if status.Promotion == nil || status.Promotion.WorkflowName == "" {
		t.Fatal("promotion workflow was not recorded")
	}
	reconcileDesired(ctx, t, instance, &image)
	setPhase(ctx, t, client, status.Promotion.WorkflowName, "Succeeded")
	registry.present["mirror/pause@"+testDigest] = true
	reconcileDesired(ctx, t, instance, &image)

	image.Metadata.Generation++
	image.Spec.Digests = nil
	reconcileDesired(ctx, t, instance, &image)
	status = &image.Status.Digests[0]
	if status.Cleanup == nil || status.Cleanup.WorkflowName == "" {
		t.Fatal("cleanup workflow was not recorded")
	}
	reconcileDesired(ctx, t, instance, &image)
	setPhase(ctx, t, client, status.Cleanup.WorkflowName, "Succeeded")
	delete(registry.present, "mirror/pause@"+testDigest)
	reconcileDesired(ctx, t, instance, &image)
	reconcileDesired(ctx, t, instance, &image)
	if len(image.Status.Digests) != 0 {
		t.Fatalf("cleanup retained %d digest statuses", len(image.Status.Digests))
	}
}

func TestAttemptIsRecordedFirst(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	client := newClient()
	instance := testController(client, &fakeRegistry{present: map[string]bool{}})
	image := testImage()

	reconcileDesired(ctx, t, instance, &image)
	if image.Status.Digests[0].Prepare == nil {
		t.Fatal("prepare attempt was not recorded")
	}
	workflows, err := client.Resource(workflowGVR).Namespace("mirrored-image").List(ctx, metav1.ListOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if len(workflows.Items) != 0 {
		t.Fatal("workflow existed before its status record")
	}

	reconcileDesired(ctx, t, instance, &image)
	workflows, err = client.Resource(workflowGVR).Namespace("mirrored-image").List(ctx, metav1.ListOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if len(workflows.Items) != 1 || workflows.Items[0].GetName() != image.Status.Digests[0].Prepare.WorkflowName {
		t.Fatal("recorded prepare workflow was not created")
	}
}

func TestCleanupWaitsForAbsence(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	client := newClient()
	registry := &fakeRegistry{present: map[string]bool{"candidates/pause@" + testDigest: true}}
	instance := testController(client, registry)
	image := testImage()
	image.Spec.Digests = nil
	image.Status.Digests = []digestStatus{{
		Digest:  testDigest,
		Cleanup: &attemptStatus{WorkflowName: "cleanup", Phase: "Succeeded"},
	}}

	reconcileDesired(ctx, t, instance, &image)
	if len(image.Status.Digests) != 1 || image.Status.Message != "cleanup: awaiting repository GC" {
		t.Fatal("cleanup completed while the candidate remained present")
	}
	delete(registry.present, "candidates/pause@"+testDigest)
	reconcileDesired(ctx, t, instance, &image)
	if len(image.Status.Digests) != 0 {
		t.Fatal("cleanup state remained after both repositories were absent")
	}
}

func TestReaddWaitsForCleanup(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	client := newClient()
	registry := &fakeRegistry{scanPhase: scanAccepted, present: map[string]bool{}}
	instance := testController(client, registry)
	image := testImage()
	oldPrepare := attemptName(&image, testDigest, prepareStage, image.Metadata.Generation, "")
	image.Metadata.Generation = 3
	image.Status.Digests = []digestStatus{{
		Digest:    testDigest,
		Tag:       "stable",
		Prepare:   &attemptStatus{WorkflowName: oldPrepare, Phase: "Succeeded", Generation: 1},
		Scan:      scanAccepted,
		Promotion: &attemptStatus{WorkflowName: "old-promotion", Phase: "Succeeded"},
		Cleanup:   &attemptStatus{WorkflowName: "old-cleanup", Phase: "Running"},
	}}
	workflow := workflowObject("old-cleanup", "Running")
	if _, err := client.Resource(workflowGVR).Namespace("mirrored-image").Create(ctx, workflow, metav1.CreateOptions{}); err != nil {
		t.Fatal(err)
	}

	reconcileDesired(ctx, t, instance, &image)
	if image.Status.Digests[0].Prepare.WorkflowName != oldPrepare {
		t.Fatal("prepare restarted while cleanup was active")
	}
	setPhase(ctx, t, client, "old-cleanup", "Succeeded")
	reconcileDesired(ctx, t, instance, &image)
	reconcileDesired(ctx, t, instance, &image)
	status := image.Status.Digests[0]
	if status.Cleanup != nil || status.Prepare != nil {
		t.Fatal("completed cleanup did not reset digest state")
	}
	reconcileDesired(ctx, t, instance, &image)
	if image.Status.Digests[0].Prepare == nil {
		t.Fatal("prepare did not restart after cleanup")
	}
	if image.Status.Digests[0].Prepare.WorkflowName == oldPrepare {
		t.Fatal("prepare reused the workflow from an earlier generation")
	}
}

func TestRetryCreatesSuccessor(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	client := newClient()
	registry := &fakeRegistry{scanPhase: scanAccepted, present: map[string]bool{}}
	instance := testController(client, registry)
	image := testImage()
	image.Status.Digests = []digestStatus{{
		Digest:  testDigest,
		Tag:     "stable",
		Prepare: &attemptStatus{WorkflowName: "failed-prepare", Phase: "Failed"},
	}}
	image.Metadata.Annotations = map[string]string{retryAnnotation: "failed-prepare"}

	retried, err := instance.reconcileRetry(ctx, &image)
	if err != nil {
		t.Fatal(err)
	}
	if !retried {
		t.Fatal("terminal retry was not handled")
	}
	attempt := image.Status.Digests[0].Prepare
	if attempt.RetryOf != "failed-prepare" || attempt.WorkflowName == "failed-prepare" {
		t.Fatalf("invalid successor: %#v", attempt)
	}
	if image.Status.LastHandledRetry != "" {
		t.Fatal("retry was acknowledged before the successor existed")
	}
	retried, err = instance.reconcileRetry(ctx, &image)
	if err != nil || !retried {
		t.Fatalf("create retry successor: retried=%v error=%v", retried, err)
	}
	if image.Status.LastHandledRetry != "failed-prepare" {
		t.Fatal("retry was not acknowledged after the successor existed")
	}
	setPhase(ctx, t, client, attempt.WorkflowName, "Succeeded")
	reconcileDesired(ctx, t, instance, &image)
	reconcileDesired(ctx, t, instance, &image)
	reconcileDesired(ctx, t, instance, &image)
	promotion := image.Status.Digests[0].Promotion
	if promotion == nil || promotion.WorkflowName == "" {
		t.Fatal("promotion was not created after the retry")
	}
	oldName := attemptName(&image, testDigest, promotionStage, image.Metadata.Generation)
	if promotion.WorkflowName == oldName {
		t.Fatal("promotion reused the predecessor prepare lineage")
	}
}

func TestRetryNameIgnoresGeneration(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	client := newClient()
	instance := testController(client, &fakeRegistry{present: map[string]bool{}})
	image := testImage()
	image.Status.Digests = []digestStatus{{
		Digest:  testDigest,
		Prepare: &attemptStatus{WorkflowName: "failed-prepare", Phase: "Failed"},
	}}
	image.Metadata.Annotations = map[string]string{retryAnnotation: "failed-prepare"}

	if _, err := instance.reconcileRetry(ctx, &image); err != nil {
		t.Fatal(err)
	}
	first := image.Status.Digests[0].Prepare.WorkflowName
	image.Metadata.Generation++
	image.Status.Digests[0].Prepare = &attemptStatus{WorkflowName: "failed-prepare", Phase: "Failed"}
	image.Status.LastHandledRetry = ""
	if _, err := instance.reconcileRetry(ctx, &image); err != nil {
		t.Fatal(err)
	}
	if second := image.Status.Digests[0].Prepare.WorkflowName; second != first {
		t.Fatalf("retry generation changed successor from %q to %q", first, second)
	}
}

func TestRunStopsWhenIdle(t *testing.T) {
	t.Parallel()
	ctx, cancel := context.WithCancel(context.Background())
	client := newClient()
	instance := testController(client, &fakeRegistry{present: map[string]bool{}})
	done := make(chan error, 1)
	go func() {
		done <- instance.run(ctx)
	}()
	time.Sleep(20 * time.Millisecond)
	cancel()
	select {
	case err := <-done:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("run returned %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("run did not stop")
	}
}

func TestRejectedCandidateIsNotPromoted(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	client := newClient()
	registry := &fakeRegistry{scanPhase: scanRejected, present: map[string]bool{}}
	instance := testController(client, registry)
	image := testImage()
	image.Status.Digests = []digestStatus{{
		Digest:  testDigest,
		Tag:     "stable",
		Prepare: &attemptStatus{WorkflowName: "prepare", Phase: "Succeeded"},
	}}

	reconcileDesired(ctx, t, instance, &image)
	status := image.Status.Digests[0]
	if status.Scan != scanRejected || status.Promotion != nil {
		t.Fatal("rejected candidate was promoted")
	}
	if image.Status.Message != "scan: Rejected" {
		t.Fatalf("unexpected status message %q", image.Status.Message)
	}
}

func TestTaglessDeclarationIsValid(t *testing.T) {
	t.Parallel()
	image := testImage()
	image.Spec.Digests[0].Tag = ""
	if err := validateImage(&image); err != nil {
		t.Fatal(err)
	}
}

func TestTagEditPreservesState(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	client := newClient()
	registry := &fakeRegistry{present: map[string]bool{"mirror/pause@" + testDigest: true}}
	instance := testController(client, registry)
	image := testImage()
	image.Spec.Digests[0].Tag = "new"
	image.Status.Digests = []digestStatus{{
		Digest:    testDigest,
		Tag:       "stable",
		Prepare:   &attemptStatus{WorkflowName: "prepare", Phase: "Succeeded"},
		Scan:      scanAccepted,
		Promotion: &attemptStatus{WorkflowName: "promotion", Phase: "Succeeded"},
	}}

	reconcileDesired(ctx, t, instance, &image)
	status := image.Status.Digests[0]
	if status.Tag != "stable" || status.Cleanup != nil {
		t.Fatalf("tag edit changed managed state: %#v", status)
	}
	workflows, err := client.Resource(workflowGVR).Namespace("mirrored-image").List(ctx, metav1.ListOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if len(workflows.Items) != 0 {
		t.Fatal("tag edit created a workflow")
	}
}

func TestFailureMessage(t *testing.T) {
	t.Parallel()
	ctx := context.Background()
	client := newClient()
	registry := &fakeRegistry{present: map[string]bool{}}
	instance := testController(client, registry)
	image := testImage()
	image.Status.Digests = []digestStatus{{
		Digest:  testDigest,
		Tag:     "stable",
		Prepare: &attemptStatus{WorkflowName: "failed-prepare", Phase: "Failed"},
	}}

	reconcileDesired(ctx, t, instance, &image)
	if image.Status.Message != "prepare: Failed" {
		t.Fatalf("unexpected status message %q", image.Status.Message)
	}
}

func TestPromotionRescan(t *testing.T) {
	for _, phase := range []string{scanPending, scanRejected, scanAccepted} {
		t.Run(phase, func(t *testing.T) {
			ctx := context.Background()
			client := newClient()
			registry := &fakeRegistry{scanPhase: phase}
			instance := testController(client, registry)
			image := testImage()
			image.Status.Digests = []digestStatus{{
				Digest: testDigest, Scan: scanAccepted,
				Prepare:   &attemptStatus{WorkflowName: "prepare", Phase: "Succeeded"},
				Promotion: &attemptStatus{WorkflowName: "promotion", Phase: "Pending"},
			}}
			err := instance.reconcileDesired(ctx, &image)
			if (err == nil) != (phase == scanAccepted) {
				t.Fatalf("scan %s: %v", phase, err)
			}
			workflows, err := client.Resource(workflowGVR).Namespace(image.Metadata.Namespace).List(ctx, metav1.ListOptions{})
			if err != nil {
				t.Fatal(err)
			}
			if (len(workflows.Items) == 1) != (phase == scanAccepted) {
				t.Fatalf("scan %s created %d workflows", phase, len(workflows.Items))
			}
			if phase == scanAccepted {
				registry.scanPhase = scanPending
				if err := instance.reconcileDesired(ctx, &image); err != nil {
					t.Fatalf("existing workflow observation requires a scan: %v", err)
				}
			}
		})
	}
}

func testController(client *fake.FakeDynamicClient, registry registryReader) *controller {
	return &controller{client: client, registry: registry, config: controllerConfig{
		prepareTemplate:   "image-prepare",
		promotionTemplate: "image-promote",
		cleanupTemplate:   "image-cleanup",
	}}
}

func newClient(objects ...runtime.Object) *fake.FakeDynamicClient {
	return fake.NewSimpleDynamicClientWithCustomListKinds(
		runtime.NewScheme(),
		map[schema.GroupVersionResource]string{
			imageGVR:    "MirroredImageList",
			workflowGVR: "WorkflowList",
		},
		objects...,
	)
}

func testImage() mirroredImage {
	return mirroredImage{
		Metadata: objectMetadata{Name: "pause", Namespace: "mirrored-image", UID: "image-uid", Generation: 1},
		Spec: mirroredImageSpec{
			Upstream: "registry.k8s.io/pause",
			Policy:   scanPolicy{RejectAtOrAbove: "High"},
			Digests:  []digestSpec{{Digest: testDigest, Tag: "stable"}},
		},
	}
}

func reconcileDesired(ctx context.Context, t *testing.T, instance *controller, image *mirroredImage) {
	t.Helper()
	if err := instance.reconcileDesired(ctx, image); err != nil {
		t.Fatal(err)
	}
}

func setPhase(ctx context.Context, t *testing.T, client *fake.FakeDynamicClient, name, phase string) {
	t.Helper()
	resource := client.Resource(workflowGVR).Namespace("mirrored-image")
	object, err := resource.Get(ctx, name, metav1.GetOptions{})
	if err != nil {
		t.Fatal(err)
	}
	if err := unstructured.SetNestedField(object.Object, phase, "status", "phase"); err != nil {
		t.Fatal(err)
	}
	if _, err := resource.Update(ctx, object, metav1.UpdateOptions{}); err != nil {
		t.Fatal(err)
	}
}

func workflowObject(name, phase string) *unstructured.Unstructured {
	return &unstructured.Unstructured{Object: map[string]any{
		"apiVersion": "argoproj.io/v1alpha1",
		"kind":       "Workflow",
		"metadata":   map[string]any{"name": name, "namespace": "mirrored-image"},
		"status":     map[string]any{"phase": phase},
	}}
}