package dbreads

import (
	"context"
	"errors"
	"net"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/benbjohnson/litestream"

	"github.com/carlosframework/platform/internal/release"
)

// countingStore counts reads of request objects and lists.
type countingStore struct {
	release.Store
	requestGets, lists atomic.Int64
}

func (s *countingStore) Get(ctx context.Context, key string) ([]byte, string, error) {
	if strings.HasSuffix(key, "/request.json") {
		s.requestGets.Add(1)
	}
	return s.Store.Get(ctx, key)
}

func (s *countingStore) List(ctx context.Context, prefix string) ([]string, error) {
	s.lists.Add(1)
	return s.Store.List(ctx, prefix)
}

// Review Focus 2: a request issued at 10:59:50 is listed at 11:00:20 under
// the previous hour's prefix, and picked.
func TestPassFindsAnEntryFromThePreviousHour(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 1)
	fx.now = time.Date(2026, 10, 22, 10, 59, 50, 0, time.UTC)
	a, _ := fx.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) {
		c.Now = func() time.Time { return time.Date(2026, 10, 22, 11, 0, 20, 0, time.UTC) }
	})
	req := fx.request("SELECT count(*) FROM t", "")
	fx.submit(req)
	a.Pass(context.Background())
	a.Wait()
	if _, _, r := fx.answer(req.ID); r.State != StateDone {
		t.Fatalf("%+v", r)
	}
}

func TestPassSkipsAnOutOfWindowEntryWithoutReadingIt(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 0)
	cs := &countingStore{Store: fx.store}
	fx.now = time.Now().UTC().Add(-3 * time.Minute).Truncate(time.Second)
	fx.submit(fx.request("SELECT 1", ""))
	a, _ := fx.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) { c.Store = cs })
	a.Pass(context.Background())
	if cs.requestGets.Load() != 0 {
		t.Fatal("an out-of-window request was read")
	}
}

// T4, the box's half: a forged queue entry and an unsigned request are
// refused, decided, and logged once.
func TestPassRefusesForgedEntriesAndUnsignedRequests(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 1)
	var logged atomic.Int64
	a, w := fx.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) {
		c.Logf = func(format string, args ...any) {
			if strings.Contains(format, "refused the request") {
				logged.Add(1)
			}
		}
	})
	forged := fx.request("SELECT 1", "")
	wire, _ := Sign(release.RootSigner(fx.rootPriv), forged)
	fx.store.Put(context.Background(), RequestKey(forged.ID), wire, "")
	issued, _ := forged.Issued()
	fx.store.Put(context.Background(), QueueKey("uk", issued.Add(-time.Second), forged.ID), nil, "")
	unsigned := fx.request("SELECT 1", "")
	fx.store.Put(context.Background(), RequestKey(unsigned.ID), []byte(`{"payload":"e30=","sig":""}`), "")
	fx.store.Put(context.Background(), QueueKey("uk", issued, unsigned.ID), nil, "")
	a.Pass(context.Background())
	a.Pass(context.Background())
	a.Wait()
	for _, id := range []string{forged.ID, unsigned.ID} {
		if _, _, err := fx.store.Get(context.Background(), PickedKey(id)); err == nil {
			t.Errorf("%s was picked", id)
		}
	}
	if logged.Load() != 2 || w.dials.Load() != 0 {
		t.Fatalf("logged %d refusals, dialed %d workers", logged.Load(), w.dials.Load())
	}
}

// T4 and codex 5: the same envelope offered again, before and after the
// agent restarts, is never run twice: picked holds.
func TestReplayIsRefusedByPickedAcrossARestart(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 1)
	a, _ := fx.agent("uk-1")
	req := fx.request("SELECT count(*) FROM t", "")
	fx.submit(req)
	a.Pass(context.Background())
	a.Wait()
	a.Pass(context.Background()) // the same agent again
	restarted, w := fx.agent("uk-1")
	restarted.Pass(context.Background())
	restarted.Wait()
	if w.dials.Load() != 0 {
		t.Fatal("a restarted agent ran a read it had already picked")
	}
}

