package lease

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	s3 "github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/smithy-go"
	"github.com/benbjohnson/litestream"
	lsfile "github.com/benbjohnson/litestream/file"
	lss3 "github.com/benbjohnson/litestream/s3"
	"github.com/superfly/ltx"
)

type lockStore struct {
	mu            sync.Mutex
	body          []byte
	etag          string
	puts, deletes int
	beforePut     func(context.Context, *s3.PutObjectInput) error
}

func (s *lockStore) GetObject(ctx context.Context, p *s3.GetObjectInput, _ ...func(*s3.Options)) (*s3.GetObjectOutput, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.body == nil {
		return nil, &smithy.GenericAPIError{Code: "NoSuchKey"}
	}
	return &s3.GetObjectOutput{Body: io.NopCloser(bytes.NewReader(bytes.Clone(s.body))), ETag: aws.String(s.etag)}, nil
}
func (s *lockStore) PutObject(ctx context.Context, p *s3.PutObjectInput, _ ...func(*s3.Options)) (*s3.PutObjectOutput, error) {
	raw, e := io.ReadAll(p.Body)
	if e != nil {
		return nil, e
	}
	if s.beforePut != nil {
		if e := s.beforePut(ctx, p); e != nil {
			return nil, e
		}
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if (aws.ToString(p.IfNoneMatch) == "*" && s.body != nil) || (p.IfMatch != nil && aws.ToString(p.IfMatch) != s.etag) {
		return nil, &smithy.GenericAPIError{Code: "PreconditionFailed"}
	}
	s.puts++
	s.etag = fmt.Sprintf("\"etag-%d\"", s.puts)
	s.body = raw
	return &s3.PutObjectOutput{ETag: aws.String(s.etag)}, nil
}
func (s *lockStore) DeleteObject(ctx context.Context, p *s3.DeleteObjectInput, _ ...func(*s3.Options)) (*s3.DeleteObjectOutput, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if p.IfMatch == nil || aws.ToString(p.IfMatch) != s.etag {
		return nil, &smithy.GenericAPIError{Code: "PreconditionFailed"}
	}
	s.deletes++
	s.body = nil
	s.etag = ""
	return &s3.DeleteObjectOutput{}, nil
}
func (s *lockStore) counts() (int, int) { s.mu.Lock(); defer s.mu.Unlock(); return s.puts, s.deletes }
func (s *lockStore) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	var err error
	switch r.Method {
	case "GET":
		var out *s3.GetObjectOutput
		out, err = s.GetObject(r.Context(), &s3.GetObjectInput{})
		if err == nil {
			w.Header().Set("ETag", aws.ToString(out.ETag))
			io.Copy(w, out.Body)
			return
		}
	case "PUT":
		p := &s3.PutObjectInput{Body: r.Body}
		if v := r.Header.Get("If-Match"); v != "" {
			p.IfMatch = aws.String(v)
		}
		if v := r.Header.Get("If-None-Match"); v != "" {
			p.IfNoneMatch = aws.String(v)
		}
		var out *s3.PutObjectOutput
		out, err = s.PutObject(r.Context(), p)
		if err == nil {
			w.Header().Set("ETag", aws.ToString(out.ETag))
			return
		}
	case "DELETE":
		_, err = s.DeleteObject(r.Context(), &s3.DeleteObjectInput{IfMatch: aws.String(r.Header.Get("If-Match"))})
		if err == nil {
			w.WriteHeader(204)
			return
		}
	}
	var api *smithy.GenericAPIError
	if errors.As(err, &api) {
		code := 412
		if api.Code == "NoSuchKey" {
			code = 404
		}
		w.Header().Set("Content-Type", "application/xml")
		w.WriteHeader(code)
		fmt.Fprintf(w, "<Error><Code>%s</Code></Error>", api.Code)
		return
	}
	w.WriteHeader(500)
}
func localHeld(t *testing.T, backend *lockStore, ttl time.Duration) *Held {
	t.Helper()
	srv := httptest.NewServer(backend)
	t.Cleanup(srv.Close)
	fakeCreds(t)
	t.Setenv("CARLOS_S3_ENDPOINT", srv.URL)
	h, e := AcquireTTL(context.Background(), "test-bucket", "one.example", "test-region", ttl)
	if e != nil {
		t.Fatal(e)
	}
	t.Cleanup(func() { _ = h.Release(context.Background()) })
	return h
}
func await(t *testing.T, ch <-chan struct{}) {
	t.Helper()
	select {
	case <-ch:
	case <-time.After(5 * time.Second):
		t.Fatal("timed out")
	}
}
func restoreCount(t *testing.T, replica, out string) int {
	t.Helper()
	r := litestream.NewReplicaWithClient(nil, lsfile.NewReplicaClient(replica))
	opt := litestream.NewRestoreOptions()
	opt.OutputPath = out
	if e := r.Restore(context.Background(), opt); e != nil {
		t.Fatal(e)
	}
	d, e := sql.Open("sqlite", "file:"+out+"?mode=ro")
	if e != nil {
		t.Fatal(e)
	}
	defer d.Close()
	var n int
	var integrity string
	if e = d.QueryRow("PRAGMA integrity_check").Scan(&integrity); e != nil || integrity != "ok" {
		t.Fatalf("integrity %q %v", integrity, e)
	}
	if e = d.QueryRow("SELECT count(*) FROM t").Scan(&n); e != nil {
		t.Fatal(e)
	}
	var contents string
	if e = d.QueryRow("SELECT group_concat(body, '|') FROM (SELECT body FROM t ORDER BY id)").Scan(&contents); e != nil {
		t.Fatal(e)
	}
	if n == 5 && contents != "post-vacuum|post-vacuum|post-vacuum|post-vacuum|post-vacuum" {
		t.Fatalf("restored content receipt differs: %q", contents)
	}
	return n
}
func TestHeldCheckpointVacuumRestoreAndTransfer(t *testing.T) {
	backend := &lockStore{}
	h := localHeld(t, backend, 10*time.Second)
	path := filepath.Join(t.TempDir(), "source.db")
	seedDB(t, path)
	replica := t.TempDir()
	h.replicaClient = lsfile.NewReplicaClient(replica)
	rn := h.rn
	receipt, e := h.Checkpoint(context.Background(), path)
	if e != nil {
		t.Fatal(e)
	}
	if receipt.CompletedAt.IsZero() || !receipt.CompletedAt.Before(receipt.SafeUntil) || h.rn != rn {
		t.Fatal("invalid checkpoint receipt/ownership")
	}
	if n := restoreCount(t, replica, filepath.Join(t.TempDir(), "old.db")); n != 50 {
		t.Fatal(n)
	}
	d, e := sql.Open("sqlite", path)
	if e != nil {
		t.Fatal(e)
	}
	_, e = d.Exec("DELETE FROM t WHERE id>5; UPDATE t SET body='post-vacuum'; PRAGMA wal_checkpoint(TRUNCATE); VACUUM; PRAGMA wal_checkpoint(TRUNCATE)")
	if e != nil {
		t.Fatal(e)
	}
	if e = d.Close(); e != nil {
		t.Fatal(e)
	}
	if _, e = h.Checkpoint(context.Background(), path); e != nil {
		t.Fatal(e)
	}
	if n := restoreCount(t, replica, filepath.Join(t.TempDir(), "new.db")); n != 5 {
		t.Fatalf("stale replica: %d", n)
	}
	if _, deletes := backend.counts(); deletes != 0 {
		t.Fatal("checkpoint released lock")
	}
	// The actual leaser sees the same still-owned lock after both checkpoints.
	if peer, e := AcquireTTL(context.Background(), "test-bucket", "one.example", "test-region", 10*time.Second); e == nil {
		peer.Release(context.Background())
		t.Fatal("acquisition gap")
	}
	r, e := h.Replicate(context.Background(), path, 0)
	if e != nil {
		t.Fatal(e)
	}
	defer r.Close(context.Background())
	if r.Lost() != h.Lost() {
		t.Fatal("transfer replaced loss channel")
	}
	select {
	case <-r.Lost():
		t.Fatal("healthy transfer signaled loss")
	default:
	}
	if r.rn != rn {
		t.Fatal("renewer replaced")
	}
	if e = h.Release(context.Background()); !errors.Is(e, ErrSpent) {
		t.Fatalf("spent release %v", e)
	}
	if _, e = h.SafeUntil(); !errors.Is(e, ErrSpent) {
		t.Fatal("spent deadline accepted")
	}
	if _, e = r.SafeUntil(); e != nil {
		t.Fatal(e)
	}
	if _, deletes := backend.counts(); deletes != 0 {
		t.Fatal("spent Held released successor")
	}
	if e = r.Close(context.Background()); e != nil {
		t.Fatal(e)
	}
	if _, e = r.SafeUntil(); e == nil {
		t.Fatal("stopped deadline accepted")
	}
	select {
	case <-r.Lost():
		t.Fatal("normal Close signaled loss")
	default:
	}
}

