package dbdoor

import (
	"context"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

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

// litestream creates <out>.tmp itself with os.Create, at the process umask,
// so a 0600 file cannot be pre-created for it. The containing directory is
// the only thing that can be private, and it must also confine the -wal
// and -shm a read-only open still needs.
func TestRestoreToUsesAPrivateDirectory(t *testing.T) {
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 2)
	parent := t.TempDir()
	got, err := RestoreTo(context.Background(), NewClient(newTestDoor(t, src).do), parent)
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(got.Dir)
	if filepath.Dir(got.Path) != got.Dir || filepath.Dir(got.Dir) != parent {
		t.Fatalf("path %s, dir %s: want the copy inside a fresh directory under %s", got.Path, got.Dir, parent)
	}
	fi, err := os.Stat(got.Dir)
	if err != nil {
		t.Fatal(err)
	}
	if perm := fi.Mode().Perm(); perm != 0o700 {
		t.Fatalf("restore directory is %04o, want 0700", perm)
	}
	if rows := countRows(t, got.Path); rows != 2*dbdoortest.Rows {
		t.Fatalf("restored %d rows, want %d", rows, 2*dbdoortest.Rows)
	}
}

// R1: the handover names the newest replicated transaction and its time,
// and the time is the LTX header's, the time of the write itself.
func TestRestoreToReportsTheNewestTransaction(t *testing.T) {
	src := t.TempDir()
	before := time.Now().Add(-time.Second)
	dbdoortest.Replicate(t, src, 3)
	door := newTestDoor(t, src)
	c := NewClient(door.do)
	files := collect(t, c, 0, 0)
	want := files[len(files)-1].MaxTXID
	got, err := RestoreTo(context.Background(), c, t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(got.Dir)
	if got.TXID != uint64(want) {
		t.Errorf("TXID = %d, want %d", got.TXID, want)
	}
	if got.At.Before(before) || got.At.After(time.Now().Add(time.Second)) {
		t.Errorf("At = %v, want the time of the last write (around now)", got.At)
	}
}

// A planned file can vanish under the restore when the writer's L0
// retention fires. ResumableReader treats ErrNotExist as FATAL and does not
// retry, so the retry has to wrap the whole Restore: a fresh plan is the
// only recovery.
func TestRestoreToRetriesTheWholeRestoreOnce(t *testing.T) {
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 2)
	door := newTestDoor(t, src)
	first := collect(t, NewClient(door.do), 0, 0)[0]
	door.gone[Span(uint64(first.MinTXID), uint64(first.MaxTXID))] = true
	door.lists = map[int]int{}
	got, err := RestoreTo(context.Background(), NewClient(door.do), t.TempDir())
	if err != nil {
		t.Fatalf("RestoreTo with a file vanishing once = %v, want success on a fresh plan", err)
	}
	defer os.RemoveAll(got.Dir)
	if door.lists[9] != 2 {
		t.Fatalf("snapshot level listed %d times, want 2 (one plan per Restore)", door.lists[9])
	}
	if rows := countRows(t, got.Path); rows != 2*dbdoortest.Rows {
		t.Fatalf("restored %d rows, want %d", rows, 2*dbdoortest.Rows)
	}
}

// Once, not forever: a file that keeps vanishing is a failure to report,
// and the private directory goes with it.
func TestRestoreToRetriesOnlyOnce(t *testing.T) {
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 2)
	door := newTestDoor(t, src)
	first := collect(t, NewClient(door.do), 0, 0)[0]
	span := Span(uint64(first.MinTXID), uint64(first.MaxTXID))
	c := NewClient(func(ctx context.Context, tail string, h http.Header) (*http.Response, error) {
		door.mu.Lock()
		door.gone[span] = true
		door.mu.Unlock()
		return door.do(ctx, tail, h)
	})
	parent := t.TempDir()
	if _, err := RestoreTo(context.Background(), c, parent); !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("RestoreTo = %v, want the vanished file reported", err)
	}
	assertEmptyDir(t, parent)
}