func TestTwoAgentsServeEachReadOnce(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 1)
	a, wa := fx.agent("uk-1")
	b, wb := fx.agent("uk-2")
	var ids []string
	for i := 0; i < 10; i++ {
		req := fx.request("SELECT count(*) FROM t", "")
		fx.submit(req)
		ids = append(ids, req.ID)
	}
	for round := 0; round < 20; round++ {
		done := make(chan struct{}, 2)
		go func() { a.Pass(context.Background()); done <- struct{}{} }()
		go func() { b.Pass(context.Background()); done <- struct{}{} }()
		<-done
		<-done
		a.Wait()
		b.Wait()
	}
	for _, id := range ids {
		if _, _, r := fx.answer(id); r.State != StateDone {
			t.Fatalf("%s: %+v", id, r)
		}
	}
	if got := wa.dials.Load() + wb.dials.Load(); got != int64(len(ids)) {
		t.Fatalf("%d worker connections for %d reads", got, len(ids))
	}
}

// Plan review finding 14 (spec §7.1, D2): a draining box beside a healthy
// one never wins a read, including one that starts draining during its
// jitter. Mutation: drop the Draining check after the jitter; the draining
// agent picks every read and this test is red. (The check before admission
// only spares a reservation; the one after the jitter is the gate.)
func TestADrainingBoxDoesNotPick(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 1)
	var flip atomic.Bool
	drain, dw := fx.agent("uk-drain", func(c *AgentConfig, _ *inProcWorkers) {
		c.Draining = flip.Load
		c.Jitter = func() time.Duration { flip.Store(true); return time.Millisecond }
	})
	healthy, _ := fx.agent("uk-1")
	for i := 0; i < 5; i++ {
		req := fx.request("SELECT 1", "")
		fx.submit(req)
		drain.Pass(context.Background()) // first: it passes the first check, then drains in its jitter
		drain.Wait()
		healthy.Pass(context.Background())
		healthy.Wait()
		if _, _, r := fx.answer(req.ID); r.Box != "uk-1" {
			t.Fatalf("read %d served by %s", i, r.Box)
		}
		flip.Store(false)
	}
	if dw.dials.Load() != 0 {
		t.Fatal("the draining box started a worker")
	}
}

// Finding 11's box half: a request already cancelled is never picked.
// Mutation: drop the cancel read in consider; the agent picks it.
func TestPassSkipsACancelledRequest(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 1)
	a, w := fx.agent("uk-1")
	req := fx.request("SELECT 1", "")
	fx.submit(req)
	fx.store.Put(context.Background(), CancelKey(req.ID), nil, "")
	a.Pass(context.Background())
	a.Wait()
	if _, _, err := fx.store.Get(context.Background(), PickedKey(req.ID)); err == nil || w.dials.Load() != 0 {
		t.Fatal("a cancelled request was picked")
	}
}

func TestPassNeedsTheCapability(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 0)
	cs := &countingStore{Store: fx.store}
	a, _ := fx.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) { c.Store = cs })
	a.probeOK.Store(false)
	a.Pass(context.Background())
	if cs.lists.Load() != 0 {
		t.Fatal("an incapable box listed the queue")
	}
}

func TestProbe(t *testing.T) {
	t.Parallel()
	for name, c := range map[string]struct {
		self func() HelloReply
		dial error
		want bool
	}{
		"a sandboxed worker":  {func() HelloReply { return HelloReply{V: ProtocolVersion, Interfaces: []string{"lo"}} }, nil, true},
		"another version":     {func() HelloReply { return HelloReply{V: 2, Interfaces: []string{"lo"}} }, nil, false},
		"a network":           {func() HelloReply { return HelloReply{V: 1, Interfaces: []string{"lo", "eth0"}} }, nil, false},
		"a visible /data":     {func() HelloReply { return HelloReply{V: 1, Interfaces: []string{"lo"}, DataEntries: 3} }, nil, false},
		"a readable host.env": {func() HelloReply { return HelloReply{V: 1, Interfaces: []string{"lo"}, EtcReadable: true} }, nil, false},
		"an observation that failed": {func() HelloReply {
			return HelloReply{V: 1, Errors: []string{"interfaces: permission denied"}}
		}, nil, false},
		"no service": {nil, errors.New("connection refused"), false},
	} {
		t.Run(name, func(t *testing.T) {
			t.Parallel()
			fx := newFixture(t, 0)
			a, _ := fx.agent("uk-1", func(cfg *AgentConfig, w *inProcWorkers) { w.opt.Self = c.self })
			if c.dial != nil {
				a.cfg.DialWorker = func(context.Context) (net.Conn, error) { return nil, c.dial }
			}
			a.probeOK.Store(!c.want)
			a.Probe(context.Background())
			if a.Capable() != c.want {
				t.Fatalf("Capable = %v", a.Capable())
			}
		})
	}
}

