package lease

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

// The parallel round (offline-move speedup): copies and deletes run up to
// mirrorWorkers(m) at once. These tests hold the round to what the serial
// one promised: the same end state, receipts for exactly what landed, the
// first failure stopping the round, and no skip a receipt cannot prove.
// Run them under -race: the round's own state is shared by its workers.

// fillReplica stores n objects under instances/h/: every third one
// multipart (the copy lands with a new ETag, so it needs a receipt), the
// rest single-part. Sizes differ so byte totals catch a double count.
func fillReplica(f *fakeMirror, n int, gen string) (keys []string, bytes int64) {
	for i := 0; i < n; i++ {
		key := fmt.Sprintf("instances/h/ltx/0/%s-%04d.ltx", gen, i)
		size := int64(1000 + i)
		if i%3 == 0 {
			f.putMultipart("src", key, fmt.Sprintf("%s-body-%d", gen, i), size)
		} else {
			f.putSized("src", key, fmt.Sprintf("%s-etag-%d", gen, i), size)
		}
		keys = append(keys, key)
		bytes += size
	}
	return keys, bytes
}

// stagedBodies is moves/h/ by key relative to it, body for body.
func stagedBodies(f *fakeMirror) map[string]string {
	f.mu.Lock()
	defer f.mu.Unlock()
	out := map[string]string{}
	for k, v := range f.buckets["dst"] {
		if strings.HasPrefix(k, "moves/h/") {
			out[strings.TrimPrefix(k, "moves/h/")] = v.body
		}
	}
	return out
}

// sourceBodies is instances/h/ the same way, lock.json aside.
func sourceBodies(f *fakeMirror) map[string]string {
	f.mu.Lock()
	defer f.mu.Unlock()
	out := map[string]string{}
	for k, v := range f.buckets["src"] {
		if rel, ok := strings.CutPrefix(k, "instances/h/"); ok && rel != "lock.json" {
			out[rel] = v.body
		}
	}
	return out
}

func sameBodies(t *testing.T, got, want map[string]string) {
	t.Helper()
	if len(got) != len(want) {
		t.Fatalf("staged %d objects, source has %d", len(got), len(want))
	}
	for k, v := range want {
		if got[k] != v {
			t.Fatalf("staged %s = %q, source %q", k, got[k], v)
		}
	}
}

func copiesIn(ops []string) int {
	n := 0
	for _, op := range ops {
		if strings.HasPrefix(op, "copy ") {
			n++
		}
	}
	return n
}

// A parallel round ends where a serial one does: every object staged body
// for body, the same counts and bytes, receipts for exactly the multipart
// objects, and a caught-up round after it. It really runs side by side,
// and never more than the store allows.
func TestMirrorPrefixParallelMatchesSerial(t *testing.T) {
	const n = 120
	var rounds [2]MirrorRound
	var books [2]map[string]copyReceipt
	for i, workers := range []int{1, 6} {
		f := newFakeMirror()
		f.workers = workers
		_, bytes := fillReplica(f, n, "a")
		f.put("src", "instances/h/lock.json", "renewal")
		// Slow copies, so the workers overlap.
		f.onCopy = func(context.Context, string) error { time.Sleep(time.Millisecond); return nil }
		r := mirrorRound(t, f, testReceipts)
		if r.Changed != n || r.BytesCopied != bytes || r.Deferred != 0 {
			t.Fatalf("workers=%d: round = %+v, want %d copies of %d bytes", workers, r, n, bytes)
		}
		sameBodies(t, stagedBodies(f), sourceBodies(f))
		if f.peak > workers || (workers > 1 && f.peak < 2) {
			t.Fatalf("workers=%d: peak concurrency %d", workers, f.peak)
		}
		book := readReceipts(t, f)
		if len(book) != n/3 {
			t.Fatalf("workers=%d: %d receipts, want one per multipart object (%d)", workers, len(book), n/3)
		}
		if again := mirrorRound(t, f, testReceipts); again.Changed != 0 {
			t.Fatalf("workers=%d: caught-up round = %+v", workers, again)
		}
		rounds[i], books[i] = r, book
	}
	if rounds[0].Changed != rounds[1].Changed || rounds[0].BytesCopied != rounds[1].BytesCopied {
		t.Fatalf("serial %+v and parallel %+v rounds differ", rounds[0], rounds[1])
	}
	for k, v := range books[0] {
		if books[1][k] != v {
			t.Fatalf("receipt %s: serial %+v, parallel %+v", k, v, books[1][k])
		}
	}
}

