package jetstream

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync"
	"testing"
	"time"

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

type streamServer struct {
	*httptest.Server

	mu      sync.Mutex
	queries []url.Values
	headers []http.Header
}

func newStreamServer(t *testing.T, script func(n int, con *websocket.Conn)) *streamServer {
	t.Helper()
	s := &streamServer{}
	s.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		con, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer con.Close()
		s.mu.Lock()
		s.queries = append(s.queries, r.URL.Query())
		s.headers = append(s.headers, r.Header)
		n := len(s.queries) - 1
		s.mu.Unlock()
		script(n, con)
	}))
	t.Cleanup(s.Close)
	return s
}

func (s *streamServer) endpoint() string {
	return "ws" + strings.TrimPrefix(s.URL, "http") + "/subscribe"
}

func (s *streamServer) query(n int) url.Values {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.queries[n]
}

func (s *streamServer) connects() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.queries)
}

func drain(con *websocket.Conn) {
	for {
		if _, _, err := con.ReadMessage(); err != nil {
			return
		}
	}
}

func testStream(t *testing.T, cfg StreamConfig) *Stream {
	t.Helper()
	if cfg.Cursor == nil {
		cfg.Cursor = func() (int64, bool, error) { return 0, false, nil }
	}
	if cfg.ReconnectDelay == 0 {
		cfg.ReconnectDelay = 10 * time.Millisecond
	}
	s, err := NewStream(cfg, slog.New(slog.NewTextHandler(io.Discard, nil)))
	if err != nil {
		t.Fatal(err)
	}
	return s
}

func run(t *testing.T, s *Stream, until func(*Event) bool) {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	err := make(chan error, 1)
	go func() {
		err <- s.Run(ctx, func(_ context.Context, e *Event) error {
			if until(e) {
				cancel()
			}
			return nil
		})
	}()
	if e := <-err; e != nil {
		t.Fatal(e)
	}
	if errors.Is(ctx.Err(), context.DeadlineExceeded) {
		t.Fatal("the stream never delivered what the test waited for")
	}
}

func TestStreamURLCarriesEachParamOnce(t *testing.T) {
	s := testStream(t, StreamConfig{
		Endpoint:    "ws://hydrant/subscribe?cursor=9&wantedEventTypes=old",
		Collections: []string{"org.tangled.track.ticket", "org.tangled.feed.comment"},
		EventTypes:  []string{"live", "historical"},
	})
	got, err := s.url(new(int64))
	if err != nil {
		t.Fatal(err)
	}
	u, err := url.Parse(got)
	if err != nil {
		t.Fatal(err)
	}
	q := u.Query()
	if strings.Count(got, "cursor=") != 1 || q.Get("cursor") != "0" {
		t.Errorf("url %q should carry exactly one cursor=0", got)
	}
	if strings.Count(got, "?") != 1 {
		t.Errorf("url %q has more than one query string", got)
	}
	if want := []string{"live", "historical"}; strings.Join(q["wantedEventTypes"], ",") != strings.Join(want, ",") {
		t.Errorf("wantedEventTypes = %v, want %v", q["wantedEventTypes"], want)
	}
	if len(q["wantedCollections"]) != 2 {
		t.Errorf("wantedCollections = %v, want both", q["wantedCollections"])
	}
}

func TestStreamDecodesLive(t *testing.T) {
	srv := newStreamServer(t, func(_ int, con *websocket.Conn) {
		for _, frame := range []string{
			`{"did":"d","time_us":1,"kind":"commit","commit":{"rev":"r","operation":"create","live":true}}`,
			`{"did":"d","time_us":2,"kind":"commit","commit":{"rev":"r","operation":"create","live":false}}`,
			`{"did":"d","time_us":3,"kind":"commit","commit":{"rev":"r","operation":"create"}}`,
		} {
			con.WriteMessage(websocket.TextMessage, []byte(frame))
		}
		drain(con)
	})
	s := testStream(t, StreamConfig{Endpoint: srv.endpoint()})

	live := map[int64]bool{}
	run(t, s, func(e *Event) bool {
		live[e.TimeUS] = e.Commit.IsLive()
		return e.TimeUS == 3
	})
	if want := map[int64]bool{1: true, 2: false, 3: true}; len(live) != 3 || live[1] != want[1] || live[2] != want[2] || live[3] != want[3] {
		t.Fatalf("live = %v, want %v (absent means live)", live, want)
	}
}

func TestStreamSkipsUndecodableFrames(t *testing.T) {
	srv := newStreamServer(t, func(_ int, con *websocket.Conn) {
		con.WriteMessage(websocket.TextMessage, []byte(`not json`))
		con.WriteMessage(websocket.TextMessage, []byte(`{"did":"d","time_us":7}`))
		drain(con)
	})
	run(t, testStream(t, StreamConfig{Endpoint: srv.endpoint()}), func(e *Event) bool { return e.TimeUS == 7 })
}

func TestStreamResumesFromTheHandlersCursorRewound(t *testing.T) {
	var mu sync.Mutex
	cursor, have := int64(0), false
	srv := newStreamServer(t, func(n int, con *websocket.Conn) {
		if n == 0 {
			con.WriteMessage(websocket.TextMessage, []byte(`{"did":"d","time_us":10000000}`))
		}
		if n == 1 {
			con.WriteMessage(websocket.TextMessage, []byte(`{"did":"d","time_us":10000001}`))
		}
		drain(con)
	})
	s := testStream(t, StreamConfig{
		Endpoint: srv.endpoint(),
		Rewind:   2 * time.Second,
		Cursor: func() (int64, bool, error) {
			mu.Lock()
			defer mu.Unlock()
			return cursor, have, nil
		},
	})

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	done := make(chan struct{})
	go func() {
		defer close(done)
		s.Run(ctx, func(_ context.Context, e *Event) error {
			mu.Lock()
			cursor, have = e.TimeUS, true
			mu.Unlock()
			if e.TimeUS == 10_000_000 {
				return errors.New("drop the connection")
			}
			cancel()
			return nil
		})
	}()
	<-done

	if got := srv.query(0).Get("cursor"); got != "0" {
		t.Errorf("empty db cursor = %q, want 0 so a fresh db replays", got)
	}
	if got := srv.query(1).Get("cursor"); got != "8000000" {
		t.Errorf("reconnect cursor = %q, want the db cursor minus the rewind", got)
	}
}