type blockedReplica struct {
	litestream.ReplicaClient
	entered chan struct{}
	release chan struct{}
	once    sync.Once
	fail    error
}

func (b *blockedReplica) WriteLTXFile(ctx context.Context, level int, min, max ltx.TXID, r io.Reader) (*ltx.FileInfo, error) {
	b.once.Do(func() { close(b.entered) })
	select {
	case <-b.release:
	case <-ctx.Done():
		return nil, ctx.Err()
	}
	if b.fail != nil {
		return nil, b.fail
	}
	return b.ReplicaClient.WriteLTXFile(ctx, level, min, max, r)
}
func TestCheckpointExclusionAndPoison(t *testing.T) {
	for _, fail := range []bool{false, true} {
		t.Run(fmt.Sprint(fail), func(t *testing.T) {
			backend := &lockStore{}
			h := localHeld(t, backend, 10*time.Second)
			path := filepath.Join(t.TempDir(), "source.db")
			seedDB(t, path)
			b := &blockedReplica{ReplicaClient: lsfile.NewReplicaClient(t.TempDir()), entered: make(chan struct{}), release: make(chan struct{})}
			if fail {
				b.fail = errors.New("synthetic replica write refusal")
			}
			h.replicaClient = b
			opCtx, cancel := context.WithTimeout(context.Background(), time.Second)
			defer cancel()
			done := make(chan error, 1)
			go func() { _, e := h.Checkpoint(opCtx, path); done <- e }()
			await(t, b.entered)
			if _, e := h.Checkpoint(context.Background(), path); !errors.Is(e, ErrBusy) {
				t.Fatal(e)
			}
			if _, e := h.Replicate(context.Background(), path, 0); !errors.Is(e, ErrBusy) {
				t.Fatal(e)
			}
			if e := h.Release(context.Background()); !errors.Is(e, ErrBusy) {
				t.Fatal(e)
			}
			close(b.release)
			e := <-done
			if fail {
				if e == nil {
					t.Fatal("failed sync accepted")
				}
				if _, e = h.Checkpoint(context.Background(), path); e == nil {
					t.Fatal("poisoned checkpoint")
				}
				if _, e = h.Replicate(context.Background(), path, 0); e == nil {
					t.Fatal("poisoned replicate")
				}
				if _, e = h.SafeUntil(); e == nil {
					t.Fatal("poisoned command deadline")
				}
			} else if e != nil {
				t.Fatal(e)
			}
			if e = h.rn.Err(); e != nil {
				t.Fatalf("ownership released on checkpoint outcome: %v", e)
			}
			if _, deletes := backend.counts(); deletes != 0 {
				t.Fatal("unexpected release")
			}
			if e = h.Release(context.Background()); e != nil {
				t.Fatal(e)
			}
		})
	}
}

