package release

import (
	"context"
	"errors"
	"math"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

type countedWorkerStore struct {
	WorkerControlStore
	puts atomic.Int64
}

func (s *countedWorkerStore) Put(ctx context.Context, key string, body []byte, etag string) (string, error) {
	s.puts.Add(1)
	return s.WorkerControlStore.Put(ctx, key, body, etag)
}

func TestWorkerAssignmentPublicationAndDisableShareCAS(t *testing.T) {
	for _, publicationFirst := range []bool{false, true} {
		t.Run(map[bool]string{false: "disable-first", true: "publication-first"}[publicationFirst], func(t *testing.T) {
			priv, pub := testKeys(t)
			sg := RootSigner(priv)
			store := NewDirStore(t.TempDir())
			ctx := t.Context()
			initial := workerTestAssignment()
			if _, err := CreateWorkerAssignment(ctx, store, sg, initial); err != nil {
				t.Fatal(err)
			}
			publisher, err := ReadWorkerAssignment(ctx, store, pub, initial.Identity)
			if err != nil {
				t.Fatal(err)
			}
			revoker, err := ReadWorkerAssignment(ctx, store, pub, initial.Identity)
			if err != nil {
				t.Fatal(err)
			}
			published := initial
			published.Revision++
			ref := workerTestRef(7)
			published.Preflight = &ref
			disabled := initial
			disabled.Revision++
			disabled.Disabled = true
			first, second := publisher, revoker
			firstAfter, secondAfter := published, disabled
			if !publicationFirst {
				first, second = revoker, publisher
				firstAfter, secondAfter = disabled, published
			}
			if _, err := AdvanceWorkerAssignment(ctx, store, sg, first, firstAfter); err != nil {
				t.Fatal(err)
			}
			if _, err := AdvanceWorkerAssignment(ctx, store, sg, second, secondAfter); !errors.Is(err, ErrCASConflict) {
				t.Fatalf("stale competing update: %v", err)
			}
			latest, err := ReadWorkerAssignment(ctx, store, pub, initial.Identity)
			if err != nil {
				t.Fatal(err)
			}
			state, err := latest.State(pub)
			if err != nil {
				t.Fatal(err)
			}
			if publicationFirst {
				if state.Preflight == nil || *state.Preflight != ref {
					t.Fatal("accepted reference lost")
				}
				state.Disabled = true
				state.Revision++
				if _, err := AdvanceWorkerAssignment(ctx, store, sg, latest, state); err != nil {
					t.Fatal(err)
				}
			} else if !state.Disabled || state.Preflight != nil {
				t.Fatal("publication accepted after disable")
			}
			// A different console can continue after the first console stops.
			recovered, err := ReadWorkerAssignment(ctx, store, pub, initial.Identity)
			if err != nil {
				t.Fatal(err)
			}
			state, err = recovered.State(pub)
			if err != nil {
				t.Fatal(err)
			}
			if !state.Disabled || (state.Preflight != nil) != publicationFirst {
				t.Fatal("revocation did not preserve accepted history")
			}
			retry := state
			retry.Disabled = false
			retry.Revision++
			retry.Preflight = &ref
			if _, err := AdvanceWorkerAssignment(ctx, store, sg, recovered, retry); !errors.Is(err, ErrWorkerControl) {
				t.Fatal("disabled assignment revived")
			}
		})
	}
}

func TestWorkerAssignmentConcurrentConsoleUpdatesHaveOneWinner(t *testing.T) {
	priv, pub := testKeys(t)
	sg := RootSigner(priv)
	store := NewDirStore(t.TempDir())
	ctx := t.Context()
	initial := workerTestAssignment()
	if _, err := CreateWorkerAssignment(ctx, store, sg, initial); err != nil {
		t.Fatal(err)
	}
	snapshots := make([]WorkerAssignmentSnapshot, 2)
	for i := range snapshots {
		var err error
		snapshots[i], err = ReadWorkerAssignment(ctx, store, pub, initial.Identity)
		if err != nil {
			t.Fatal(err)
		}
	}
	ready := make(chan struct{})
	results := make(chan error, 2)
	var wg sync.WaitGroup
	for i := range snapshots {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			<-ready
			after := initial
			after.Revision++
			ref := workerTestRef(byte(i + 7))
			after.Preflight = &ref
			_, err := AdvanceWorkerAssignment(ctx, store, sg, snapshots[i], after)
			results <- err
		}(i)
	}
	close(ready)
	wg.Wait()
	close(results)
	wins, conflicts := 0, 0
	for err := range results {
		if err == nil {
			wins++
		} else if errors.Is(err, ErrCASConflict) {
			conflicts++
		} else {
			t.Fatal(err)
		}
	}
	if wins != 1 || conflicts != 1 {
		t.Fatalf("wins=%d conflicts=%d", wins, conflicts)
	}
}

