package dbreads

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io"
	"os"
	"reflect"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/benbjohnson/litestream"
	lsfile "github.com/benbjohnson/litestream/file"
	"github.com/superfly/ltx"

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

func serveOne(t *testing.T, a *Agent, fx *fixture, req Request) ([]byte, Summary, Result) {
	t.Helper()
	e := fx.submit(req)
	a.consider(context.Background(), e)
	a.Wait()
	return fx.answer(req.ID)
}

func TestReadServesAQuery(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 3)
	a, w := fx.agent("uk-1")
	rows, s, r := serveOne(t, a, fx, fx.request("SELECT count(*) AS n FROM t", ""))
	if r.State != StateDone || r.Code != "" || string(rows) != "n\n30\n" || r.Box != "uk-1" {
		t.Fatalf("result %+v, rows %q", r, rows)
	}
	plan, err := dbdoor.PlanFrom(context.Background(), lsfile.NewReplicaClient(fx.replica), time.Time{})
	if err != nil || s.TXID != plan.TXID || !s.At.Equal(plan.At) {
		t.Fatalf("summary %+v, plan %+v %v", s, plan, err)
	}
	if w.dials.Load() != 1 {
		t.Fatalf("%d worker connections for one read", w.dials.Load())
	}
}

// vanishing makes full-file opens of the first planned file answer
// os.ErrNotExist, `times` times: litestream's compaction removing an L0
// file under the plan.
type vanishing struct {
	dbdoor.ReplicaReader
	mu    sync.Mutex
	times int
}

func (v *vanishing) OpenLTXFile(ctx context.Context, level int, lo, hi ltx.TXID, off, size int64) (io.ReadCloser, error) {
	v.mu.Lock()
	gone := off == 0 && size == 0 && v.times > 0
	if gone {
		v.times--
	}
	v.mu.Unlock()
	if gone {
		return nil, os.ErrNotExist
	}
	return v.ReplicaReader.OpenLTXFile(ctx, level, lo, hi, off, size)
}

// Review Focus 3: one quiet re-plan, then changed-twice, and no chunk.
func TestReadRePlansOnceThenSaysChangedTwice(t *testing.T) {
	t.Parallel()
	for _, c := range []struct {
		times int
		want  string
	}{{1, ""}, {2, CodeChangedTwice}} {
		fx := newFixture(t, 2)
		v := &vanishing{ReplicaReader: lsfile.NewReplicaClient(fx.replica), times: c.times}
		a, _ := fx.agent("uk-1", func(cfg *AgentConfig, _ *inProcWorkers) {
			cfg.OpenReplica = func(context.Context, string) (dbdoor.ReplicaReader, error) { return v, nil }
		})
		req := fx.request("SELECT count(*) FROM t", "")
		_, _, r := serveOne(t, a, fx, req)
		if r.Code != c.want {
			t.Fatalf("vanished %d times: %+v", c.times, r)
		}
		if c.want != "" && len(fx.chunkKeys(req.ID)) != 0 {
			t.Fatal("a chunk exists for a read that never ran")
		}
	}
}

// Plan review finding 10: after a vanished file, the retry waits for the
// first worker's unit to stop before it dials again, and keeps that unit's
// reservation until then. 12 GiB of 20 fits one floor at a time.
// Mutations: drop the settle call in read (the second dial sees a live unit,
// and the retired floor makes the new one no-room); or drop retry's move of
// the old reservation (TestAdmissionRetryHoldsTheOldReservation goes red).
func TestRetryWaitsForTheFirstWorkerToStop(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 2)
	v := &vanishing{ReplicaReader: lsfile.NewReplicaClient(fx.replica), times: 1}
	a, w := fx.agent("uk-1", func(c *AgentConfig, w *inProcWorkers) {
		c.OpenReplica = func(context.Context, string) (dbdoor.ReplicaReader, error) { return v, nil }
		c.Free = func() (uint64, uint64, error) { return 12 << 30, 20 << 30, nil }
		w.linger = 300 * time.Millisecond
	})
	_, _, r := serveOne(t, a, fx, fx.request("SELECT count(*) FROM t", ""))
	if r.State != StateDone {
		t.Fatalf("result %+v", r)
	}
	w.mu.Lock()
	defer w.mu.Unlock()
	if len(w.liveAtDial) != 2 || w.liveAtDial[1] != 0 {
		t.Fatalf("live units at each dial: %v; the retry dialed beside the first worker", w.liveAtDial)
	}
}

