package dbreads

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"sync"
	"testing"
	"time"

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

	"github.com/carlosframework/platform/internal/dbdoor"
	"github.com/carlosframework/platform/internal/dbdoor/dbdoortest"
)

// sendPullPlan sends a pull header and the planned files; it returns the
// files' bytes concatenated, which is what the worker must send back.
// net.Pipe to the real worker code: not acceptance (Task 5's root-runner
// test is).
func sendPullPlan(t *testing.T, w io.Writer, dir string, corrupt bool) (dbdoor.RestorePlan, []byte) {
	t.Helper()
	rd := lsfile.NewReplicaClient(dir)
	plan, err := dbdoor.PlanFrom(context.Background(), rd, time.Time{})
	if err != nil {
		t.Fatal(err)
	}
	if err := writeJSONFrame(w, frameHeader, planHeader{V: ProtocolVersion, Op: OpPull, Files: plan.Planned, Target: plan.TXID, MaxRestored: DefaultMaxRestored}); err != nil {
		t.Fatal(err)
	}
	var all bytes.Buffer
	for i, f := range plan.Planned {
		rc, err := rd.OpenLTXFile(context.Background(), f.Level, ltx.TXID(f.MinTXID), ltx.TXID(f.MaxTXID), 0, 0)
		if err != nil {
			t.Fatal(err)
		}
		b, _ := io.ReadAll(rc)
		rc.Close()
		if corrupt && i == 0 {
			b[len(b)/2] ^= 0xff
		}
		all.Write(b)
		if err := writeFrame(w, frameFile, b); err != nil {
			t.Fatal(err)
		}
	}
	return plan, all.Bytes()
}

func TestWorkerHoldsThenSendsThePlannedFiles(t *testing.T) {
	t.Parallel()
	dir := t.TempDir()
	dbdoortest.Replicate(t, dir, 3)
	conn, done := workerPipe(t, WorkerOptions{})
	plan, want := sendPullPlan(t, conn, dir, false)
	var ready readyMsg
	readTyped(t, conn, frameReady, &ready)
	if ready.TXID != plan.TXID {
		t.Fatalf("ready at %d", ready.TXID)
	}
	if err := writeJSONFrame(conn, frameRun, runMsg{}); err != nil {
		t.Fatal(err)
	}
	var got bytes.Buffer
	for {
		typ, p, err := readFrame(conn)
		if err != nil {
			t.Fatal(err)
		}
		if typ == frameOutput {
			got.Write(p)
			continue
		}
		var end endMsg
		if typ != frameEnd || json.Unmarshal(p, &end) != nil || end.Code != "" {
			t.Fatalf("end %q %s", typ, p)
		}
		break
	}
	if !bytes.Equal(got.Bytes(), want) {
		t.Fatalf("sent back %d bytes, held %d", got.Len(), len(want))
	}
	if err := <-done; err != nil {
		t.Fatal(err)
	}
}

// Review Focus 2: a corrupt file is refused before ready, never shipped.
// Mutation: holdAndSend skips verifyLTX (red: ready, not check-failed).
func TestWorkerHoldRefusesACorruptFile(t *testing.T) {
	t.Parallel()
	dir := t.TempDir()
	dbdoortest.Replicate(t, dir, 2)
	conn, done := workerPipe(t, WorkerOptions{})
	sendPullPlan(t, conn, dir, true)
	var end endMsg
	readTyped(t, conn, frameEnd, &end)
	if end.Code != CodeCheckFailed {
		t.Fatalf("end = %+v", end)
	}
	<-done
}

func TestWorkerHoldStopsWithoutSending(t *testing.T) {
	t.Parallel()
	dir := t.TempDir()
	dbdoortest.Replicate(t, dir, 1)
	conn, done := workerPipe(t, WorkerOptions{})
	sendPullPlan(t, conn, dir, false)
	readTyped(t, conn, frameReady, nil)
	if err := writeFrame(conn, frameStop, nil); err != nil {
		t.Fatal(err)
	}
	if err := <-done; err != nil {
		t.Fatal(err)
	}
}

func (fx *fixture) pullRequest(at string) Request {
	req := fx.request("", at)
	req.Op = OpPull
	return req
}

