package dbdoor

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"

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

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

var (
	// ErrNothingReplicated: the replica holds no restorable state yet.
	// litestream says so two ways (ErrNoSnapshots, ErrTxNotAvailable), the
	// pair lease.Restore already treats as "nothing to restore".
	ErrNothingReplicated = errors.New("dbdoor: nothing has been replicated yet")
	// ErrCheckFailed: the copy restored, but lease.CheckDB refused it. It
	// is never handed over.
	ErrCheckFailed = errors.New("dbdoor: the restored copy failed its check")
	// ErrHistoryGone: a restore target older than the replica's retained
	// history (seeded-instances spec §4 step 4).
	ErrHistoryGone = errors.New("dbdoor: history before the target is no longer kept")
	// ErrBeforeHistory: an as-of restore asked for a time earlier than
	// anything the replica keeps. RestoreAsOf returns it as a
	// *BeforeHistoryError naming the oldest point there is.
	ErrBeforeHistory = errors.New("dbdoor: the time is earlier than the history the replica keeps")
	// ErrNoTimes: an as-of restore met a listing with no file times (the
	// claim-mode door sends none). litestream reads a zero time as "before
	// T" and would restore the newest copy, so this refuses instead.
	ErrNoTimes = errors.New("dbdoor: the replica's listing carries no times, so it cannot be read as of a time")
)

// EarliestAsOf is the earliest time RestoreAsOf asks litestream for. Any
// earlier request is moved to it: it is before every replica's history,
// and unlike Go's zero time it cannot mean "the newest copy".
var EarliestAsOf = time.Unix(0, 0).UTC()

// BeforeHistoryError is ErrBeforeHistory with the oldest point the replica
// keeps: the earliest time an as-of restore can ask for, the first whole
// second after the oldest file a restore could start from was listed.
type BeforeHistoryError struct {
	Oldest time.Time
}

func (e *BeforeHistoryError) Error() string {
	return fmt.Sprintf("%v (the oldest point kept is %s)", ErrBeforeHistory, e.Oldest.UTC().Format(time.RFC3339))
}

func (e *BeforeHistoryError) Is(target error) bool { return target == ErrBeforeHistory }

// Restored is a private copy of a replica on this machine.
type Restored struct {
	// Path is the restored database. It lives inside Dir.
	Path string
	// Dir is the private 0700 directory holding Path and any -wal/-shm a
	// read-only open makes. The caller removes it.
	Dir string
	// TXID is the newest replicated transaction in the copy, and At is
	// that transaction's time from its LTX header: the time of the write,
	// not of a later compaction (spec §8.1).
	TXID uint64
	At   time.Time
}

// restoredName is the copy's file name inside its private directory.
const restoredName = "replica.db"

// checkDB is lease.CheckDBContext, the one check a disaster recovery leans on. A
// variable only so a test can make it refuse.
var checkDB = lease.CheckDBContext

// RestoreTo restores the replica c reads into a new 0700 directory under
// parent (os.TempDir() when empty) and proves it with lease.CheckDB.
//
// The three conditions spec §12 rests on hold here: c is not a
// ReplicaClientV3, RestoreOptions.Follow is never set, and CheckDB runs on
// the result rather than relying on litestream's own IntegrityCheck.
//
// A planned file can vanish under the restore when the writer's L0
// retention fires, and litestream's ResumableReader treats that as fatal.
// So the retry wraps the WHOLE Restore, once: a fresh plan is the only
// recovery. On any failure the directory is removed, so no copy of
// production data is left behind.
func RestoreTo(ctx context.Context, c *Client, parent string) (*Restored, error) {
	return restoreAt(ctx, c, parent, 0, time.Time{})
}

// RestoreToTarget is RestoreTo at a target transaction: 0 is the latest.
//
// A nonzero target restores exactly to target or not at all. litestream
// v0.5.14's CalcRestorePlan skips every file whose MaxTXID is past target
// and refuses unless the chain it builds ends exactly at target, so the
// copy never reaches past it and Restored.TXID is target. A target no
// retained chain ends at (older than the retained history, inside a
// compacted file, or past the latest transaction) is refused as
// ErrHistoryGone. Ask for a PlanRestore's TXID: that restore produces
// exactly the plan's RestoredBytes.
func RestoreToTarget(ctx context.Context, c *Client, parent string, target uint64) (*Restored, error) {
	return restoreAt(ctx, c, parent, target, time.Time{})
}