// Plan review finding 15: a worker that ends before ready (its service
// failed after the socket accepted) asks for a probe at once. Mutation:
// drop the askProbe after a pre-ready sendFile or readFrame failure;
// probeAfter stays empty and this test is red.
func TestAWorkerThatNeverStartsAsksForAProbe(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 1)
	a, _ := fx.agent("uk-1", func(_ *AgentConfig, w *inProcWorkers) { w.failStart.Store(true) })
	_, _, r := serveOne(t, a, fx, fx.request("SELECT 1", ""))
	if r.Code != CodeInterrupted || len(a.probeAfter) != 1 {
		t.Fatalf("result %+v, probe asked %d", r, len(a.probeAfter))
	}
}

// stageHook changes the record at one stage of the read: planning (the
// first listing), download (the first full-file open), or after the last
// byte has been read from the bucket (EOF of the last planned file).
type stageHook struct {
	dbdoor.ReplicaReader
	stage, lastSpan string
	once            sync.Once
	change          func()
}

func (h *stageHook) LTXFiles(ctx context.Context, level int, seek ltx.TXID, m bool) (ltx.FileIterator, error) {
	if h.stage == "planning" {
		h.once.Do(h.change)
	}
	return h.ReplicaReader.LTXFiles(ctx, level, seek, m)
}

func (h *stageHook) OpenLTXFile(ctx context.Context, level int, lo, hi ltx.TXID, off, size int64) (io.ReadCloser, error) {
	rc, err := h.ReplicaReader.OpenLTXFile(ctx, level, lo, hi, off, size)
	if err != nil || off != 0 || size != 0 {
		return rc, err
	}
	switch {
	case h.stage == "download":
		h.once.Do(h.change)
	case h.stage == "after the last byte" && dbdoor.Span(uint64(lo), uint64(hi)) == h.lastSpan:
		return &eofHook{ReadCloser: rc, at: func() { h.once.Do(h.change) }}, nil
	}
	return rc, nil
}

type eofHook struct {
	io.ReadCloser
	at func()
}

func (e *eofHook) Read(p []byte) (int, error) {
	n, err := e.ReadCloser.Read(p)
	if err == io.EOF {
		e.at()
	}
	return n, err
}