// answerPull opens a pull's chunks as the CLI does: the ltx stream.
func (fx *fixture) answerPull(id string) ([]byte, Summary, Result) {
	fx.t.Helper()
	raw, _, err := fx.store.Get(context.Background(), ResultKey(id))
	if err != nil {
		fx.t.Fatal(err)
	}
	var r Result
	if err := json.Unmarshal(raw, &r); err != nil {
		fx.t.Fatal(err)
	}
	s, err := OpenSummary(fx.reply, id, r.Sealed)
	if err != nil {
		fx.t.Fatal(err)
	}
	var out bytes.Buffer
	for seq := int64(1); seq <= r.LastSeq; seq++ {
		sealed, _, _ := fx.store.Get(context.Background(), ChunkKey(id, seq))
		plain, err := OpenChunk(fx.reply, id, seq, StreamLTX, seq == r.LastSeq, sealed)
		if err != nil {
			fx.t.Fatalf("chunk %d: %v", seq, err)
		}
		out.Write(plain)
	}
	return out.Bytes(), s, r
}

// plannedBytes is dir's planned files, concatenated in plan order.
func plannedBytes(t *testing.T, dir string) (dbdoor.RestorePlan, []byte) {
	t.Helper()
	rd := lsfile.NewReplicaClient(dir)
	plan, err := dbdoor.PlanFrom(context.Background(), rd, time.Time{})
	if err != nil {
		t.Fatal(err)
	}
	var want bytes.Buffer
	for _, f := range plan.Planned {
		rc, err := rd.OpenLTXFile(context.Background(), f.Level, ltx.TXID(f.MinTXID), ltx.TXID(f.MaxTXID), 0, 0)
		if err != nil {
			t.Fatal(err)
		}
		io.Copy(&want, rc)
		rc.Close()
	}
	return plan, want.Bytes()
}

// The agent's pull ships the planned files whole, in plan order, as the
// ltx stream in chunks of at most PullChunkBytes, with the file list in the
// summary, and logs plan 1's two timing lines (plan review finding 26).
// Mutation: drop the timing lines (red).
func TestPullShipsThePlannedFiles(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 4)
	var logMu sync.Mutex
	var logs []string
	a, _ := fx.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) {
		c.Logf = func(format string, args ...any) {
			logMu.Lock()
			defer logMu.Unlock()
			logs = append(logs, fmt.Sprintf(format, args...))
		}
	})
	req := fx.pullRequest("")
	fx.submit(req)
	a.Pass(context.Background())
	a.Wait()
	got, s, r := fx.answerPull(req.ID)
	plan, want := plannedBytes(t, fx.replica)
	if r.State != StateDone || !bytes.Equal(got, want) || s.TXID != plan.TXID || len(s.Files) != len(plan.Planned) || s.Bytes != plan.LTXBytes {
		t.Fatalf("result %+v summary %+v: %d bytes, want %d", r, s, len(got), len(want))
	}
	for i, f := range s.Files {
		if f.Level != plan.Planned[i].Level || f.MinTXID != plan.Planned[i].MinTXID || f.MaxTXID != plan.Planned[i].MaxTXID || f.Size != plan.Planned[i].Size {
			t.Fatalf("file %d is %+v, planned %+v", i, f, plan.Planned[i])
		}
	}
	for seq := int64(1); seq < r.LastSeq; seq++ {
		sealed, _, _ := fx.store.Get(context.Background(), ChunkKey(req.ID, seq))
		if len(sealed)-SealOverhead > PullChunkBytes {
			t.Fatalf("chunk %d is %d bytes", seq, len(sealed))
		}
	}
	// Task 1's measurement reads these two lines (plan review finding 26).
	logMu.Lock()
	all := strings.Join(logs, "\n")
	logMu.Unlock()
	if !strings.Contains(all, "and held through") || !strings.Contains(all, "confirmed chunks in") {
		t.Fatalf("the pull's timing lines are missing:\n%s", all)
	}
}