// RestoreAsOf is RestoreTo stopped at a point in time: the copy holds every
// file replicated before at and nothing after it. A zero at is the newest
// copy, which is RestoreTo.
//
// The plan is litestream's own (Timestamp set, TXID not). It keeps a file
// only if its listed time is before at, and the door lists S3's
// LastModified, the upload time, which is never earlier than the writes a
// file holds. at is rounded down to its second first (restoreOnce says
// why). So the copy can be older than at, never newer. Restored.At
// says what it actually holds.
func RestoreAsOf(ctx context.Context, c *Client, parent string, at time.Time) (*Restored, error) {
	// A requested time that rounds down to Go's zero time would read as
	// "the newest copy" to litestream. Anything that early is before every
	// replica's history, so it is asked as the Unix epoch, which still is.
	if !at.IsZero() && at.Truncate(time.Second).Before(EarliestAsOf) {
		at = EarliestAsOf
	}
	return restoreAt(ctx, c, parent, 0, at)
}

// restoreAt is the one restore behind RestoreTo, RestoreToTarget and
// RestoreAsOf. At most one of target and at is set: a target transaction
// (RestoreToTarget) or an as-of time (RestoreAsOf). Neither is the newest
// copy. Each keeps its own refusal: a target no chain ends at is
// ErrHistoryGone, and a time before the kept history is ErrBeforeHistory.
func restoreAt(ctx context.Context, c *Client, parent string, target uint64, at time.Time) (*Restored, error) {
	dir, err := os.MkdirTemp(parent, "carlos-db-")
	if err != nil {
		return nil, err
	}
	// MkdirTemp already creates 0700; say it anyway, because the privacy of
	// everything restored rests on this one mode.
	if err := os.Chmod(dir, 0o700); err != nil {
		os.RemoveAll(dir)
		return nil, err
	}
	out := filepath.Join(dir, restoredName)
	var rec *recorder
	for attempt := 0; ; attempt++ {
		rec = &recorder{Client: c, dir: dir, at: at}
		err = restoreOnce(ctx, rec, out, target, at)
		if err == nil {
			break
		}
		removeCopy(out)
		rec.mu.Lock()
		checkFailed := rec.checkFailed
		rec.mu.Unlock()
		if checkFailed {
			os.RemoveAll(dir)
			if cerr := ctx.Err(); cerr != nil {
				return nil, cerr
			}
			return nil, ErrCheckFailed
		}
		vanished := errors.Is(err, os.ErrNotExist) || rec.sawVanish()
		if attempt == 0 && vanished && ctx.Err() == nil {
			continue
		}
		os.RemoveAll(dir)
		// A stop reaches litestream's local half as a broken read, which
		// its compactor reports without wrapping. Name the real cause.
		if cerr := ctx.Err(); cerr != nil {
			return nil, cerr
		}
		switch {
		case errors.Is(err, ErrNoTimes):
			return nil, ErrNoTimes
		case errors.Is(err, litestream.ErrNoSnapshots) || errors.Is(err, litestream.ErrTxNotAvailable):
			if target != 0 {
				return nil, ErrHistoryGone
			}
			if oldest := rec.oldestStart(); !at.IsZero() && !oldest.IsZero() {
				if !oldest.Before(at.Truncate(time.Second)) {
					// The first whole second after the oldest file: the
					// earliest --at that, rounded down, still includes it.
					return nil, &BeforeHistoryError{Oldest: oldest.Truncate(time.Second).Add(time.Second)}
				}
				return nil, err
			}
			return nil, ErrNothingReplicated
		case vanished && !errors.Is(err, os.ErrNotExist):
			return nil, fmt.Errorf("%w: %w", os.ErrNotExist, err)
		}
		return nil, err
	}
	if err := ctx.Err(); err != nil {
		os.RemoveAll(dir)
		return nil, err
	}
	if err := checkDB(ctx, out); err != nil {
		os.RemoveAll(dir)
		if cerr := ctx.Err(); cerr != nil {
			return nil, cerr
		}
		return nil, fmt.Errorf("%w: %v", ErrCheckFailed, err)
	}
	// A stop that lands as the check finishes still withdraws the copy.
	if err := ctx.Err(); err != nil {
		os.RemoveAll(dir)
		return nil, err
	}
	txid, at := rec.newest()
	return &Restored{Path: out, Dir: dir, TXID: txid, At: at}, nil
}

func restoreOnce(ctx context.Context, client litestream.ReplicaClient, out string, target uint64, at time.Time) error {
	opt := litestream.NewRestoreOptions()
	opt.OutputPath = out
	// A target transaction (RestoreToTarget); zero is none.
	opt.TXID = ltx.TXID(target)
	// S3 lists LastModified to the whole second, rounded down: a file
	// uploaded at 14:00:00.7 lists as 14:00:00, before a T of 14:00:00.5,
	// though it holds writes after T. Rounding T down to its second keeps
	// it out: a listed time before a whole second means an upload before it.
	// A zero at stays zero: no time bound.
	opt.Timestamp = at.Truncate(time.Second)
	// Follow stays false (spec §12 condition ii); IntegrityCheck is left at
	// litestream's default and CheckDB runs after (condition iii).
	return litestream.NewReplicaWithClient(nil, client).Restore(ctx, opt)
}

