package depot

import (
	"context"
	"io"
	"sort"
	"sync/atomic"
	"testing"
	"time"
)

// stubBackend lets tests script Probe behavior per call.
type stubBackend struct {
	probe func(ctx context.Context, sha string) (ProbeState, bool, error)
}

func (s *stubBackend) Name() string { return "stub" }
func (s *stubBackend) Probe(ctx context.Context, sha string) (ProbeState, bool, error) {
	return s.probe(ctx, sha)
}
func (s *stubBackend) Get(ctx context.Context, sha, dest string) error { return nil }
func (s *stubBackend) Put(ctx context.Context, sha, src string) error  { return nil }

func sweepTestCfg() Config {
	return Config{ProbeJobs: 4, ConfirmRetries: 3, ConfirmMax: 200}
}

func TestSweepAllPresent(t *testing.T) {
	b := &stubBackend{probe: func(ctx context.Context, sha string) (ProbeState, bool, error) {
		return Present, false, nil
	}}
	shas := []string{"bbb", "aaa", "ccc"}
	res, err := Sweep(context.Background(), b, shas, sweepTestCfg(), io.Discard)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	want := []string{"aaa", "bbb", "ccc"}
	sort.Strings(want)
	if len(res.Present) != 3 || len(res.Missing) != 0 || len(res.Unconfirmed) != 0 {
		t.Fatalf("got %+v", res)
	}
	for i, sha := range res.Present {
		if sha != want[i] {
			t.Fatalf("expected sorted output, got %v", res.Present)
		}
	}
}

func TestSweepMissing(t *testing.T) {
	b := &stubBackend{probe: func(ctx context.Context, sha string) (ProbeState, bool, error) {
		return Absent, false, nil
	}}
	shas := []string{"aaa"}
	res, err := Sweep(context.Background(), b, shas, sweepTestCfg(), io.Discard)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(res.Missing) != 1 || res.Missing[0] != "aaa" {
		t.Fatalf("got %+v", res)
	}
	if len(res.Present) != 0 || len(res.Unconfirmed) != 0 {
		t.Fatalf("got %+v", res)
	}
}

func TestSweepThrottledUnconfirmed(t *testing.T) {
	var sleeps []time.Duration
	orig := confirmSleep
	confirmSleep = func(d time.Duration) { sleeps = append(sleeps, d) }
	defer func() { confirmSleep = orig }()

	b := &stubBackend{probe: func(ctx context.Context, sha string) (ProbeState, bool, error) {
		return Unconfirmed, true, nil
	}}
	shas := []string{"aaa"}
	cfg := sweepTestCfg()
	res, err := Sweep(context.Background(), b, shas, cfg, io.Discard)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(res.Unconfirmed) != 1 || res.Unconfirmed[0] != "aaa" {
		t.Fatalf("got %+v", res)
	}
	if len(res.Missing) != 0 {
		t.Fatalf("must never collapse into missing: %+v", res)
	}
	if len(sleeps) != 2 || sleeps[0] != 1*time.Second || sleeps[1] != 2*time.Second {
		t.Fatalf("expected sleeps [1s 2s], got %v", sleeps)
	}
}

func TestSweepTransientRecovers(t *testing.T) {
	orig := confirmSleep
	confirmSleep = func(d time.Duration) {}
	defer func() { confirmSleep = orig }()

	var calls int32
	b := &stubBackend{probe: func(ctx context.Context, sha string) (ProbeState, bool, error) {
		n := atomic.AddInt32(&calls, 1)
		if n == 1 {
			return Unconfirmed, true, nil // initial parallel probe: 428
		}
		return Present, false, nil // serial re-probe: 206
	}}
	res, err := Sweep(context.Background(), b, []string{"aaa"}, sweepTestCfg(), io.Discard)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(res.Present) != 1 || res.Present[0] != "aaa" {
		t.Fatalf("got %+v", res)
	}
}

func TestSweepConfirmMaxCap(t *testing.T) {
	orig := confirmSleep
	var sleepCalls int32
	confirmSleep = func(d time.Duration) { atomic.AddInt32(&sleepCalls, 1) }
	defer func() { confirmSleep = orig }()

	b := &stubBackend{probe: func(ctx context.Context, sha string) (ProbeState, bool, error) {
		return Unconfirmed, true, nil
	}}
	shas := []string{"a", "b", "c", "d", "e"}
	cfg := sweepTestCfg()
	cfg.ConfirmMax = 2
	res, err := Sweep(context.Background(), b, shas, cfg, io.Discard)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(res.Unconfirmed) != 5 {
		t.Fatalf("expected all 5 unconfirmed, got %+v", res)
	}
	if len(res.Missing) != 0 || len(res.Present) != 0 {
		t.Fatalf("got %+v", res)
	}
	if sleepCalls != 0 {
		t.Fatalf("expected zero serial re-probes/sleeps when over ConfirmMax, got %d calls", sleepCalls)
	}
}

func TestSweepParallelBounded(t *testing.T) {
	var inFlight int32
	var maxInFlight int32
	shas := make([]string, 100)
	for i := range shas {
		shas[i] = string(rune('a'+i%26)) + string(rune('A'+i/26))
	}
	b := &stubBackend{probe: func(ctx context.Context, sha string) (ProbeState, bool, error) {
		n := atomic.AddInt32(&inFlight, 1)
		for {
			cur := atomic.LoadInt32(&maxInFlight)
			if n <= cur || atomic.CompareAndSwapInt32(&maxInFlight, cur, n) {
				break
			}
		}
		time.Sleep(time.Millisecond)
		atomic.AddInt32(&inFlight, -1)
		return Present, false, nil
	}}
	cfg := sweepTestCfg()
	cfg.ProbeJobs = 4
	_, err := Sweep(context.Background(), b, shas, cfg, io.Discard)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if maxInFlight > 4 {
		t.Fatalf("expected max concurrency <= 4, got %d", maxInFlight)
	}
}