// "Nothing replicated yet" is its own answer, for both of the errors
// litestream gives it (spec §10), and leaves nothing behind.
func TestRestoreToNamesAnEmptyReplica(t *testing.T) {
	parent := t.TempDir()
	_, err := RestoreTo(context.Background(), NewClient(newTestDoor(t, t.TempDir()).do), parent)
	if !errors.Is(err, ErrNothingReplicated) {
		t.Fatalf("RestoreTo of an empty replica = %v, want ErrNothingReplicated", err)
	}
	assertEmptyDir(t, parent)
}

// A copy that fails lease.CheckDB is never handed over (spec §10). A
// replica litestream wrote always holds its own _litestream_seq table and
// a corrupt one fails inside Restore, so the check's refusal is injected:
// what is under test is that RestoreTo honours it.
func TestRestoreToRefusesACopyThatFailsTheCheck(t *testing.T) {
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 1)
	var checked string
	defer func(old func(context.Context, string) error) { checkDB = old }(checkDB)
	checkDB = func(_ context.Context, path string) error {
		checked = path
		return errors.New("integrity_check: page 2 is never used")
	}
	parent := t.TempDir()
	_, err := RestoreTo(context.Background(), NewClient(newTestDoor(t, src).do), parent)
	if !errors.Is(err, ErrCheckFailed) {
		t.Fatalf("RestoreTo of a copy failing its check = %v, want ErrCheckFailed", err)
	}
	if filepath.Base(checked) != restoredName {
		t.Fatalf("checked %q, want the restored copy", checked)
	}
	assertEmptyDir(t, parent)
}

func assertEmptyDir(t *testing.T, dir string) {
	t.Helper()
	ents, err := os.ReadDir(dir)
	if err != nil {
		t.Fatal(err)
	}
	if len(ents) != 0 {
		t.Fatalf("%s holds %d entries after a failed restore, want none (no copy of production left behind)", dir, len(ents))
	}
}

// countingBody tracks how many door responses are open at once.
type countingBody struct {
	io.ReadCloser
	open *atomic.Int32
	once sync.Once
}

func (b *countingBody) Close() error {
	b.once.Do(func() { b.open.Add(-1) })
	return b.ReadCloser.Close()
}

// The door bounds a person's concurrent file streams, and litestream's
// compactor holds every planned file open until the restore ends. Streaming
// straight through would need one slot per file, so RestoreTo downloads each
// file whole into its private directory before handing it on: one stream at
// a time, whatever the plan's length.
func TestRestoreToHoldsOneStreamAtATime(t *testing.T) {
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 4)
	door := newTestDoor(t, src)
	var open, peak atomic.Int32
	do := func(ctx context.Context, tail string, h http.Header) (*http.Response, error) {
		res, err := door.do(ctx, tail, h)
		if err != nil || !strings.Contains(strings.TrimPrefix(tail, "/ltx/"), "/") {
			return res, err // listings are read and closed at once
		}
		n := open.Add(1)
		for p := peak.Load(); n > p && !peak.CompareAndSwap(p, n); p = peak.Load() {
		}
		res.Body = &countingBody{ReadCloser: res.Body, open: &open}
		return res, nil
	}
	got, err := RestoreTo(context.Background(), NewClient(do), t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(got.Dir)
	if p := peak.Load(); p != 1 {
		t.Fatalf("restore held %d file streams open at once, want 1", p)
	}
	if rows := countRows(t, got.Path); rows != 4*dbdoortest.Rows {
		t.Fatalf("restored %d rows, want %d", rows, 4*dbdoortest.Rows)
	}
	if ents, _ := os.ReadDir(got.Dir); len(ents) > 3 {
		t.Fatalf("downloaded files left in the private directory: %v", ents)
	}
}

// flakyBody hands over at most limit bytes and then fails the way a reset
// connection does.
type flakyBody struct {
	io.ReadCloser
	left int
}

func (b *flakyBody) Read(p []byte) (int, error) {
	if b.left <= 0 {
		return 0, io.ErrUnexpectedEOF
	}
	if len(p) > b.left {
		p = p[:b.left]
	}
	n, err := b.ReadCloser.Read(p)
	b.left -= n
	return n, err
}

// flaky wraps a door so every file response breaks after limit bytes.
func flaky(door *testDoor, limit int, gets *atomic.Int32) Doer {
	return func(ctx context.Context, tail string, h http.Header) (*http.Response, error) {
		res, err := door.do(ctx, tail, h)
		if err != nil || !strings.Contains(strings.TrimPrefix(tail, "/ltx/"), "/") {
			return res, err
		}
		gets.Add(1)
		res.Body = &flakyBody{ReadCloser: res.Body, left: limit}
		return res, nil
	}
}