// The versioned mode under workers: every copied object's version lands in
// Versions, and the receipts carry them, so the next round names them all
// without copying.
func TestMirrorPrefixParallelVersionedRound(t *testing.T) {
	f := &fakeVersionedMirror{fakeMirror: newFakeMirror()}
	f.workers = 8
	keys, _ := fillReplica(f.fakeMirror, 64, "a")
	r1, err := MirrorPrefix(context.Background(), f, "src", "instances/h/", "dst", "moves/h/", testReceipts)
	if err != nil || len(r1.Versions) != len(keys) {
		t.Fatalf("round 1 = %d versions, %v; want %d", len(r1.Versions), err, len(keys))
	}
	seen := map[string]bool{}
	for _, v := range r1.Versions {
		if seen[v] {
			t.Fatalf("version %s handed to two copies", v)
		}
		seen[v] = true
	}
	r2, err := MirrorPrefix(context.Background(), f, "src", "instances/h/", "dst", "moves/h/", testReceipts)
	if err != nil || r2.Changed != 0 || len(r2.Versions) != len(keys) {
		t.Fatalf("round 2 = %+v, %v; want nothing copied and every version named", r2, err)
	}
	for k, v := range r1.Versions {
		if r2.Versions[k] != v {
			t.Fatalf("%s: round 2 names %q, round 1 copied %q", k, r2.Versions[k], v)
		}
	}
}

// The first failure stops the round: it is the error, copies still in
// flight see their context cancelled, nothing new starts, and every copy
// that landed keeps its receipt, so the next round copies only what did
// not land.
func TestMirrorPrefixParallelFailureMidRound(t *testing.T) {
	f := newFakeMirror()
	f.workers = 4
	// All multipart: a landed copy is skipped next round only by its receipt.
	for i := 0; i < 40; i++ {
		f.putMultipart("src", fmt.Sprintf("instances/h/ltx/1/%04d.ltx", i), fmt.Sprintf("b%d", i), 100)
	}
	bad := "instances/h/ltx/1/0010.ltx"
	var cancelled atomic.Int32
	later := make(chan struct{}, 64) // a copy after the bad one has started
	f.onCopy = func(ctx context.Context, key string) error {
		switch {
		case key == bad:
			// Fail only once a later copy is in flight, so the cancel
			// has something to reach.
			select {
			case <-later:
			case <-time.After(5 * time.Second):
			}
			return errors.New("injected copy failure")
		case key > bad:
			later <- struct{}{}
			select {
			case <-ctx.Done():
				cancelled.Add(1)
				return ctx.Err()
			case <-time.After(5 * time.Second):
				return errors.New("the round did not cancel a copy in flight")
			}
		}
		return nil
	}
	r, err := MirrorPrefix(context.Background(), f, "src", "instances/h/", "dst", "moves/h/", testReceipts)
	if err == nil || !strings.Contains(err.Error(), bad) || errors.Is(err, context.Canceled) {
		t.Fatalf("round error = %v, want the injected failure on %s", err, bad)
	}
	landed := stagedBodies(f)
	if r.Changed != len(landed) || r.BytesCopied != int64(len(landed))*100 {
		t.Fatalf("round = %+v, but %d objects landed", r, len(landed))
	}
	if len(landed) == 0 || len(landed) > 10 {
		t.Fatalf("%d objects landed, want those before the failure only", len(landed))
	}
	if cancelled.Load() == 0 {
		t.Fatal("no copy in flight was cancelled")
	}
	// Nothing past the workers in flight at the failure ever started.
	if started := copiesIn(f.ops) + int(cancelled.Load()) + 1; started > 10+1+f.workers {
		t.Fatalf("%d copies started, want the round to stop dispatching at the failure", started)
	}
	book := readReceipts(t, f)
	for rel := range landed {
		if _, ok := book[rel]; !ok {
			t.Fatalf("%s landed without a receipt", rel)
		}
	}
	if len(book) != len(landed) {
		t.Fatalf("%d receipts for %d landed copies", len(book), len(landed))
	}
	f.onCopy = nil
	before := len(f.ops)
	next := mirrorRound(t, f, testReceipts)
	if next.Changed != 40-len(landed) {
		t.Fatalf("next round = %+v, want exactly the %d that did not land", next, 40-len(landed))
	}
	for _, op := range f.ops[before:] {
		for rel := range landed {
			if strings.HasPrefix(op, "copy instances/h/"+rel+" ") {
				t.Fatalf("a landed copy was recopied: %s", op)
			}
		}
	}
	sameBodies(t, stagedBodies(f), sourceBodies(f))
}