// T9: a record change at every stage ends with its code, and no chunk is
// ever written (codex 1).
func TestRecordChangeAtEveryStageWritesNoChunk(t *testing.T) {
	t.Parallel()
	no := false
	changes := map[string]struct {
		change func(fx *fixture)
		want   string
	}{
		"moved": {func(fx *fixture) { r := fx.record(); r.Region = "ie"; fx.putRecord(r) }, CodeRecordChanged},
		"moving": {func(fx *fixture) {
			r := fx.record()
			r.Move = &provision.MoveState{ID: "mv_x", To: "ie", State: "requested"}
			fx.putRecord(r)
		}, CodeMoving},
		"recreated in the account": {func(fx *fixture) { r := fx.record(); r.CreatedAt = r.CreatedAt.Add(time.Hour); fx.putRecord(r) }, CodeRecordChanged},
		"recreated in another account": {func(fx *fixture) {
			created := fx.record().CreatedAt
			fx.store.Delete(context.Background(), provision.RecordKey(fxAccount, fxApp, fxHost))
			fx.putRecordUnder("bac", provision.Record{Region: "uk", CreatedAt: created})
		}, CodeRecordChanged},
		"class changed": {func(fx *fixture) { r := fx.record(); r.Hibernate = &no; fx.putRecord(r) }, CodeRecordChanged},
	}
	for _, stage := range []string{"before the pick", "planning", "download", "after the last byte"} {
		for name, c := range changes {
			t.Run(stage+"/"+name, func(t *testing.T) {
				t.Parallel()
				fx := newFixture(t, 2)
				plan, err := dbdoor.PlanFrom(context.Background(), lsfile.NewReplicaClient(fx.replica), time.Time{})
				if err != nil {
					t.Fatal(err)
				}
				last := plan.Planned[len(plan.Planned)-1]
				hook := &stageHook{ReplicaReader: lsfile.NewReplicaClient(fx.replica), stage: stage,
					lastSpan: dbdoor.Span(last.MinTXID, last.MaxTXID), change: func() { c.change(fx) }}
				a, _ := fx.agent("uk-1", func(cfg *AgentConfig, _ *inProcWorkers) {
					cfg.OpenReplica = func(context.Context, string) (dbdoor.ReplicaReader, error) { return hook, nil }
				})
				req := fx.request("SELECT count(*) FROM t", "")
				if stage == "before the pick" {
					c.change(fx)
				}
				_, _, r := serveOne(t, a, fx, req)
				if r.State != StateRefused || r.Code != c.want {
					t.Fatalf("result %+v, want refused %s", r, c.want)
				}
				if keys := fx.chunkKeys(req.ID); len(keys) != 0 {
					t.Fatalf("chunks written: %v", keys)
				}
			})
		}
	}
}

func TestReadTooLargeAndNoRoom(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 2)
	a, _ := fx.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) { c.MaxRestored = 1024 })
	_, s, r := serveOne(t, a, fx, fx.request("SELECT 1", ""))
	if r.Code != CodeTooLarge || s.Cap != 1024 || s.Restored <= 1024 {
		t.Fatalf("too large: %+v %+v", r, s)
	}
	fx2 := newFixture(t, 2)
	b, _ := fx2.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) {
		calls := 0
		c.Free = func() (uint64, uint64, error) {
			calls++
			if calls == 1 { // the floor fits; the plan's growth does not
				return 11 << 30, 20 << 30, nil
			}
			return 5<<30 + 1, 20 << 30, nil
		}
	})
	_, _, r2 := serveOne(t, b, fx2, fx2.request("SELECT 1", ""))
	if r2.Code != CodeNoRoom {
		t.Fatalf("no room: %+v", r2)
	}
}

// Review Focus 4: the cap is exact, a row bigger than a chunk is split, and
// the summary covers exactly the kept bytes.
func TestChunkWriterCapIsExact(t *testing.T) {
	t.Parallel()
	for _, extra := range []int{0, 1} {
		fx := newFixture(t, 0)
		max := int64(3*ChunkBytes + 7)
		cw := newChunkWriter(context.Background(), fx.store, "dbr_aaaaaaaaaaaaaaaaaaaaaaaaaa", fx.reply.PublicKey(), StreamRows, time.Hour, max, ChunkBytes)
		big := bytes.Repeat([]byte("x"), 2*ChunkBytes+5)
		rest := bytes.Repeat([]byte("y"), int(max)-len(big)+extra)
		cw.Write(big)
		_, err := cw.Write(rest)
		if (extra == 1) != (err == errCapped) || cw.capped() != (extra == 1) {
			t.Fatalf("extra %d: Write = %v, capped %v", extra, err, cw.capped())
		}
		last, n, sum, err := cw.close()
		if err != nil || n != max {
			t.Fatalf("extra %d: kept %d of %d: %v", extra, n, max, err)
		}
		want := sha256.Sum256(append(append([]byte{}, big...), rest[:int(max)-len(big)]...))
		if sum != hex.EncodeToString(want[:]) {
			t.Fatal("the summary hash is not the kept bytes'")
		}
		for seq := int64(1); seq <= last; seq++ {
			sealed, _, err := fx.store.Get(context.Background(), ChunkKey("dbr_aaaaaaaaaaaaaaaaaaaaaaaaaa", seq))
			if err != nil || len(sealed)-SealOverhead > ChunkBytes {
				t.Fatalf("chunk %d: %d bytes, %v", seq, len(sealed), err)
			}
		}
	}
}