func TestBlockedRenewalDeadlinePrecedesTakeover(t *testing.T) {
	entered := make(chan struct{})
	var once sync.Once
	backend := &lockStore{beforePut: func(ctx context.Context, p *s3.PutObjectInput) error {
		if p.IfMatch != nil {
			once.Do(func() { close(entered) })
			<-ctx.Done()
			return ctx.Err()
		}
		return nil
	}}
	// marginForTTL gives short test TTLs a margin of ttl/10. At 900ms that
	// was 90ms, less than a loaded CI runner's scheduling delay (#268); 3s
	// gives 300ms and checks the same property. Production's margin is 5s.
	ttl := 3 * time.Second
	h := localHeld(t, backend, ttl)
	until, e := h.SafeUntil()
	if e != nil {
		t.Fatal(e)
	}
	await(t, entered)
	// A worker observes cancellation without waiting for the blocked remote call.
	await(t, h.Lost())
	if e = h.Err(); e == nil {
		t.Fatal("blocked renewal remained healthy")
	}
	if time.Since(until) > ttl/10 {
		t.Fatal("loss notification consumed termination margin")
	}
	backend.mu.Lock()
	var original litestream.Lease
	json.Unmarshal(backend.body, &original)
	backend.mu.Unlock()
	if !time.Now().Before(original.ExpiresAt) {
		t.Fatal("termination did not precede lease expiry")
	}
	if peer, e := AcquireTTL(context.Background(), "test-bucket", "one.example", "test-region", ttl); e == nil {
		peer.Release(context.Background())
		t.Fatal("peer took live lock")
	}
	time.Sleep(time.Until(original.ExpiresAt) + 20*time.Millisecond)
	// A separate real leaser can take over the expired CAS object. Its acquisition
	// uses If-Match, so disable the synthetic blocked-renewal hook first.
	backend.beforePut = nil
	peer, e := AcquireTTL(context.Background(), "test-bucket", "one.example", "test-region", ttl)
	if e != nil {
		t.Fatal(e)
	}
	defer peer.Release(context.Background())
	if e = h.Release(context.Background()); e == nil {
		t.Fatal("stale release erased peer")
	}
	if e = peer.Err(); e != nil {
		t.Fatal(e)
	}
}