// A pull over the agent's cap ends too-large with the planned bytes and the
// cap, and a record change after the last byte left the bucket (T9 for a
// pull) writes no chunk and never sends the worker run. Mutation: the
// agent's pull sends run before check B (red: the recorded frames hold a
// run).
func TestPullTooLargeAndNoChunkBeforeCheckB(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 2)
	a, _ := fx.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) { c.MaxPull = 100 })
	req := fx.pullRequest("")
	fx.submit(req)
	a.Pass(context.Background())
	a.Wait()
	_, s, r := fx.answerPull(req.ID)
	if r.Code != CodeTooLarge || s.Cap != 100 || s.Restored <= 100 {
		t.Fatalf("%+v %+v", r, s)
	}
	fx2 := newFixture(t, 2)
	plan, _ := plannedBytes(t, fx2.replica)
	last := plan.Planned[len(plan.Planned)-1]
	hook := &stageHook{ReplicaReader: lsfile.NewReplicaClient(fx2.replica), stage: "after the last byte",
		lastSpan: dbdoor.Span(last.MinTXID, last.MaxTXID), change: func() { r := fx2.record(); r.Region = "ie"; fx2.putRecord(r) }}
	b, w := fx2.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) {
		c.OpenReplica = func(context.Context, string) (dbdoor.ReplicaReader, error) { return hook, nil }
	})
	req2 := fx2.pullRequest("")
	fx2.submit(req2)
	b.Pass(context.Background())
	b.Wait()
	if _, _, r := fx2.answerPull(req2.ID); r.Code != CodeRecordChanged || len(fx2.chunkKeys(req2.ID)) != 0 {
		t.Fatalf("%+v, chunks %v", r, fx2.chunkKeys(req2.ID))
	}
	// The worker never heard run: check B failed, so the agent said stop
	// (plan review finding 24). Without this, an agent that sent run first
	// and refused the files afterwards would pass the two checks above.
	sent := w.sentFrames()
	if !bytes.Contains(sent, []byte{frameStop}) {
		t.Fatalf("the agent never said stop: %q", sent)
	}
	if bytes.Contains(sent, []byte{frameRun}) {
		t.Fatalf("the agent sent run before check B passed: %q", sent)
	}
}

// Plan review finding 24: NewAgent itself clamps the pull's cap, whatever a
// box's settings hand it. Mutation: NewAgent keeps a MaxPull above
// MaxPullBytes (red).
func TestNewAgentClampsMaxPull(t *testing.T) {
	t.Parallel()
	fx := newFixture(t, 1)
	for give, want := range map[int64]int64{0: MaxPullBytes, -1: MaxPullBytes, 1 << 20: 1 << 20, MaxPullBytes + 1: MaxPullBytes} {
		a, _ := fx.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) { c.MaxPull = give })
		if a.cfg.MaxPull != want {
			t.Errorf("MaxPull %d became %d, want %d", give, a.cfg.MaxPull, want)
		}
	}
	// Never more than the worker holds: PlanBytesFactor times MaxRestored.
	a, _ := fx.agent("uk-1", func(c *AgentConfig, _ *inProcWorkers) { c.MaxRestored = 256 << 20 })
	if a.cfg.MaxPull != PlanBytesFactor*(256<<20) {
		t.Errorf("MaxPull %d with a 256 MiB restored cap", a.cfg.MaxPull)
	}
}

// relabelled is the bytes of dir's first planned file, a valid snapshot at
// transaction 1, under the coordinates of a file that ends at 2.
func relabelled(t *testing.T, dir string) (dbdoor.LTXFile, []byte) {
	t.Helper()
	rd := lsfile.NewReplicaClient(dir)
	rc, err := rd.OpenLTXFile(context.Background(), 0, 1, 1, 0, 0)
	if err != nil {
		t.Fatal(err)
	}
	b, _ := io.ReadAll(rc)
	rc.Close()
	return dbdoor.LTXFile{Level: 0, MinTXID: 1, MaxTXID: 2, Size: int64(len(b))}, b
}

// Astra round 1 on plan 3, finding 1: a valid snapshot under another file's
// coordinates is refused before ready, for a pull and for a query alike: the
// worker checks each file's LTX header against its plan entry, so a copy
// can never be reported at a transaction it does not hold. Mutation: drop
// the header comparison (red: the query restores and answers ready at 2;
// the pull verifies and answers ready).
func TestWorkerRefusesARelabelledSnapshot(t *testing.T) {
	t.Parallel()
	dir := t.TempDir()
	dbdoortest.Replicate(t, dir, 2)
	for _, op := range []string{OpQuery, OpPull} {
		lf, b := relabelled(t, dir)
		conn, done := workerPipe(t, WorkerOptions{})
		if err := writeJSONFrame(conn, frameHeader, planHeader{V: ProtocolVersion, Op: op, Files: []dbdoor.LTXFile{lf}, Target: 2, MaxRestored: DefaultMaxRestored}); err != nil {
			t.Fatal(err)
		}
		if err := writeFrame(conn, frameFile, b); err != nil {
			t.Fatal(err)
		}
		typ, p, err := readFrame(conn)
		var end endMsg
		if err != nil || typ != frameEnd || json.Unmarshal(p, &end) != nil || end.Code != CodeCheckFailed {
			t.Errorf("%s: a relabelled snapshot answered %q %s %v", op, typ, p, err)
			writeFrame(conn, frameStop, nil)
		}
		<-done
	}
}