// The output cap end to end: done, output-cap, and exactly the cap kept.
func TestReadOutputCap(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 2)
	a, _ := fx.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) { c.MaxOutput = 100 })
	rows, s, r := serveOne(t, a, fx, fx.request("SELECT body FROM t", ""))
	if r.State != StateDone || r.Code != CodeOutputCap || len(rows) != 100 || s.Bytes != 100 {
		t.Fatalf("result %+v, %d bytes", r, len(rows))
	}
}

func TestCancelStopsTheOwner(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 2)
	hang := make(chan struct{})
	t.Cleanup(func() { close(hang) })
	a, _ := fx.agent("uk-1", func(c *AgentConfig, w *inProcWorkers) {
		c.CancelEvery = 20 * time.Millisecond
		w.opt.Restore = func(ctx context.Context, _ litestream.ReplicaClient, _ string, _ uint64) error {
			select {
			case <-ctx.Done():
			case <-hang:
			}
			return ctx.Err()
		}
	})
	req := fx.request("SELECT 1", "")
	e := fx.submit(req)
	a.consider(context.Background(), e)
	if _, err := fx.store.Put(context.Background(), CancelKey(req.ID), nil, ""); err != nil {
		t.Fatal(err)
	}
	done := make(chan struct{})
	go func() { a.Wait(); close(done) }()
	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("the owner did not stop within 5 s of cancel")
	}
	if _, _, r := fx.answer(req.ID); r.State != StateCancelled {
		t.Fatalf("result %+v", r)
	}
}

// refusedChunk never lets one chunk's create land: every attempt answers an
// error and the object stays absent.
type refusedChunk struct {
	release.Store
	key string
}

func (s refusedChunk) Put(ctx context.Context, key string, body []byte, ifMatch string) (string, error) {
	if key == s.key {
		return "", errors.New("connection reset")
	}
	return s.Store.Put(ctx, key, body, ifMatch)
}

// Plan review finding 9: a chunk that is never confirmed is never named by
// the result. The read ends interrupted, and every chunk through its
// last_seq exists. Mutation: flushLocked increments seq before the create is
// confirmed; the result then names chunk 2 and this test is red.
func TestResultNamesOnlyConfirmedChunks(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 1)
	req := fx.request("SELECT hex(randomblob(300000))", "") // about 600 KB: three chunks
	a, _ := fx.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) {
		c.Store = refusedChunk{Store: fx.store, key: ChunkKey(req.ID, 2)}
		c.OwnerDeadline = 2 * time.Second
	})
	a.consider(context.Background(), fx.submit(req))
	a.Wait()
	raw, _, err := fx.store.Get(context.Background(), ResultKey(req.ID))
	if err != nil {
		t.Fatal(err)
	}
	var r Result
	json.Unmarshal(raw, &r)
	if r.State != StateFailed || r.Code != CodeInterrupted || r.LastSeq != 1 {
		t.Fatalf("result %+v", r)
	}
	for seq := int64(1); seq <= r.LastSeq; seq++ {
		if _, _, err := fx.store.Get(context.Background(), ChunkKey(req.ID, seq)); err != nil {
			t.Fatalf("the result names chunk %d, which does not exist", seq)
		}
	}
}

// strictStore fails the test on any protocol delete or replacing write
// under db-reads/ (T16, the box's half).
type strictStore struct {
	release.Store
	t *testing.T
}