func TestWorkerAssignmentPreservesAcceptedExecutionOnRevocation(t *testing.T) {
	priv, pub := testKeys(t)
	sg := RootSigner(priv)
	store := &countedWorkerStore{WorkerControlStore: NewDirStore(t.TempDir())}
	ctx := t.Context()
	state := workerTestAssignment()
	snapshot, err := CreateWorkerAssignment(ctx, store, sg, state)
	if err != nil {
		t.Fatal(err)
	}
	for step := 0; step < 4; step++ {
		ref := workerTestRef(byte(step + 10))
		switch step {
		case 0:
			state.Preflight = &ref
		case 1:
			state.Compatibility = &ref
		case 2:
			state.Execution = &ref
		case 3:
			state.Resume = &ref
		}
		state.Revision++
		snapshot, err = AdvanceWorkerAssignment(ctx, store, sg, snapshot, state)
		if err != nil {
			t.Fatal(err)
		}
	}
	state.Disabled = true
	state.Revision++
	snapshot, err = AdvanceWorkerAssignment(ctx, store, sg, snapshot, state)
	if err != nil {
		t.Fatal(err)
	}
	got, err := snapshot.State(pub)
	if err != nil {
		t.Fatal(err)
	}
	if !sameWorkerReference(got.Execution, state.Execution) || !sameWorkerReference(got.Resume, state.Resume) {
		t.Fatal("accepted outcome lost during revocation")
	}
	// Accessors return independent decoded values, not mutable snapshot state.
	got.Execution.ID = workerTestRef(20).ID
	again, err := snapshot.State(pub)
	if err != nil || *again.Execution != *state.Execution {
		t.Fatal("caller mutated retained signed snapshot")
	}
	puts := store.puts.Load()
	if _, err := AdvanceWorkerAssignment(ctx, store, sg, snapshot, state); !errors.Is(err, ErrWorkerControl) {
		t.Fatal("identical retry rewrote state")
	}
	if store.puts.Load() != puts {
		t.Fatal("no-op retry reached storage")
	}
	if _, err := CreateWorkerAssignment(ctx, store, sg, workerTestAssignment()); !errors.Is(err, ErrCASConflict) {
		t.Fatal("create retry reset accepted state")
	}
}

