package jetstream

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/gorilla/websocket"
	"github.com/klauspost/compress/zstd"
)

type fakeCursorDB struct {
	saved    int64
	savedErr error

	mu  sync.Mutex
	got int64
}

func (f *fakeCursorDB) GetLastTimeUs() (int64, error) { return f.saved, f.savedErr }
func (f *fakeCursorDB) SaveLastTimeUs(us int64) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.got = us
	return nil
}

func cursorFor(db DB) int64 {
	j := &JetstreamClient{db: db, wantedDids: make(Set[string])}
	return *j.getLastTimeUs(context.Background())
}

func newGapClient(db DB, gaps *int) *JetstreamClient {
	j := &JetstreamClient{db: db, l: slog.Default(), wantedDids: make(Set[string])}
	j.SetGapHandler(func(context.Context) { *gaps++ })
	return j
}

const twoDaysUs = int64(2 * 24 * 60 * 60 * 1000 * 1000)

func TestStaleCursorIsHonored(t *testing.T) {
	old := time.Now().UnixMicro() - 5*twoDaysUs
	if got := cursorFor(&fakeCursorDB{saved: old}); got != old {
		t.Fatalf("stale cursor must be honored, not snapped forward: got %d, want %d", got, old)
	}
}

func TestZeroCursorIsHonored(t *testing.T) {
	if got := cursorFor(&fakeCursorDB{saved: 0}); got != 0 {
		t.Fatalf("cursor 0 must replay from the start: got %d, want 0", got)
	}
}

func TestMissingCursorConnectsLive(t *testing.T) {
	var gaps int
	j := newGapClient(&fakeCursorDB{savedErr: errors.New("no row")}, &gaps)
	if got := j.resume(context.Background()); got != nil {
		t.Fatalf("missing cursor must connect live, got %d", *got)
	}
	_ = j.withDidFilter(func(context.Context, *Event) error { return nil })(context.Background(), &Event{Did: "did:plc:a", TimeUS: 5})
	if gaps != 0 {
		t.Fatal("a first boot has nothing to catch up on")
	}
}

func TestResumeRewindsASecondFromTheLastEvent(t *testing.T) {
	var gaps int
	j := newGapClient(&fakeCursorDB{saved: 10_000_001}, &gaps)
	if got := j.resume(context.Background()); got == nil || *got != 9_000_001 {
		t.Fatalf("cursor = %v, want 9000001", got)
	}
}

func TestResumeKeepsTheStoredCursorWithoutAGapHandler(t *testing.T) {
	j := &JetstreamClient{db: &fakeCursorDB{saved: 10_000_001}, l: slog.Default(), wantedDids: make(Set[string])}
	if got := j.resume(context.Background()); got == nil || *got != 10_000_001 {
		t.Fatalf("cursor = %v, want the stored one", got)
	}
}

func firstEventAfterResume(t *testing.T, saved, first int64) int {
	t.Helper()
	var gaps int
	j := newGapClient(&fakeCursorDB{saved: saved}, &gaps)
	j.resume(context.Background())
	wrapped := j.withDidFilter(func(context.Context, *Event) error { return nil })
	_ = wrapped(context.Background(), &Event{Did: "did:plc:a", TimeUS: first})
	_ = wrapped(context.Background(), &Event{Did: "did:plc:a", TimeUS: first + 1_000_000})
	return gaps
}

func TestAGapIsReportedWhenTheStreamStartsAfterTheLastEvent(t *testing.T) {
	// last event was 10_000_000, so the stored cursor is 10_000_001
	if got := firstEventAfterResume(t, 10_000_001, 10_000_001); got != 1 {
		t.Fatalf("gap handler ran %d times, want once", got)
	}
}

func TestNoGapWhenTheStreamReplaysTheLastEvent(t *testing.T) {
	for name, first := range map[string]int64{
		"the last event itself": 10_000_000,
		"an earlier event":      9_500_000,
	} {
		if got := firstEventAfterResume(t, 10_000_001, first); got != 0 {
			t.Fatalf("%s: gap handler ran %d times, want none", name, got)
		}
	}
}