// Review Focus 5 and spec §7.1: a probe beside an orphan, or with no unit
// listing, is skipped and the last result stands.
func TestProbeIsSkippedWhileAdmissionIsBlocked(t *testing.T) {
	t.Parallel()
	for _, last := range []bool{true, false} {
		fx := newFixture(t, 0)
		a, _ := fx.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) {
			c.Units = func(context.Context) (int, error) { return 1, nil } // an orphan
		})
		a.probeOK.Store(last)
		a.Probe(context.Background())
		if a.Capable() != last {
			t.Fatalf("a skipped probe changed the capability from %v", last)
		}
	}
}

// T13, the unit half: a box running two reads keeps the capability; its
// probe gets the third connection.
func TestProbeOnTheThirdConnection(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 1)
	hang := make(chan struct{})
	a, w := fx.agent("uk-1", func(c *AgentConfig, w *inProcWorkers) {
		w.opt.Restore = func(ctx context.Context, _ litestream.ReplicaClient, _ string, _ uint64) error {
			<-hang
			return errors.New("released")
		}
	})
	for i := 0; i < 2; i++ {
		fx.submit(fx.request("SELECT 1", ""))
	}
	a.Pass(context.Background())
	deadline := time.Now().Add(5 * time.Second)
	for w.live.Load() < 2 && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	a.Probe(context.Background())
	if !a.Capable() {
		t.Fatal("a busy box lost the capability")
	}
	close(hang)
	a.Wait()
}

func TestRunPassesOnANudge(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 1)
	// The start probe's worker unit lingers as systemd's deactivation (and a
	// loaded CI box) makes it: main's CI caught the nudge's pass meeting it
	// as an orphan, refusing admission, and nothing passing again within
	// the test's bound (ListEvery is 30 s).
	a, w := fx.agent("uk-1", func(_ *AgentConfig, w *inProcWorkers) { w.linger = 300 * time.Millisecond })
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	nudge := make(chan struct{}, 1)
	go a.Run(ctx, nudge)
	// A pass that meets that unit refuses admission by design (spec §5.7)
	// and the read waits for a later pass. This test is about the nudge, so
	// it nudges once the start probe has run and its unit is gone.
	ready := time.Now().Add(10 * time.Second)
	for !(w.dials.Load() >= 1 && w.live.Load() == 0 && a.Capable()) {
		if time.Now().After(ready) {
			t.Fatal("the start probe never finished")
		}
		time.Sleep(10 * time.Millisecond)
	}
	req := fx.request("SELECT count(*) FROM t", "")
	fx.submit(req)
	nudge <- struct{}{}
	deadline := time.Now().Add(10 * time.Second)
	for time.Now().Before(deadline) {
		if _, _, err := fx.store.Get(context.Background(), ResultKey(req.ID)); err == nil {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("a nudge did not serve the read")
}

// Plan review finding 29, through the real loop: a read whose worker never
// starts asks for a probe while that worker's unit is still stopping, so
// admission skips it; the probe stays owed and runs once the unit is gone,
// and the capability is withdrawn well inside ProbeEvery. Mutation: Run
// calls Probe once on probeAfter and drops it when skipped (no retry):
// Capable stays true until ProbeEvery and this test is red.
func TestAProbeSkippedWhileTheUnitLingersRunsOnceItIsGone(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 1)
	a, w := fx.agent("uk-1", func(_ *AgentConfig, w *inProcWorkers) { w.linger = 700 * time.Millisecond })
	ctx, cancel := context.WithCancel(context.Background())
	nudge := make(chan struct{}, 1)
	done := make(chan struct{})
	go func() { a.Run(ctx, nudge); close(done) }()
	defer func() { cancel(); <-done }()
	until := func(what string, ok func() bool) {
		t.Helper()
		deadline := time.Now().Add(10 * time.Second)
		for !ok() {
			if time.Now().After(deadline) {
				t.Fatalf("never: %s", what)
			}
			time.Sleep(20 * time.Millisecond)
		}
	}
	// The start probe has run and its unit is gone.
	until("the start probe", func() bool { return w.dials.Load() >= 1 && w.live.Load() == 0 && a.Capable() })
	w.failStart.Store(true)
	req := fx.request("SELECT 1", "")
	fx.submit(req)
	nudge <- struct{}{}
	until("the read's result", func() bool {
		_, _, err := fx.store.Get(context.Background(), ResultKey(req.ID))
		return err == nil
	})
	until("the owed probe withdrawing the capability", func() bool { return !a.Capable() })
}