func TestWorkerAssignmentInvalidTransitionNeverWrites(t *testing.T) {
	priv, pub := testKeys(t)
	sg := RootSigner(priv)
	ctx := t.Context()
	for name, change := range map[string]func(*WorkerAssignmentState){
		"same-revision":           func(s *WorkerAssignmentState) { s.Revision-- },
		"skip-revision":           func(s *WorkerAssignmentState) { s.Revision++ },
		"identity":                func(s *WorkerAssignmentState) { s.Identity.Box = "other-box" },
		"expiry":                  func(s *WorkerAssignmentState) { s.Expires++ },
		"replace-preflight":       func(s *WorkerAssignmentState) { r := workerTestRef(20); s.Preflight = &r },
		"remove-preflight":        func(s *WorkerAssignmentState) { s.Preflight = nil },
		"two-new-references":      func(s *WorkerAssignmentState) { r := workerTestRef(20); s.Compatibility = &r },
		"publish-and-disable":     func(s *WorkerAssignmentState) { s.Disabled = true },
		"resume-before-execution": func(s *WorkerAssignmentState) { s.Resume = s.Execution; s.Execution = nil },
	} {
		t.Run(name, func(t *testing.T) {
			store := &countedWorkerStore{WorkerControlStore: NewDirStore(t.TempDir())}
			state := workerTestAssignment()
			snapshot, err := CreateWorkerAssignment(ctx, store, sg, state)
			if err != nil {
				t.Fatal(err)
			}
			preflight := workerTestRef(10)
			state.Preflight = &preflight
			state.Revision++
			snapshot, err = AdvanceWorkerAssignment(ctx, store, sg, snapshot, state)
			if err != nil {
				t.Fatal(err)
			}
			after := state
			after.Revision++
			execution := workerTestRef(11)
			after.Execution = &execution
			change(&after)
			puts := store.puts.Load()
			if _, err := AdvanceWorkerAssignment(ctx, store, sg, snapshot, after); !errors.Is(err, ErrWorkerControl) {
				t.Fatal("invalid transition accepted")
			}
			if store.puts.Load() != puts {
				t.Fatal("invalid transition reached storage")
			}
			loaded, err := ReadWorkerAssignment(ctx, store, pub, state.Identity)
			if err != nil {
				t.Fatal(err)
			}
			got, err := loaded.State(pub)
			if err != nil || got.Revision != state.Revision || got.Execution != nil {
				t.Fatal("invalid transition changed state")
			}
		})
	}
	state := workerTestAssignment()
	state.Revision = math.MaxInt64
	after := state
	after.Revision = math.MinInt64
	after.Disabled = true
	if validWorkerAssignmentAdvance(state, after) {
		t.Fatal("revision overflow accepted")
	}
}

func TestWorkerAssignmentWrongScopeAndSignerRefuse(t *testing.T) {
	priv, pub := testKeys(t)
	sg := RootSigner(priv)
	ctx := t.Context()
	store := &countedWorkerStore{WorkerControlStore: NewDirStore(t.TempDir())}
	state := workerTestAssignment()
	snapshot, err := CreateWorkerAssignment(ctx, store, sg, state)
	if err != nil {
		t.Fatal(err)
	}
	wrong := state.Identity
	wrong.Host = "other.example.com"
	if _, err := ReadWorkerAssignment(ctx, store, pub, wrong); !errors.Is(err, ErrWorkerControl) {
		t.Fatal("wrong scope accepted")
	}
	otherPriv, _ := testKeys(t)
	_, _, scoped := scopedFixture(t, []string{"canary"}, time.Hour)
	after := state
	after.Revision++
	after.Disabled = true
	puts := store.puts.Load()
	for _, signer := range []Signer{{}, RootSigner(otherPriv), scoped} {
		if _, err := AdvanceWorkerAssignment(ctx, store, signer, snapshot, after); !errors.Is(err, ErrWorkerControl) {
			t.Fatal("wrong signer accepted")
		}
	}
	if _, err := AdvanceWorkerAssignment(ctx, store, sg, WorkerAssignmentSnapshot{}, after); !errors.Is(err, ErrWorkerControl) {
		t.Fatal("constructed empty snapshot accepted")
	}
	if store.puts.Load() != puts {
		t.Fatal("unauthorized transition wrote storage")
	}
}

func TestWorkerAssignmentCancelledOperationDoesNotWrite(t *testing.T) {
	priv, pub := testKeys(t)
	sg := RootSigner(priv)
	store := &countedWorkerStore{WorkerControlStore: NewDirStore(t.TempDir())}
	state := workerTestAssignment()
	snapshot, err := CreateWorkerAssignment(t.Context(), store, sg, state)
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(t.Context())
	cancel()
	after := state
	after.Revision++
	after.Disabled = true
	puts := store.puts.Load()
	if _, err := AdvanceWorkerAssignment(ctx, store, sg, snapshot, after); !errors.Is(err, ErrWorkerControl) {
		t.Fatal("cancelled write accepted")
	}
	if _, err := CreateWorkerAssignment(ctx, store, sg, state); !errors.Is(err, ErrWorkerControl) {
		t.Fatal("cancelled create accepted")
	}
	if _, err := ReadWorkerAssignment(ctx, store, pub, state.Identity); !errors.Is(err, ErrWorkerControl) {
		t.Fatal("cancelled read accepted")
	}
	if store.puts.Load() != puts {
		t.Fatal("cancelled operation reached storage")
	}
}