func TestLateRenewalCannotResurrect(t *testing.T) {
	gate := make(chan struct{})
	entered := make(chan struct{})
	var once sync.Once
	backend := &lockStore{beforePut: func(ctx context.Context, p *s3.PutObjectInput) error {
		if p.IfMatch != nil {
			once.Do(func() { close(entered) })
			<-gate
		}
		return nil
	}}
	leaser := lss3.NewLeaser()
	leaser.SetClient(backend)
	leaser.TTL = 300 * time.Millisecond
	started := time.Now()
	lock, e := leaser.AcquireLease(context.Background())
	if e != nil {
		t.Fatal(e)
	}
	rn := startRenewer(leaser, lock, started, 30*time.Millisecond)
	await(t, entered)
	await(t, rn.ctx.Done())
	if !errors.Is(rn.Err(), ErrLost) {
		t.Fatal("deadline not latched")
	}
	close(gate)
	await(t, rn.done)
	if !errors.Is(rn.Err(), ErrLost) {
		t.Fatal("late acknowledgement resurrected ownership")
	}
	if _, e = rn.SafeUntil(); e == nil {
		t.Fatal("late acknowledgement extended deadline")
	}
	leaser.ReleaseLease(context.Background(), rn.stop())
}

func TestAcquireDeadlineStartsBeforeRemoteRequest(t *testing.T) {
	entered := make(chan time.Time, 1)
	backend := &lockStore{beforePut: func(ctx context.Context, p *s3.PutObjectInput) error {
		if p.IfMatch == nil {
			entered <- time.Now()
			time.Sleep(120 * time.Millisecond)
		}
		return nil
	}}
	ttl := time.Second
	h := localHeld(t, backend, ttl)
	beforeAck := <-entered
	until, e := h.SafeUntil()
	if e != nil {
		t.Fatal(e)
	}
	if until.After(beforeAck.Add(ttl - marginForTTL(ttl))) {
		t.Fatal("acquisition latency extended safe deadline")
	}
	if until == until.Round(0) {
		t.Fatal("deadline lost monotonic clock")
	}
}