func TestLastSeenTracksEveryEventPreFilter(t *testing.T) {
	j := &JetstreamClient{wantedDids: Set[string]{"did:plc:boltless": {}}}
	wrapped := j.withDidFilter(func(context.Context, *Event) error { return nil })

	_ = wrapped(context.Background(), &Event{Did: "did:plc:akshay", TimeUS: 100})
	if got := j.lastSeenUs.Load(); got != 101 {
		t.Fatalf("filtered-out event must still advance last-seen: got %d, want 101", got)
	}

	_ = wrapped(context.Background(), &Event{Did: "did:plc:boltless", TimeUS: 200})
	if got := j.lastSeenUs.Load(); got != 201 {
		t.Fatalf("matching event must advance last-seen: got %d, want 201", got)
	}
}

func TestExemptCollectionShouldBypassDidFilter(t *testing.T) {
	j := &JetstreamClient{
		wantedDids:      Set[string]{"did:plc:boltless": {}},
		unfilteredNsids: make(Set[string]),
	}
	j.ExemptCollection("sh.tangled.repo")

	var processed int
	wrapped := j.withDidFilter(func(context.Context, *Event) error {
		processed++
		return nil
	})

	_ = wrapped(context.Background(), &Event{
		Did: "did:plc:akshay", TimeUS: 100,
		Commit: &Commit{Collection: "sh.tangled.repo"},
	})
	if processed != 1 {
		t.Fatal("An exempt collection must reach this processor from an unwanted DID; enrollment is decided later, not by this filter.")
	}

	_ = wrapped(context.Background(), &Event{
		Did: "did:plc:akshay", TimeUS: 200,
		Commit: &Commit{Collection: "sh.tangled.repo.pull"},
	})
	if processed != 1 {
		t.Fatal("A non-exempt collection must still be filtered by DID.")
	}
}

func TestLastSeenAdvancesAfterProcessing(t *testing.T) {
	j := &JetstreamClient{wantedDids: make(Set[string])}

	var seenDuringProcess int64
	wrapped := j.withDidFilter(func(context.Context, *Event) error {
		seenDuringProcess = j.lastSeenUs.Load()
		return nil
	})

	_ = wrapped(context.Background(), &Event{Did: "did:plc:boltless", TimeUS: 500})

	if seenDuringProcess == 501 {
		t.Fatal("cursor advanced before the event finished processing; a crash mid-process would skip it")
	}
	if got := j.lastSeenUs.Load(); got != 501 {
		t.Fatalf("cursor must advance once processing returns: got %d, want 501", got)
	}
}

func TestShutdownLeavesTheCursorBeforeTheCutEvent(t *testing.T) {
	j := &JetstreamClient{wantedDids: make(Set[string])}
	ctx, cancel := context.WithCancel(context.Background())
	wrapped := j.withDidFilter(func(ctx context.Context, _ *Event) error {
		cancel()
		return ctx.Err()
	})

	_ = wrapped(ctx, &Event{Did: "did:plc:boltless", TimeUS: 700})
	if got := j.lastSeenUs.Load(); got != 0 {
		t.Fatalf("cursor moved past an event shutdown cut short: got %d, want it untouched", got)
	}
}

func TestAFailingEventStillAdvancesTheCursor(t *testing.T) {
	j := &JetstreamClient{wantedDids: make(Set[string])}
	wrapped := j.withDidFilter(func(context.Context, *Event) error {
		return errors.New("bad record")
	})

	_ = wrapped(context.Background(), &Event{Did: "did:plc:boltless", TimeUS: 800})
	if got := j.lastSeenUs.Load(); got != 801 {
		t.Fatalf("a record that fails on its own must not stall the cursor: got %d, want 801", got)
	}
}

// nothing listens on port 1, so every connect fails at once
const refusedEndpoint = "ws://127.0.0.1:1/subscribe"

func startClient(t *testing.T, endpoint string, db DB, waitForDid bool) (*JetstreamClient, context.CancelFunc) {
	t.Helper()
	j, err := NewJetstreamClient(endpoint, nil, nil, slog.Default(), db, waitForDid, false)
	if err != nil {
		t.Fatal(err)
	}
	j.lastSeenUs.Store(42)
	ctx, cancel := context.WithCancel(context.Background())
	if err := j.StartJetstream(ctx, func(context.Context, *Event) error { return nil }); err != nil {
		t.Fatal(err)
	}
	return j, cancel
}

func requireStopsWithSavedCursor(t *testing.T, j *JetstreamClient, db *fakeCursorDB, msg string) {
	t.Helper()
	select {
	case <-j.Done():
	case <-time.After(5 * time.Second):
		t.Fatal(msg)
	}
	db.mu.Lock()
	defer db.mu.Unlock()
	if db.got != 42 {
		t.Fatalf("cursor saved at stop = %d, want 42", db.got)
	}
}