// Over a link that resets every few hundred bytes, a download keeps what it
// has and asks for the rest, so a file larger than any one connection
// survives still arrives.
func TestRestoreToResumesAFileAcrossResets(t *testing.T) {
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 2)
	var gets atomic.Int32
	got, err := RestoreTo(context.Background(), NewClient(flaky(newTestDoor(t, src), 100, &gets)), t.TempDir())
	if err != nil {
		t.Fatalf("RestoreTo over a resetting link = %v, want it to resume and finish", err)
	}
	defer os.RemoveAll(got.Dir)
	if rows := countRows(t, got.Path); rows != 2*dbdoortest.Rows {
		t.Fatalf("restored %d rows, want %d", rows, 2*dbdoortest.Rows)
	}
}

// A link that delivers nothing is given up on after a bounded number of
// attempts, not retried forever.
func TestRestoreToGivesUpOnALinkThatMakesNoProgress(t *testing.T) {
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 1)
	var gets atomic.Int32
	parent := t.TempDir()
	if _, err := RestoreTo(context.Background(), NewClient(flaky(newTestDoor(t, src), 0, &gets)), parent); err == nil {
		t.Fatal("RestoreTo over a dead link succeeded")
	}
	if n := gets.Load(); n > 40 {
		t.Fatalf("made %d attempts at a link that never delivered a byte", n)
	}
	assertEmptyDir(t, parent)
}

// Ctrl-C after the last file has arrived still stops the restore. The
// local half (litestream's compaction over the downloaded files) does not
// check the context itself, so the downloaded files do, and RestoreTo
// checks again before CheckDB. Nothing is handed over, nothing left.
func TestRestoreToStopsWhenCancelledAfterTheDownloads(t *testing.T) {
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 3)
	door := newTestDoor(t, src)
	files := len(collect(t, NewClient(door.do), 0, 0))
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var gets atomic.Int32
	do := func(c context.Context, tail string, h http.Header) (*http.Response, error) {
		res, err := door.do(c, tail, h)
		if err == nil && strings.Contains(strings.TrimPrefix(tail, "/ltx/"), "/") && int(gets.Add(1)) == files {
			cancel() // the last download is in hand; only local work is left
		}
		return res, err
	}
	parent := t.TempDir()
	if _, err := RestoreTo(ctx, NewClient(do), parent); !errors.Is(err, context.Canceled) {
		t.Fatalf("RestoreTo cancelled after its downloads = %v, want context.Canceled", err)
	}
	assertEmptyDir(t, parent)
}

// answer is a canned door response.
func answer(status int, retryAfter string) *http.Response {
	rec := httptest.NewRecorder()
	if retryAfter != "" {
		rec.Header().Set("Retry-After", retryAfter)
	}
	rec.WriteHeader(status)
	io.WriteString(rec, `{"error":"busy"}`)
	return rec.Result()
}

// rangeStart is the offset a request's Range header asks for.
func rangeStart(h http.Header) int64 {
	off, _, _ := ParseRange(h.Get("Range"))
	return off
}