func TestAcquireLateSuccessGrantsNoHandle(t *testing.T) {
	backend := &lockStore{beforePut: func(ctx context.Context, p *s3.PutObjectInput) error {
		if p.IfMatch == nil {
			time.Sleep(350 * time.Millisecond)
		}
		return nil
	}}
	srv := httptest.NewServer(backend)
	defer srv.Close()
	fakeCreds(t)
	t.Setenv("CARLOS_S3_ENDPOINT", srv.URL)
	h, e := AcquireTTL(context.Background(), "test-bucket", "one.example", "test-region", 300*time.Millisecond)
	if e == nil || h != nil {
		if h != nil {
			h.Release(context.Background())
		}
		t.Fatal("late acquisition granted authority")
	}
	if _, deletes := backend.counts(); deletes != 1 {
		t.Fatal("known late acquisition was not conditionally cleaned up")
	}
}

func TestCheckpointLossCancelsBlockedSync(t *testing.T) {
	path := filepath.Join(t.TempDir(), "source.db")
	seedDB(t, path)
	b := &blockedReplica{ReplicaClient: lsfile.NewReplicaClient(t.TempDir()), entered: make(chan struct{}), release: make(chan struct{})}
	entered := make(chan struct{})
	var once sync.Once
	backend := &lockStore{beforePut: func(ctx context.Context, p *s3.PutObjectInput) error {
		if p.IfMatch != nil {
			// Keep renewal healthy until the upload is actually blocked.
			// Setup and filesystem latency must not consume the test lease.
			select {
			case <-b.entered:
			default:
				return nil
			}
			once.Do(func() { close(entered) })
			<-ctx.Done()
			return ctx.Err()
		}
		return nil
	}}
	h := localHeld(t, backend, 600*time.Millisecond)
	h.replicaClient = b
	done := make(chan error, 1)
	go func() { _, e := h.Checkpoint(context.Background(), path); done <- e }()
	await(t, b.entered)
	await(t, entered)
	select {
	case e := <-done:
		if e == nil {
			t.Fatal("lost lease returned receipt")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("loss did not cancel checkpoint")
	}
	if _, deletes := backend.counts(); deletes != 0 {
		t.Fatal("checkpoint released after loss")
	}
	if _, e := h.SafeUntil(); e == nil {
		t.Fatal("failed checkpoint usable")
	}
}

func TestCheckpointOpenFailureLeavesOwnedPoisonedHandle(t *testing.T) {
	backend := &lockStore{}
	h := localHeld(t, backend, 10*time.Second)
	// A directory is never a SQLite database; the real store Open must refuse.
	if receipt, e := h.Checkpoint(context.Background(), t.TempDir()); e == nil || !receipt.CompletedAt.IsZero() {
		t.Fatal("failed open returned receipt")
	}
	if e := h.rn.Err(); e != nil {
		t.Fatal("open failure lost held renewal")
	}
	if _, e := h.Replicate(context.Background(), "unused", 0); e == nil {
		t.Fatal("open failure did not poison")
	}
	if _, deletes := backend.counts(); deletes != 0 {
		t.Fatal("open failure released lease")
	}
}

func TestRenewDeadlineStartsBeforeRemoteRequest(t *testing.T) {
	entered := make(chan time.Time, 1)
	backend := &lockStore{beforePut: func(ctx context.Context, p *s3.PutObjectInput) error {
		if p.IfMatch != nil {
			select {
			case entered <- time.Now():
			default:
			}
			time.Sleep(100 * time.Millisecond)
		}
		return nil
	}}
	leaser := lss3.NewLeaser()
	leaser.SetClient(backend)
	leaser.TTL = time.Second
	started := time.Now()
	lock, e := leaser.AcquireLease(context.Background())
	if e != nil {
		t.Fatal(e)
	}
	rn := startRenewer(leaser, lock, started, 100*time.Millisecond)
	defer func() { leaser.ReleaseLease(context.Background(), rn.stop()) }()
	initial, e := rn.SafeUntil()
	if e != nil {
		t.Fatal(e)
	}
	var began time.Time
	select {
	case began = <-entered:
	case <-time.After(2 * time.Second):
		t.Fatal("renewal never started")
	}
	timeout := time.Now().Add(time.Second)
	for {
		until, e := rn.SafeUntil()
		if e != nil {
			t.Fatal(e)
		}
		if until.After(initial) {
			if until.After(began.Add(900 * time.Millisecond)) {
				t.Fatal("renewal latency extended deadline")
			}
			if until == until.Round(0) {
				t.Fatal("renewed deadline lost monotonic clock")
			}
			break
		}
		if time.Now().After(timeout) {
			t.Fatal("renewal never completed")
		}
		time.Sleep(time.Millisecond)
	}
}

type closeFailureReplica struct {
	litestream.ReplicaClient
	path    string
	writes  int
	failure error
}

func (c *closeFailureReplica) WriteLTXFile(ctx context.Context, level int, min, max ltx.TXID, r io.Reader) (*ltx.FileInfo, error) {
	c.writes++
	if c.writes > 1 {
		return nil, c.failure
	}
	info, e := c.ReplicaClient.WriteLTXFile(ctx, level, min, max, r)
	if e != nil {
		return nil, e
	}
	// Test-only interference after the explicit sync's last upload gives the
	// real Close a fresh tail to flush. Maintenance callers must prove stopped
	// writers; this deliberately hostile fixture isolates Close error handling.
	d, e := sql.Open("sqlite", c.path)
	if e != nil {
		return nil, e
	}
	_, e = d.Exec("INSERT INTO t(body) VALUES('late fixture tail')")
	closeErr := d.Close()
	if e != nil {
		return nil, e
	}
	if closeErr != nil {
		return nil, closeErr
	}
	return info, nil
}
func TestCheckpointChecksCloseAfterSuccessfulSync(t *testing.T) {
	backend := &lockStore{}
	h := localHeld(t, backend, 10*time.Second)
	path := filepath.Join(t.TempDir(), "source.db")
	seedDB(t, path)
	sentinel := errors.New("synthetic final tail upload failure")
	client := &closeFailureReplica{ReplicaClient: lsfile.NewReplicaClient(t.TempDir()), path: path, failure: sentinel}
	h.replicaClient = client
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	receipt, e := h.Checkpoint(ctx, path)
	if !errors.Is(e, sentinel) || !receipt.CompletedAt.IsZero() || client.writes < 2 {
		t.Fatalf("close error not preserved: %v writes=%d receipt=%+v", e, client.writes, receipt)
	}
	if e = h.rn.Err(); e != nil {
		t.Fatalf("close failure stopped ownership: %v", e)
	}
}

func TestRenewalFailureSignalsLostBeforeDeadline(t *testing.T) {
	backend := &lockStore{beforePut: func(ctx context.Context, p *s3.PutObjectInput) error {
		if p.IfMatch != nil {
			return &smithy.GenericAPIError{Code: "PreconditionFailed"}
		}
		return nil
	}}
	h := localHeld(t, backend, 3*time.Second)
	until, e := h.SafeUntil()
	if e != nil {
		t.Fatal(e)
	}
	loss := h.Lost()
	await(t, loss)
	if !time.Now().Before(until) {
		t.Fatal("renewal refusal was signaled only at deadline")
	}
	if e = h.Err(); e == nil {
		t.Fatal("loss notification without failed health")
	}
	if _, e = h.SafeUntil(); e == nil {
		t.Fatal("loss notification left safe deadline usable")
	}
	if h.Lost() != loss {
		t.Fatal("loss channel replaced")
	}
}

func TestHealthyReleaseDoesNotSignalLost(t *testing.T) {
	h := localHeld(t, &lockStore{}, 300*time.Millisecond)
	until, e := h.SafeUntil()
	if e != nil {
		t.Fatal(e)
	}
	loss := h.Lost()
	if e = h.Release(context.Background()); e != nil {
		t.Fatal(e)
	}
	if _, e = h.SafeUntil(); !errors.Is(e, ErrSpent) {
		t.Fatal("released handle remained usable")
	}
	time.Sleep(time.Until(until) + 30*time.Millisecond)
	select {
	case <-loss:
		t.Fatal("normal release masqueraded as loss")
	default:
	}
	if h.Err() == nil {
		t.Fatal("spent handle remained healthy")
	}
	select {
	case <-loss:
		t.Fatal("late health check invented loss after release")
	default:
	}
}