func TestDoneClosesOnlyAfterTheFinalCursorSave(t *testing.T) {
	db := &fakeCursorDB{savedErr: errors.New("no row")}
	j, cancel := startClient(t, refusedEndpoint, db, false)
	// long enough to be waiting out a failed connect when the cancel lands
	time.Sleep(50 * time.Millisecond)
	cancel()
	requireStopsWithSavedCursor(t, j, db, "the reconnect loop kept going after its context was cancelled")
}

func TestDoneClosesWhenCancelledBeforeAnyDidArrives(t *testing.T) {
	j, cancel := startClient(t, refusedEndpoint, &fakeCursorDB{}, true)
	cancel()
	select {
	case <-j.Done():
	case <-time.After(5 * time.Second):
		t.Fatal("a client still waiting for its first did ignored shutdown")
	}
}

func TestDoneClosesEvenWhenTheStreamIsQuiet(t *testing.T) {
	connected := make(chan struct{})
	release := make(chan struct{})
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer conn.Close()
		close(connected)
		<-release
	}))
	defer srv.Close()
	defer close(release)

	db := &fakeCursorDB{savedErr: errors.New("no row")}
	j, cancel := startClient(t, "ws"+strings.TrimPrefix(srv.URL, "http")+"/subscribe", db, false)
	<-connected
	cancel()
	requireStopsWithSavedCursor(t, j, db, "a stream with no traffic held shutdown hostage, so the final cursor save never ran")
}

func frameServer(t *testing.T, frames ...string) *streamServer {
	t.Helper()
	enc, err := zstd.NewWriter(nil, zstd.WithEncoderDict(zstdDictionary))
	if err != nil {
		t.Fatal(err)
	}
	var srv *streamServer
	srv = newStreamServer(t, func(n int, con *websocket.Conn) {
		srv.mu.Lock()
		compressed := srv.headers[n].Get("Socket-Encoding") == "zstd"
		srv.mu.Unlock()
		for _, f := range frames {
			if compressed {
				con.WriteMessage(websocket.BinaryMessage, enc.EncodeAll([]byte(f), nil))
			} else {
				con.WriteMessage(websocket.TextMessage, []byte(f))
			}
		}
		drain(con)
	})
	return srv
}

func commitFrame(did string, us int, collection string, live bool) string {
	return fmt.Sprintf(`{"did":%q,"time_us":%d,"kind":"commit","commit":{"operation":"create","collection":%q,"rkey":"k","live":%t}}`, did, us, collection, live)
}

func runClient(t *testing.T, j *JetstreamClient, stop func(*Event) bool) {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := j.StartJetstream(ctx, func(_ context.Context, e *Event) error {
		if stop(e) {
			cancel()
		}
		return nil
	}); err != nil {
		t.Fatal(err)
	}
	<-j.Done()
	if errors.Is(ctx.Err(), context.DeadlineExceeded) {
		t.Fatal("the client never delivered what the test waited for")
	}
}

func newTestClient(t *testing.T, srv *streamServer, db DB, collections ...string) *JetstreamClient {
	t.Helper()
	j, err := NewJetstreamClient(srv.endpoint(), collections, nil, slog.Default(), db, false, false)
	if err != nil {
		t.Fatal(err)
	}
	return j
}

func TestClientDeliversInOrderAndOneAtATime(t *testing.T) {
	var frames []string
	for i := 1; i <= 50; i++ {
		frames = append(frames, commitFrame("did:plc:a", i, "c", true))
	}
	j := newTestClient(t, frameServer(t, frames...), &fakeCursorDB{savedErr: errors.New("no row")})

	var inflight atomic.Int32
	var got []int64
	runClient(t, j, func(e *Event) bool {
		if inflight.Add(1) != 1 {
			t.Error("two events were being handled at once")
		}
		time.Sleep(time.Millisecond)
		inflight.Add(-1)
		got = append(got, e.TimeUS)
		return e.TimeUS == 50
	})
	for i, us := range got {
		if us != int64(i+1) {
			t.Fatalf("event %d arrived as time_us %d, want %d", i, us, i+1)
		}
	}
}