// A reopen of the rest of a file that meets a busy door (503, or 429 with
// Retry-After) is retried in place, keeping the bytes already downloaded,
// rather than thrown back to litestream, which would start the file again.
func TestRestoreToRetriesABusyReopenWithoutLosingProgress(t *testing.T) {
	defer func(old time.Duration) { dbDownloadBackoff = old }(dbDownloadBackoff)
	dbDownloadBackoff = time.Millisecond
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 2)
	door := newTestDoor(t, src)
	var mu sync.Mutex
	busy := map[string]bool{}
	var restarts atomic.Int32
	reached := map[string]int64{}
	do := func(ctx context.Context, tail string, h http.Header) (*http.Response, error) {
		if !strings.Contains(strings.TrimPrefix(tail, "/ltx/"), "/") {
			return door.do(ctx, tail, h)
		}
		off := rangeStart(h)
		key := tail + "@" + strconv.FormatInt(off, 10)
		mu.Lock()
		if off < reached[tail] {
			restarts.Add(1)
		}
		reached[tail] = max(reached[tail], off)
		first := !busy[key]
		busy[key] = true
		mu.Unlock()
		if off > 0 && first {
			return answer(http.StatusServiceUnavailable, ""), nil
		}
		res, err := door.do(ctx, tail, h)
		if err == nil {
			res.Body = &flakyBody{ReadCloser: res.Body, left: 100}
		}
		return res, err
	}
	got, err := RestoreTo(context.Background(), NewClient(do), t.TempDir())
	if err != nil {
		t.Fatalf("RestoreTo past busy reopens = %v, want it to retry in place and finish", err)
	}
	defer os.RemoveAll(got.Dir)
	if n := restarts.Load(); n != 0 {
		t.Fatalf("%d requests went back to an earlier offset: progress was thrown away", n)
	}
	if rows := countRows(t, got.Path); rows != 2*dbdoortest.Rows {
		t.Fatalf("restored %d rows, want %d", rows, 2*dbdoortest.Rows)
	}
}

// A 429's Retry-After is honoured before the next attempt.
func TestRestoreToHonoursRetryAfter(t *testing.T) {
	defer func(old time.Duration) { dbDownloadBackoff = old }(dbDownloadBackoff)
	dbDownloadBackoff = time.Millisecond
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 1)
	door := newTestDoor(t, src)
	var mu sync.Mutex
	var refusedAt, nextAt time.Time
	do := func(ctx context.Context, tail string, h http.Header) (*http.Response, error) {
		if !strings.Contains(strings.TrimPrefix(tail, "/ltx/"), "/") {
			return door.do(ctx, tail, h)
		}
		mu.Lock()
		defer mu.Unlock()
		if rangeStart(h) > 0 && refusedAt.IsZero() {
			refusedAt = time.Now()
			return answer(http.StatusTooManyRequests, "1"), nil
		}
		if !refusedAt.IsZero() && nextAt.IsZero() {
			nextAt = time.Now()
		}
		res, err := door.do(ctx, tail, h)
		if err == nil && rangeStart(h) == 0 {
			res.Body = &flakyBody{ReadCloser: res.Body, left: 100}
		}
		return res, err
	}
	got, err := RestoreTo(context.Background(), NewClient(do), t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(got.Dir)
	if refusedAt.IsZero() || nextAt.Sub(refusedAt) < 900*time.Millisecond {
		t.Fatalf("retried %v after a 429 with Retry-After: 1", nextAt.Sub(refusedAt))
	}
}

// The wait between attempts ends the moment the person stops.
func TestRestoreToBackoffIsCancellable(t *testing.T) {
	defer func(old time.Duration) { dbDownloadBackoff = old }(dbDownloadBackoff)
	dbDownloadBackoff = time.Hour
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 1)
	door := newTestDoor(t, src)
	do := func(ctx context.Context, tail string, h http.Header) (*http.Response, error) {
		if strings.Contains(strings.TrimPrefix(tail, "/ltx/"), "/") && rangeStart(h) > 0 {
			return answer(http.StatusServiceUnavailable, ""), nil
		}
		res, err := door.do(ctx, tail, h)
		if err == nil && strings.Contains(strings.TrimPrefix(tail, "/ltx/"), "/") {
			res.Body = &flakyBody{ReadCloser: res.Body, left: 100}
		}
		return res, err
	}
	ctx, cancel := context.WithCancel(context.Background())
	time.AfterFunc(200*time.Millisecond, cancel)
	start := time.Now()
	parent := t.TempDir()
	if _, err := RestoreTo(ctx, NewClient(do), parent); !errors.Is(err, context.Canceled) {
		t.Fatalf("RestoreTo cancelled during a backoff = %v, want context.Canceled", err)
	}
	if d := time.Since(start); d > 5*time.Second {
		t.Fatalf("the stop took %v to land", d)
	}
	assertEmptyDir(t, parent)
}