func removeCopy(out string) {
	for _, p := range []string{out, out + ".tmp", out + "-wal", out + "-shm"} {
		_ = os.Remove(p)
	}
}

// recorder sees every file a Restore opens, for two reasons: a file that
// answered os.ErrNotExist is what makes the restore worth a second plan,
// whatever litestream wrapped it in on the way out; and the header of the
// newest file opened is the copy's "as of" (R1).
type recorder struct {
	*Client
	dir         string // the private directory each file is downloaded into
	mu          sync.Mutex
	vanished    bool
	checkFailed bool
	at          time.Time // the as-of time; zero for the newest copy
	oldest      time.Time // earliest time of a file a restore can start from
	maxTXID     ltx.TXID
	header      *bytes.Buffer // first ltx.HeaderSize bytes of the maxTXID file
}

// LTXFiles passes the listing through, and for an as-of restore checks it
// first: every file must carry a time, and the earliest file a restore can
// start from (any snapshot, or a file beginning at the first transaction)
// is remembered so a time before it can be named.
func (r *recorder) LTXFiles(ctx context.Context, level int, seek ltx.TXID, useMetadata bool) (ltx.FileIterator, error) {
	itr, err := r.Client.LTXFiles(ctx, level, seek, useMetadata)
	if err != nil || r.at.IsZero() {
		return itr, err
	}
	var infos []*ltx.FileInfo
	for itr.Next() {
		infos = append(infos, itr.Item())
	}
	err = itr.Err()
	if cerr := itr.Close(); err == nil {
		err = cerr
	}
	if err != nil {
		return nil, err
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	for _, f := range infos {
		if f.CreatedAt.IsZero() {
			return nil, ErrNoTimes
		}
		if level == litestream.SnapshotLevel || f.MinTXID == 1 {
			if r.oldest.IsZero() || f.CreatedAt.Before(r.oldest) {
				r.oldest = f.CreatedAt
			}
		}
	}
	return ltx.NewFileInfoSliceIterator(infos), nil
}

func (r *recorder) oldestStart() time.Time {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.oldest
}

func (r *recorder) OpenLTXFile(ctx context.Context, level int, minTXID, maxTXID ltx.TXID, offset, size int64) (io.ReadCloser, error) {
	// The first open goes through download's loop too, so a full door's
	// 429 waits out its Retry-After here rather than failing back to
	// litestream, whose own quick retries ignore it.
	rc, err := r.download(ctx, nil, level, minTXID, maxTXID, offset, size)
	if err != nil {
		return nil, err
	}
	if offset != 0 {
		return rc, nil
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if maxTXID < r.maxTXID {
		return rc, nil
	}
	r.maxTXID = maxTXID
	r.header = &bytes.Buffer{}
	return &headerTee{ReadCloser: rc, buf: r.header}, nil
}

// dbDownloadMaxStalls is how many reopens in a row may deliver no byte
// before a download gives up. Progress resets the count, so a slow but
// moving link always finishes.
const dbDownloadMaxStalls = 5

// download reads one file whole into the private directory and closes the
// stream before the restore reads a byte of it. litestream's compactor
// holds every planned file open until the restore ends, so streaming
// straight through would hold one door stream per planned file, and the
// door bounds a person's concurrent streams (apidb.go). Downloaded, a
// restore holds one stream at a time.
//
// A connection that drops part way is resumed here, from the bytes already
// kept, rather than handed to litestream: its resumable reader would start
// the file again from the requested offset, and a large file over a flaky
// link would never finish. The file goes when the restore closes it; on
// failure, with the whole directory.
func (r *recorder) download(ctx context.Context, rc io.ReadCloser, level int, minTXID, maxTXID ltx.TXID, offset, size int64) (io.ReadCloser, error) {
	f, err := os.CreateTemp(r.dir, "ltx-*")
	if err != nil {
		if rc != nil {
			rc.Close()
		}
		return nil, err
	}
	fail := func(err error) (io.ReadCloser, error) {
		f.Close()
		os.Remove(f.Name())
		return nil, err
	}
	var written int64
	stalls, wait := 0, dbDownloadBackoff
	for {
		if rc != nil {
			n, cerr := io.Copy(f, rc)
			rc.Close()
			rc = nil
			written += n
			if cerr == nil {
				break
			}
			if n > 0 {
				stalls, wait = 0, dbDownloadBackoff
			} else if stalls++; stalls >= dbDownloadMaxStalls {
				return fail(cerr)
			}
		}
		if ctx.Err() != nil {
			return fail(ctx.Err())
		}
		var rest int64 // 0: to the end, as RangeHeader reads it
		if size > 0 {
			if rest = size - written; rest <= 0 {
				break
			}
		}
		next, err := r.Client.OpenLTXFile(ctx, level, minTXID, maxTXID, offset+written, rest)
		if err == nil {
			rc = next
			continue
		}
		switch {
		case errors.Is(err, os.ErrNotExist):
			r.mu.Lock()
			r.vanished = true
			r.mu.Unlock()
			return fail(err)
		case ctx.Err() != nil:
			return fail(ctx.Err())
		case !transientOpen(err):
			return fail(err)
		}
		// A busy door or a dropped connection: wait and ask again for the
		// same remaining range, keeping every byte already here. It counts
		// as an attempt that delivered nothing.
		if stalls++; stalls >= dbDownloadMaxStalls {
			return fail(err)
		}
		d := wait
		var door *DoorError
		if errors.As(err, &door) && door.RetryAfter > 0 {
			d = door.RetryAfter
		}
		wait = min(wait*2, dbDownloadBackoffMax)
		select {
		case <-ctx.Done():
			return fail(ctx.Err())
		case <-time.After(d):
		}
	}
	if _, err := f.Seek(0, io.SeekStart); err != nil {
		return fail(err)
	}
	if err := r.verifyPinnedSnapshot(ctx, f, level, minTXID, maxTXID, offset, written); err != nil {
		r.mu.Lock()
		r.checkFailed = true
		r.mu.Unlock()
		return fail(err)
	}
	return &downloaded{ctx: ctx, File: f}, nil
}

func (r *recorder) verifyPinnedSnapshot(ctx context.Context, f *os.File, level int, lo, hi ltx.TXID, offset, written int64) error {
	r.handles.mu.Lock()
	meta := r.handles.files[GetTail(level, uint64(lo), uint64(hi))]
	r.handles.mu.Unlock()
	if meta.SnapshotCRC64 == nil {
		return nil
	}
	if offset != 0 || written != meta.Size {
		return ErrCheckFailed
	}
	dec := ltx.NewDecoder(&downloaded{ctx: ctx, File: f})
	if err := dec.Verify(); err != nil || uint64(dec.Trailer().FileChecksum) != *meta.SnapshotCRC64 {
		return ErrCheckFailed
	}
	_, err := f.Seek(0, io.SeekStart)
	return err
}

// downloaded is a file read once and removed on Close.
type downloaded struct {
	ctx context.Context
	*os.File
}

// Read checks the restore's context first. Once every file is downloaded,
// litestream's compaction reads only these files and checks no context of
// its own, so this is where Ctrl-C reaches it.
func (d *downloaded) Read(p []byte) (int, error) {
	if err := d.ctx.Err(); err != nil {
		return 0, err
	}
	return d.File.Read(p)
}

func (d *downloaded) Close() error {
	err := d.File.Close()
	if rerr := os.Remove(d.File.Name()); err == nil {
		err = rerr
	}
	return err
}

func (r *recorder) sawVanish() bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.vanished
}

// newest returns the copy's newest transaction and the time its LTX header
// records. A header that cannot be read leaves At zero rather than guessed.
func (r *recorder) newest() (uint64, time.Time) {
	r.mu.Lock()
	defer r.mu.Unlock()
	var h ltx.Header
	if r.header == nil || r.header.Len() < ltx.HeaderSize || h.UnmarshalBinary(r.header.Bytes()[:ltx.HeaderSize]) != nil {
		return uint64(r.maxTXID), time.Time{}
	}
	return uint64(r.maxTXID), time.UnixMilli(h.Timestamp).UTC()
}

// headerTee copies the first ltx.HeaderSize bytes read through it.
type headerTee struct {
	io.ReadCloser
	buf *bytes.Buffer
}

func (t *headerTee) Read(p []byte) (int, error) {
	n, err := t.ReadCloser.Read(p)
	if need := ltx.HeaderSize - t.buf.Len(); need > 0 && n > 0 {
		t.buf.Write(p[:min(n, need)])
	}
	return n, err
}

// dbDownloadBackoff is the first wait before retrying a reopen that met a
// busy door or a dropped connection; each further wait doubles, up to
// dbDownloadBackoffMax. A 429's Retry-After overrides it.
var dbDownloadBackoff = 250 * time.Millisecond

const dbDownloadBackoffMax = 8 * time.Second

// transientOpen reports a reopen failure worth waiting out: the door busy
// or briefly unavailable (429, 502, 503, 504), or no answer at all (a
// transport error). Any other refusal from the door is final.
func transientOpen(err error) bool {
	var door *DoorError
	if errors.As(err, &door) {
		switch door.Status {
		case http.StatusTooManyRequests, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout:
			return true
		}
		return false
	}
	return true
}