func (s strictStore) Put(ctx context.Context, key string, body []byte, ifMatch string) (string, error) {
	if strings.HasPrefix(key, Prefix) && ifMatch != "" {
		s.t.Errorf("a replacing write of %s", key)
	}
	return s.Store.Put(ctx, key, body, ifMatch)
}

func (s strictStore) Delete(ctx context.Context, key string) error {
	if strings.HasPrefix(key, Prefix) {
		s.t.Errorf("a delete of %s", key)
	}
	return s.Store.Delete(ctx, key)
}

func TestReadNeverDeletesOrReplaces(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 2)
	a, _ := fx.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) { c.Store = strictStore{Store: fx.store, t: t} })
	if _, _, r := serveOne(t, a, fx, fx.request("SELECT count(*) FROM t", "")); r.State != StateDone {
		t.Fatalf("%+v", r)
	}
}

func TestOwnerDeadlineEndsAHungRead(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 1)
	var started atomic.Bool
	a, _ := fx.agent("uk-1", func(c *AgentConfig, w *inProcWorkers) {
		c.OwnerDeadline = 200 * time.Millisecond
		w.opt.Restore = func(ctx context.Context, _ litestream.ReplicaClient, _ string, _ uint64) error {
			started.Store(true)
			<-ctx.Done()
			return ctx.Err()
		}
	})
	_, _, r := serveOne(t, a, fx, fx.request("SELECT 1", ""))
	if !started.Load() || r.State != StateFailed || r.Code != CodeInterrupted {
		t.Fatalf("result %+v", r)
	}
}

// Astra round 1, finding 2 (T9 extended): every replica-derived ending (a
// failed restore's error text, the restored size, the oldest point kept,
// "nothing replicated", "no room", "changed twice") waits for check B. When
// the record changed after check A, the read ends refused record-changed
// and its sealed summary carries none of that replica's details.
func TestReplicaDetailWaitsForCheckB(t *testing.T) {
	t.Parallel()
	recreate := func(fx *fixture) { r := fx.record(); r.CreatedAt = r.CreatedAt.Add(time.Hour); fx.putRecord(r) }
	for name, c := range map[string]struct {
		batches int
		at      string
		vanish  int
		mut     func(*AgentConfig, *inProcWorkers)
	}{
		"a failing restore": {batches: 2, mut: func(_ *AgentConfig, w *inProcWorkers) {
			w.opt.Restore = func(context.Context, litestream.ReplicaClient, string, uint64) error {
				return errors.New("malformed database schema (secret_customer_ledger)")
			}
		}},
		"too large":          {batches: 2, mut: func(c *AgentConfig, _ *inProcWorkers) { c.MaxRestored = 1024 }},
		"before the history": {batches: 2, at: "2000-01-01T00:00:00Z"},
		"nothing replicated": {batches: 0},
		"no room": {batches: 2, mut: func(c *AgentConfig, _ *inProcWorkers) {
			calls := 0
			c.Free = func() (uint64, uint64, error) {
				calls++
				if calls == 1 {
					return 11 << 30, 20 << 30, nil
				}
				return 5<<30 + 1, 20 << 30, nil
			}
		}},
		"changed twice": {batches: 2, vanish: 2},
	} {
		t.Run(name, func(t *testing.T) {
			t.Parallel()
			fx := newFixture(t, c.batches)
			var rd dbdoor.ReplicaReader = &stageHook{ReplicaReader: lsfile.NewReplicaClient(fx.replica), stage: "planning", change: func() { recreate(fx) }}
			if c.vanish > 0 {
				// The record changes during the second attempt's download,
				// after that attempt's check A: only a check before
				// "changed twice" can see it.
				rd = &vanishThenChange{ReplicaReader: &vanishing{ReplicaReader: lsfile.NewReplicaClient(fx.replica), times: c.vanish}, at: 2, change: func() { recreate(fx) }}
			}
			a, _ := fx.agent("uk-1", func(cfg *AgentConfig, w *inProcWorkers) {
				cfg.OpenReplica = func(context.Context, string) (dbdoor.ReplicaReader, error) { return rd, nil }
				if c.mut != nil {
					c.mut(cfg, w)
				}
			})
			_, s, r := serveOne(t, a, fx, fx.request("SELECT 1", c.at))
			if r.State != StateRefused || r.Code != CodeRecordChanged {
				t.Fatalf("result %+v, want refused %s", r, CodeRecordChanged)
			}
			if !reflect.DeepEqual(s, Summary{}) {
				t.Fatalf("the refused read's summary carries the replica's details: %+v", s)
			}
		})
	}
}