// A door that stays busy is given up on: only attempts that deliver no
// byte count, and there are a bounded number of them.
func TestRestoreToGivesUpOnAPermanentlyBusyDoor(t *testing.T) {
	defer func(old time.Duration) { dbDownloadBackoff = old }(dbDownloadBackoff)
	dbDownloadBackoff = time.Millisecond
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 1)
	door := newTestDoor(t, src)
	var reopens atomic.Int32
	do := func(ctx context.Context, tail string, h http.Header) (*http.Response, error) {
		if strings.Contains(strings.TrimPrefix(tail, "/ltx/"), "/") && rangeStart(h) > 0 {
			reopens.Add(1)
			return answer(http.StatusServiceUnavailable, ""), nil
		}
		res, err := door.do(ctx, tail, h)
		if err == nil && strings.Contains(strings.TrimPrefix(tail, "/ltx/"), "/") {
			res.Body = &flakyBody{ReadCloser: res.Body, left: 100}
		}
		return res, err
	}
	parent := t.TempDir()
	if _, err := RestoreTo(context.Background(), NewClient(do), parent); err == nil {
		t.Fatal("RestoreTo against a door that never recovers succeeded")
	}
	if n := reopens.Load(); n == 0 || n > 100 {
		t.Fatalf("%d reopens against a permanently busy door", n)
	}
	assertEmptyDir(t, parent)
}

// A stop that lands while the copy is being checked still withdraws it:
// RestoreTo checks the context after CheckDB and hands nothing over.
func TestRestoreToStopsWhenCancelledDuringTheCheck(t *testing.T) {
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 1)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	defer func(old func(context.Context, string) error) { checkDB = old }(checkDB)
	var sawCtx bool
	checkDB = func(c context.Context, path string) error {
		sawCtx = c == ctx
		cancel() // the stop arrives as the check finishes
		return nil
	}
	parent := t.TempDir()
	if _, err := RestoreTo(ctx, NewClient(newTestDoor(t, src).do), parent); !errors.Is(err, context.Canceled) {
		t.Fatalf("RestoreTo cancelled during the check = %v, want context.Canceled", err)
	}
	if !sawCtx {
		t.Fatal("the check did not get the restore's context")
	}
	assertEmptyDir(t, parent)
}

// The FIRST open of a file meeting a full door (429) waits and asks again,
// the same as a reopen does, instead of failing straight back to
// litestream, whose own retries ignore Retry-After.
func TestRestoreToRetriesABusyFirstOpen(t *testing.T) {
	defer func(old time.Duration) { dbDownloadBackoff = old }(dbDownloadBackoff)
	dbDownloadBackoff = time.Millisecond
	src := t.TempDir()
	dbdoortest.Replicate(t, src, 2)
	door := newTestDoor(t, src)
	var mu sync.Mutex
	refused := map[string]int{}
	do := func(ctx context.Context, tail string, h http.Header) (*http.Response, error) {
		if !strings.Contains(strings.TrimPrefix(tail, "/ltx/"), "/") {
			return door.do(ctx, tail, h)
		}
		mu.Lock()
		n := refused[tail]
		refused[tail] = n + 1
		mu.Unlock()
		if rangeStart(h) == 0 && n < 4 {
			return answer(http.StatusTooManyRequests, ""), nil
		}
		return door.do(ctx, tail, h)
	}
	got, err := RestoreTo(context.Background(), NewClient(do), t.TempDir())
	if err != nil {
		t.Fatalf("RestoreTo past four busy first opens per file = %v, want it to wait and finish", err)
	}
	defer os.RemoveAll(got.Dir)
	if rows := countRows(t, got.Path); rows != 2*dbdoortest.Rows {
		t.Fatalf("restored %d rows, want %d", rows, 2*dbdoortest.Rows)
	}
}

func TestDownloadTempFileFailure(t *testing.T) {
	for _, supplied := range []bool{false, true} {
		t.Run(strconv.FormatBool(supplied), func(t *testing.T) {
			var rc io.ReadCloser
			var open atomic.Int32
			if supplied {
				open.Store(1)
				rc = &countingBody{ReadCloser: io.NopCloser(strings.NewReader("unused")), open: &open}
			}
			r := &recorder{dir: filepath.Join(t.TempDir(), "missing")}
			got, err := r.download(t.Context(), rc, 0, 1, 1, 0, 0)
			if got != nil || !errors.Is(err, os.ErrNotExist) {
				t.Fatalf("download = %v, %v", got, err)
			}
			if open.Load() != 0 {
				t.Fatal("supplied reader leaked")
			}
		})
	}
}