func TestClientDidFilterChangesAtRuntime(t *testing.T) {
	srv := frameServer(t,
		commitFrame("did:plc:a", 1, "c", true),
		commitFrame("did:plc:b", 2, "c", true),
		commitFrame("did:plc:a", 3, "c", true),
		commitFrame("did:plc:b", 4, "c", true),
		commitFrame("did:plc:a", 5, "c", true),
	)
	j := newTestClient(t, srv, &fakeCursorDB{savedErr: errors.New("no row")})
	j.AddDid("did:plc:b")

	var got []int64
	runClient(t, j, func(e *Event) bool {
		got = append(got, e.TimeUS)
		if e.TimeUS == 2 {
			j.AddDid("did:plc:a")
			j.RemoveDid("did:plc:b")
		}
		return e.TimeUS == 5
	})
	if want := []int64{2, 3, 5}; !slices.Equal(got, want) {
		t.Fatalf("delivered %v, want %v", got, want)
	}
}

func TestClientUnfilteredCollectionIgnoresTheDidFilter(t *testing.T) {
	srv := frameServer(t,
		commitFrame("did:plc:a", 1, "other", true),
		commitFrame("did:plc:a", 2, "open", true),
		commitFrame("did:plc:b", 3, "other", true),
	)
	j := newTestClient(t, srv, &fakeCursorDB{savedErr: errors.New("no row")})
	j.AddDid("did:plc:b")
	j.ExemptCollection("open")

	var got []int64
	runClient(t, j, func(e *Event) bool {
		got = append(got, e.TimeUS)
		return e.TimeUS == 3
	})
	if len(got) != 2 || got[0] != 2 || got[1] != 3 {
		t.Fatalf("delivered %v, want the exempt event 2 and the wanted did's event 3", got)
	}
}

// only deliberi sends wantedEventTypes, a plain jetstream client must not
func TestClientAsksForNoEventTypes(t *testing.T) {
	srv := frameServer(t, commitFrame("did:plc:a", 1, "c", true))
	j := newTestClient(t, srv, &fakeCursorDB{savedErr: errors.New("no row")}, "c")

	runClient(t, j, func(e *Event) bool { return true })
	if q := srv.query(0); q.Has("wantedEventTypes") {
		t.Fatalf("query = %v, want no wantedEventTypes", q)
	}
}

func TestClientSavesTheCursorOnTheTimer(t *testing.T) {
	srv := frameServer(t, commitFrame("did:plc:a", 10, "c", true))
	db := &fakeCursorDB{savedErr: errors.New("no row")}
	j := newTestClient(t, srv, db)
	j.saveEvery = 10 * time.Millisecond

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	if err := j.StartJetstream(ctx, func(context.Context, *Event) error { return nil }); err != nil {
		t.Fatal(err)
	}
	deadline := time.Now().Add(5 * time.Second)
	for {
		db.mu.Lock()
		got := db.got
		db.mu.Unlock()
		if got == 11 {
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("saved cursor = %d, want 11 from the timer before any shutdown", got)
		}
		time.Sleep(5 * time.Millisecond)
	}
}

func TestClientReportsAGapAndResumesARewoundCursor(t *testing.T) {
	// cursor is one past 10_000_000, but the server's history starts after it
	srv := frameServer(t, commitFrame("did:plc:a", 20_000_000, "c", true))
	j := newTestClient(t, srv, &fakeCursorDB{saved: 10_000_001})
	gaps := make(chan struct{}, 2)
	j.SetGapHandler(func(context.Context) { gaps <- struct{}{} })

	runClient(t, j, func(e *Event) bool { return true })
	if len(gaps) != 1 {
		t.Fatalf("gap handler ran %d times, want once", len(gaps))
	}
	if got := srv.query(0).Get("cursor"); got != "9000001" {
		t.Fatalf("connected with cursor %q, want the stored one rewound a second", got)
	}
}

func TestClientFailingEventsDontDropTheConnection(t *testing.T) {
	srv := frameServer(t,
		commitFrame("did:plc:a", 1, "c", true),
		commitFrame("did:plc:a", 2, "c", true),
	)
	j := newTestClient(t, srv, &fakeCursorDB{savedErr: errors.New("no row")})

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	var seen atomic.Int32
	_ = j.StartJetstream(ctx, func(_ context.Context, e *Event) error {
		if seen.Add(1) == 2 {
			cancel()
		}
		return errors.New("bad record")
	})
	<-j.Done()
	if seen.Load() != 2 || srv.connects() != 1 {
		t.Fatalf("saw %d events over %d connections, want 2 over 1", seen.Load(), srv.connects())
	}
}