// vanishThenChange runs change on the at'th whole-file open, before the
// open answers; over vanishing, that open is one that vanishes.
type vanishThenChange struct {
	dbdoor.ReplicaReader
	mu     sync.Mutex
	n, at  int
	change func()
}

func (v *vanishThenChange) OpenLTXFile(ctx context.Context, level int, lo, hi ltx.TXID, off, size int64) (io.ReadCloser, error) {
	if off == 0 && size == 0 {
		v.mu.Lock()
		v.n++
		fire := v.n == v.at
		v.mu.Unlock()
		if fire {
			v.change()
		}
	}
	return v.ReplicaReader.OpenLTXFile(ctx, level, lo, hi, off, size)
}

// slowPicked lands the picked create and then answers late: an election
// whose response takes delay (Astra round 1, finding 3).
type slowPicked struct {
	release.Store
	delay time.Duration
}

func (s slowPicked) Put(ctx context.Context, key string, body []byte, ifMatch string) (string, error) {
	etag, err := s.Store.Put(ctx, key, body, ifMatch)
	if strings.HasSuffix(key, "/picked") {
		time.Sleep(s.delay)
	}
	return etag, err
}

// The owner works to the deadline it published in picked, not to a fresh
// one from whenever the election's answer arrived (Astra round 1, finding
// 3): a hung read ends by picked.Deadline, and a read whose election
// answered after its deadline starts no work at all.
func TestServeKeepsThePublishedDeadline(t *testing.T) {
	t.Parallel()
	for name, c := range map[string]struct {
		deadline, delay time.Duration
		dials           int64
	}{
		"a late answer inside the deadline": {deadline: 600 * time.Millisecond, delay: 300 * time.Millisecond, dials: 1},
		"an answer after the deadline":      {deadline: 200 * time.Millisecond, delay: 300 * time.Millisecond, dials: 0},
	} {
		t.Run(name, func(t *testing.T) {
			t.Parallel()
			fx := newFixture(t, 1)
			a, w := fx.agent("uk-1", func(cfg *AgentConfig, w *inProcWorkers) {
				cfg.Store = slowPicked{Store: fx.store, delay: c.delay}
				cfg.OwnerDeadline = c.deadline
				w.opt.Restore = func(ctx context.Context, _ litestream.ReplicaClient, _ string, _ uint64) error {
					<-ctx.Done()
					return ctx.Err()
				}
			})
			req := fx.request("SELECT 1", "")
			_, _, r := serveOne(t, a, fx, req)
			raw, _, err := fx.store.Get(context.Background(), PickedKey(req.ID))
			if err != nil {
				t.Fatal(err)
			}
			var p Picked
			if err := json.Unmarshal(raw, &p); err != nil {
				t.Fatal(err)
			}
			if r.State != StateFailed || r.Code != CodeInterrupted || w.dials.Load() != c.dials {
				t.Fatalf("result %+v, %d worker connections", r, w.dials.Load())
			}
			if over := r.Ended.Sub(p.Deadline); over > 150*time.Millisecond {
				t.Fatalf("the read ran %s past the deadline it published", over)
			}
		})
	}
}