// A caller's context that ends mid-round is the round's error, never a
// silent partial success.
func TestMirrorPrefixParallelContextEnds(t *testing.T) {
	f := newFakeMirror()
	f.workers = 2
	fillReplica(f, 30, "a")
	ctx, cancel := context.WithCancel(context.Background())
	var n atomic.Int32
	f.onCopy = func(context.Context, string) error {
		if n.Add(1) == 3 {
			cancel()
		}
		return nil
	}
	if _, err := MirrorPrefix(ctx, f, "src", "instances/h/", "dst", "moves/h/", testReceipts); !errors.Is(err, context.Canceled) {
		t.Fatalf("round error = %v, want the context's", err)
	}
	if len(stagedBodies(f)) == 30 {
		t.Fatal("every copy ran after the context ended")
	}
}

// The pre-copy (offline move stage): rounds while the instance serves, then
// the seal's round. The seal finds what stage copied and copies only what is
// new since; its receipts prove every multipart skip.
func TestMirrorPrefixReceiptsSkipAfterAPreCopy(t *testing.T) {
	f := newFakeMirror()
	_, staged := fillReplica(f, 90, "a")
	if r := mirrorRound(t, f, testReceipts); r.BytesCopied != staged {
		t.Fatalf("stage round = %+v, want %d bytes", r, staged)
	}
	// The instance kept writing: new LTX files.
	_, fresh := fillReplica(f, 6, "b")
	before := len(f.ops)
	seal := mirrorRound(t, f, testReceipts)
	if seal.BytesCopied != fresh || seal.Changed != 6 || seal.Deferred != 0 {
		t.Fatalf("seal round = %+v, want only the 6 new objects (%d bytes)", seal, fresh)
	}
	for _, op := range f.ops[before:] {
		if strings.HasPrefix(op, "copy ") && !strings.Contains(op, "/b-") {
			t.Fatalf("the seal recopied a pre-copied object: %s", op)
		}
	}
	sameBodies(t, stagedBodies(f), sourceBodies(f))
}

// Between stage and seal, litestream's retention and compaction delete
// replica files and write new ones. The seal's round deletes what the source
// no longer has and drops those receipts, so the staged listing is the
// source's again (the seal's sameListing).
func TestMirrorPrefixSourceDeletionsBetweenStageAndSeal(t *testing.T) {
	f := newFakeMirror()
	keys, _ := fillReplica(f, 30, "a")
	mirrorRound(t, f, testReceipts)
	// Compaction: twelve level-0 files folded into one level-1 file.
	for _, k := range keys[:12] {
		f.remove("src", k)
	}
	f.putMultipart("src", "instances/h/ltx/1/compacted.ltx", "compacted", 12000)
	r := mirrorRound(t, f, testReceipts)
	if r.Changed != 13 || r.BytesCopied != 12000 {
		t.Fatalf("seal round = %+v, want 12 deletes and the one copy", r)
	}
	sameBodies(t, stagedBodies(f), sourceBodies(f))
	for rel := range readReceipts(t, f) {
		if _, ok := sourceBodies(f)[rel]; !ok {
			t.Fatalf("a receipt outlived its source object: %s", rel)
		}
	}
	if again := mirrorRound(t, f, testReceipts); again.Changed != 0 {
		t.Fatalf("round after = %+v", again)
	}
}