func TestStreamOnlyReadsTheCursorOncePerConnect(t *testing.T) {
	var reads int
	srv := newStreamServer(t, func(_ int, con *websocket.Conn) {
		con.WriteMessage(websocket.TextMessage, []byte(`{"did":"d","time_us":5}`))
		drain(con)
	})
	s := testStream(t, StreamConfig{
		Endpoint: srv.endpoint(),
		Cursor:   func() (int64, bool, error) { reads++; return 0, false, nil },
	})
	run(t, s, func(e *Event) bool { return true })
	if reads != 1 {
		t.Fatalf("cursor source read %d times for one connection, want 1", reads)
	}
}

func TestStreamCursorErrorStopsTheConnect(t *testing.T) {
	srv := newStreamServer(t, func(int, *websocket.Conn) { t.Error("connected despite an unreadable cursor") })
	s := testStream(t, StreamConfig{
		Endpoint: srv.endpoint(),
		Cursor:   func() (int64, bool, error) { return 0, false, errors.New("db down") },
	})
	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()
	if err := s.Run(ctx, func(context.Context, *Event) error { return nil }); err != nil {
		t.Fatal(err)
	}
}

func TestStreamReadTimeoutReconnects(t *testing.T) {
	srv := newStreamServer(t, func(n int, con *websocket.Conn) {
		if n == 0 {
			drain(con)
			return
		}
		con.WriteMessage(websocket.TextMessage, []byte(`{"did":"d","time_us":9}`))
		drain(con)
	})
	s := testStream(t, StreamConfig{Endpoint: srv.endpoint(), ReadTimeout: 100 * time.Millisecond})
	run(t, s, func(e *Event) bool { return e.TimeUS == 9 })
}

func TestStreamPingsKeepAConnectionAlive(t *testing.T) {
	srv := newStreamServer(t, func(n int, con *websocket.Conn) {
		for range 6 {
			time.Sleep(60 * time.Millisecond)
			con.WriteControl(websocket.PingMessage, nil, time.Now().Add(time.Second))
		}
		con.WriteMessage(websocket.TextMessage, []byte(`{"did":"d","time_us":4}`))
		drain(con)
	})
	s := testStream(t, StreamConfig{Endpoint: srv.endpoint(), ReadTimeout: 150 * time.Millisecond})
	run(t, s, func(e *Event) bool { return e.TimeUS == 4 })
	if n := len(srv.queries); n != 1 {
		t.Fatalf("connected %d times, pings should have kept the first alive", n)
	}
}

func TestStreamDecodesZstdWithTheJetstreamDictionary(t *testing.T) {
	enc, err := zstd.NewWriter(nil, zstd.WithEncoderDict(zstdDictionary))
	if err != nil {
		t.Fatal(err)
	}
	srv := newStreamServer(t, func(_ int, con *websocket.Conn) {
		frame := enc.EncodeAll([]byte(`{"did":"d","time_us":6,"kind":"commit","commit":{"live":false}}`), nil)
		con.WriteMessage(websocket.BinaryMessage, frame)
		drain(con)
	})
	s := testStream(t, StreamConfig{Endpoint: srv.endpoint(), Compress: true, UserAgent: "test"})
	run(t, s, func(e *Event) bool { return e.TimeUS == 6 && !e.Commit.IsLive() })
	if got := srv.headers[0].Get("Socket-Encoding"); got != "zstd" {
		t.Errorf("Socket-Encoding = %q, want zstd", got)
	}
	if got := srv.headers[0].Get("User-Agent"); got != "test" {
		t.Errorf("User-Agent = %q, want test", got)
	}
}

func TestStreamDecodesZstdWithoutADictionary(t *testing.T) {
	enc, err := zstd.NewWriter(nil)
	if err != nil {
		t.Fatal(err)
	}
	srv := newStreamServer(t, func(_ int, con *websocket.Conn) {
		frame := enc.EncodeAll([]byte(`{"did":"d","time_us":7,"kind":"commit","commit":{"live":false}}`), nil)
		con.WriteMessage(websocket.BinaryMessage, frame)
		drain(con)
	})
	s := testStream(t, StreamConfig{Endpoint: srv.endpoint(), Compress: true})
	run(t, s, func(e *Event) bool { return e.TimeUS == 7 && !e.Commit.IsLive() })
}

func TestStreamWithoutACursorTailsWhenAsked(t *testing.T) {
	srv := newStreamServer(t, func(_ int, con *websocket.Conn) { drain(con) })
	s := testStream(t, StreamConfig{Endpoint: srv.endpoint(), TailWhenNoCursor: true})
	ctx, cancel := context.WithCancel(context.Background())
	go s.Run(ctx, func(context.Context, *Event) error { return nil })
	for srv.connects() == 0 {
		time.Sleep(time.Millisecond)
	}
	cancel()
	if _, ok := srv.query(0)["cursor"]; ok {
		t.Fatal("a stream with no cursor should connect without one")
	}
}