// A staged prefix left by an earlier move is never trusted as such. A stale
// object skips only when its ETag equals the source's (the same bytes) or a
// receipt for this exact pair still matches both sides; anything else is
// copied again, and every key the source lacks is deleted.
func TestMirrorPrefixStaleStagedPrefixIsNotTrusted(t *testing.T) {
	f := newFakeMirror()
	f.putMultipart("src", "instances/h/ltx/1/a.ltx", "new-a", 500)
	f.putMultipart("src", "instances/h/ltx/1/b.ltx", "new-b", 500)
	f.putSized("src", "instances/h/ltx/0/c.ltx", "new-c", 50)
	f.putSized("src", "instances/h/ltx/0/d.ltx", "same-d", 50)
	// The stale prefix: a under the old move's copy of other bytes, b with
	// the same size as the source and a stale receipt, c changed, d byte
	// identical, and an extra the source never had.
	f.putSized("dst", "moves/h/ltx/1/a.ltx", "old-a", 500)
	f.putSized("dst", "moves/h/ltx/1/b.ltx", "old-b", 500)
	f.putSized("dst", "moves/h/ltx/0/c.ltx", "old-c", 50)
	f.putSized("dst", "moves/h/ltx/0/d.ltx", "same-d", 50)
	f.putSized("dst", "moves/h/ltx/0/extra.ltx", "old-x", 50)
	// An old book for this pair whose receipt for b names the source as it
	// is now but a destination ETag the staged object does not carry.
	stale := (`{"v":1,"src_bucket":"src","src_prefix":"instances/h/","dst_bucket":"dst","dst_prefix":"moves/h/",` +
		`"objects":{"ltx/1/b.ltx":{"src_etag":"new-b-27","src_size":500,"dst_etag":"new-b"},"ltx/0/extra.ltx":{"src_etag":"x","src_size":50,"dst_etag":"old-x"}}}`)
	f.store("src", testReceipts.Key, fakeObject{MirrorObject{ETag: "book", Size: int64(len(stale))}, stale})

	r := mirrorRound(t, f, testReceipts)
	if r.Changed != 4 || r.BytesCopied != 1050 {
		t.Fatalf("round = %+v, want a, b and c copied and the extra deleted", r)
	}
	sameBodies(t, stagedBodies(f), sourceBodies(f))
	book := readReceipts(t, f)
	if _, ok := book["ltx/0/extra.ltx"]; ok {
		t.Fatal("the stale receipt of a deleted key survived")
	}
	if got := book["ltx/1/b.ltx"]; got.DstETag != "new-b" || got.SrcETag != "new-b-27" {
		t.Fatalf("b's receipt = %+v", got)
	}

	// A book written for another pair (another move's target) is no
	// evidence at all, even when its receipt would match.
	g := newFakeMirror()
	g.putMultipart("src", "instances/h/ltx/1/a.ltx", "a", 500)
	g.putSized("dst", "moves/h/ltx/1/a.ltx", "a", 500)
	other := `{"v":1,"src_bucket":"src","src_prefix":"instances/h/","dst_bucket":"other-dst","dst_prefix":"moves/h/",` +
		`"objects":{"ltx/1/a.ltx":{"src_etag":"a-27","src_size":500,"dst_etag":"a"}}}`
	g.store("src", testReceipts.Key, fakeObject{MirrorObject{ETag: "book", Size: int64(len(other))}, other})
	if r := mirrorRound(t, g, testReceipts); r.Changed != 1 {
		t.Fatalf("round over another pair's book = %+v, want the object copied again", r)
	}
}

// forEachParallel on its own: an error stops dispatch, and an already-ended
// context starts nothing.
func TestForEachParallel(t *testing.T) {
	var ran atomic.Int32
	err := forEachParallel(context.Background(), 1, 10, func(_ context.Context, i int) error {
		ran.Add(1)
		if i == 3 {
			return errors.New("boom")
		}
		return nil
	})
	if err == nil || err.Error() != "boom" || ran.Load() > 5 {
		t.Fatalf("err %v after %d calls, want boom after at most 5", err, ran.Load())
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	ran.Store(0)
	err = forEachParallel(ctx, 4, 10, func(context.Context, int) error { ran.Add(1); return nil })
	if !errors.Is(err, context.Canceled) || ran.Load() > 4 {
		t.Fatalf("err %v after %d calls on an ended context", err, ran.Load())
	}
	if err := forEachParallel(context.Background(), 8, 0, nil); err != nil {
		t.Fatal(err)
	}
}

// Astra round 1, finding 3: no network request is made while the round's
// accounting is locked. A receipts read that stalls once a copy has landed
// (the lazy load a receipt write used to make under the lock) must not keep
// another copy's failure from stopping the round: the book is read before
// any copy starts.
func TestMirrorPrefixAReceiptsReadNeverBlocksTheFirstFailure(t *testing.T) {
	f := newFakeMirror()
	f.workers = 2
	f.putMultipart("src", "instances/h/ltx/1/a.ltx", "a", 100) // lands, needs a receipt
	f.putSized("src", "instances/h/ltx/0/b.ltx", "b", 10)      // fails
	var landed atomic.Bool
	release := make(chan struct{})
	t.Cleanup(func() { close(release) })
	f.onCopy = func(_ context.Context, key string) error {
		if strings.HasSuffix(key, "b.ltx") {
			// Fail only once a's copy has landed and its receipt is due.
			for deadline := time.Now().Add(5 * time.Second); !landed.Load() && time.Now().Before(deadline); time.Sleep(time.Millisecond) {
			}
			time.Sleep(20 * time.Millisecond)
			return errors.New("injected copy failure")
		}
		landed.Store(true)
		return nil
	}
	f.onGet = func(string) {
		if landed.Load() {
			<-release // a receipts read that stalls after a copy
		}
	}
	done := make(chan error, 1)
	go func() {
		_, err := MirrorPrefix(context.Background(), f, "src", "instances/h/", "dst", "moves/h/", testReceipts)
		done <- err
	}()
	select {
	case err := <-done:
		if err == nil || !strings.Contains(err.Error(), "injected copy failure") {
			t.Fatalf("round error = %v, want the injected failure", err)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("a receipts read made under the round's lock kept the first failure from stopping the round")
	}
}