// Command warp-loadtest drives a homeserver the way a crowd of clients
// would, and reports how it held up.
//
// It registers users, creates rooms and has every user join them, then for
// the length of the run each user sends messages at a steady rate while
// syncing in a long-poll loop, the way a client sits in /sync. Every message
// carries the time it was sent, so a user's sync can measure how long the
// server took to deliver what somebody else said — the number a user of a
// chat server actually feels — alongside the latency of each request.
//
// It speaks the plain Client-Server API and works against any homeserver,
// not only Warp. Point it at a test server: it creates accounts and rooms
// it never cleans up. Rate limiting on the server has to be off, or the
// run measures the rate limiter.
//
//	warp-loadtest -url http://localhost:8008 -users 50 -rooms 10 -rate 0.5 -duration 2m
package main

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"slices"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"text/tabwriter"
	"time"
)

type options struct {
	baseURL       string
	users         int
	rooms         int
	rate          float64
	duration      time.Duration
	prefix        string
	password      string
	regToken      string
	sliding       bool
	setupParallel int
}

func main() {
	var o options
	flag.StringVar(&o.baseURL, "url", "http://localhost:8008", "the homeserver's client API base URL")
	flag.IntVar(&o.users, "users", 20, "how many users to register")
	flag.IntVar(&o.rooms, "rooms", 5, "how many rooms to create; every user joins every room")
	flag.Float64Var(&o.rate, "rate", 0.2, "messages per second per user")
	flag.DurationVar(&o.duration, "duration", time.Minute, "how long to send for")
	flag.StringVar(&o.prefix, "prefix", "", "localpart prefix for the users (default: random)")
	flag.StringVar(&o.password, "password", "loadtest-password-1234", "the users' password")
	flag.StringVar(&o.regToken, "registration-token", "", "a registration token, for a server that requires one")
	flag.BoolVar(&o.sliding, "sliding", false, "sync with Simplified Sliding Sync instead of /sync")
	flag.IntVar(&o.setupParallel, "setup-parallelism", 8, "concurrent requests while registering and joining")
	flag.Parse()

	if o.prefix == "" {
		var b [3]byte
		_, _ = rand.Read(b[:])
		o.prefix = "lt" + hex.EncodeToString(b[:])
	}
	o.baseURL = strings.TrimRight(o.baseURL, "/")

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
	defer stop()
	if err := run(ctx, o); err != nil {
		fmt.Fprintln(os.Stderr, "warp-loadtest:", err)
		os.Exit(1)
	}
}

func run(ctx context.Context, o options) error {
	stats := newStats()
	c := &client{base: o.baseURL, http: &http.Client{Timeout: 90 * time.Second}, stats: stats}

	fmt.Printf("registering %d users as %s-N on %s\n", o.users, o.prefix, o.baseURL)
	users := make([]*user, o.users)
	if err := parallel(ctx, o.users, o.setupParallel, func(i int) error {
		u, err := c.register(ctx, fmt.Sprintf("%s-%d", o.prefix, i), o.password, o.regToken)
		users[i] = u
		return err
	}); err != nil {
		return err
	}

	fmt.Printf("creating %d rooms and joining every user to each\n", o.rooms)
	rooms := make([]string, o.rooms)
	for i := range rooms {
		roomID, err := c.createRoom(ctx, users[0], fmt.Sprintf("load test %s #%d", o.prefix, i))
		if err != nil {
			return err
		}
		rooms[i] = roomID
	}
	if err := parallel(ctx, len(users)-1, o.setupParallel, func(i int) error {
		for _, roomID := range rooms {
			if err := c.join(ctx, users[i+1], roomID); err != nil {
				return err
			}
		}
		return nil
	}); err != nil {
		return err
	}

	fmt.Printf("sending %.2f msg/s per user for %s (%s)\n", o.rate, o.duration, map[bool]string{true: "sliding sync", false: "/sync"}[o.sliding])

	var wg sync.WaitGroup
	syncCtx, stopSync := context.WithCancel(ctx)
	for _, u := range users {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if o.sliding {
				c.slidingLoop(syncCtx, u, len(rooms))
			} else {
				c.syncLoop(syncCtx, u)
			}
		}()
	}
	// The syncs need a moment to reach the top of the stream before the
	// first message, or the first delivery latencies would include the
	// initial sync.
	time.Sleep(2 * time.Second)

	runCtx, cancel := context.WithTimeout(ctx, o.duration)
	defer cancel()
	var senders sync.WaitGroup
	start := time.Now()
	for i, u := range users {
		senders.Add(1)
		go func() {
			defer senders.Done()
			c.sendLoop(runCtx, u, rooms, o.rate, float64(i)/float64(len(users)))
		}()
	}
	senders.Wait()
	elapsed := time.Since(start)

	// Give the last messages time to be delivered before the syncs stop.
	time.Sleep(3 * time.Second)
	stopSync()
	wg.Wait()

	stats.report(os.Stdout, elapsed)
	return nil
}

// parallel runs fn for 0..n-1, at most width at once, stopping at the
// first error.
func parallel(ctx context.Context, n, width int, fn func(int) error) error {
	sem := make(chan struct{}, max(width, 1))
	var (
		wg       sync.WaitGroup
		mu       sync.Mutex
		firstErr error
	)
	for i := range n {
		if ctx.Err() != nil {
			break
		}
		sem <- struct{}{}
		wg.Add(1)
		go func() {
			defer func() { <-sem; wg.Done() }()
			if err := fn(i); err != nil {
				mu.Lock()
				if firstErr == nil {
					firstErr = err
				}
				mu.Unlock()
			}
		}()
	}
	wg.Wait()
	if firstErr != nil {
		return firstErr
	}
	return ctx.Err()
}

type user struct {
	id    string
	token string
}

type client struct {
	base  string
	http  *http.Client
	stats *stats
	txn   atomic.Int64
}

// do sends one request, timing it under op, and decodes the response into
// out when there is one.
func (c *client) do(ctx context.Context, op, method, path, token string, body, out any) (int, error) {
	var reader io.Reader
	if body != nil {
		raw, err := json.Marshal(body)
		if err != nil {
			return 0, err
		}
		reader = bytes.NewReader(raw)
	}
	req, err := http.NewRequestWithContext(ctx, method, c.base+path, reader)
	if err != nil {
		return 0, err
	}
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	if token != "" {
		req.Header.Set("Authorization", "Bearer "+token)
	}

	start := time.Now()
	resp, err := c.http.Do(req)
	if err != nil {
		if ctx.Err() == nil {
			c.stats.fail(op)
		}
		return 0, err
	}
	defer func() { _ = resp.Body.Close() }()
	raw, err := io.ReadAll(resp.Body)
	c.stats.observe(op, time.Since(start), resp.StatusCode < 400 || resp.StatusCode == http.StatusUnauthorized && op == "register")
	if err != nil {
		return resp.StatusCode, err
	}
	if out != nil && len(raw) > 0 {
		_ = json.Unmarshal(raw, out)
	}
	if resp.StatusCode >= 400 && (resp.StatusCode != http.StatusUnauthorized || op != "register") {
		return resp.StatusCode, fmt.Errorf("%s %s: %d %s", method, path, resp.StatusCode, strings.TrimSpace(string(raw)))
	}
	return resp.StatusCode, nil
}

// register walks the user-interactive registration flow, answering the
// dummy and registration-token stages.
func (c *client) register(ctx context.Context, localpart, password, token string) (*user, error) {
	body := map[string]any{"username": localpart, "password": password, "inhibit_login": false}
	var challenge struct {
		Session string `json:"session"`
		Flows   []struct {
			Stages []string `json:"stages"`
		} `json:"flows"`
		UserID      string `json:"user_id"`
		AccessToken string `json:"access_token"`
	}
	if _, err := c.do(ctx, "register", http.MethodPost, "/_matrix/client/v3/register", "", body, &challenge); err != nil {
		return nil, err
	}
	if challenge.AccessToken != "" {
		return &user{id: challenge.UserID, token: challenge.AccessToken}, nil
	}
	if len(challenge.Flows) == 0 {
		return nil, errors.New("register: the server offered no flow")
	}
	for _, stage := range challenge.Flows[0].Stages {
		auth := map[string]any{"type": stage, "session": challenge.Session}
		switch stage {
		case "m.login.dummy":
		case "m.login.registration_token":
			if token == "" {
				return nil, errors.New("register: the server wants a registration token; pass -registration-token")
			}
			auth["token"] = token
		default:
			return nil, fmt.Errorf("register: cannot complete the %s stage", stage)
		}
		body["auth"] = auth
		var next struct {
			UserID      string `json:"user_id"`
			AccessToken string `json:"access_token"`
		}
		if _, err := c.do(ctx, "register", http.MethodPost, "/_matrix/client/v3/register", "", body, &next); err != nil {
			return nil, err
		}
		if next.AccessToken != "" {
			return &user{id: next.UserID, token: next.AccessToken}, nil
		}
	}
	return nil, errors.New("register: the flow ended without an access token")
}

func (c *client) createRoom(ctx context.Context, u *user, name string) (string, error) {
	var out struct {
		RoomID string `json:"room_id"`
	}
	_, err := c.do(ctx, "createRoom", http.MethodPost, "/_matrix/client/v3/createRoom", u.token,
		map[string]any{"name": name, "preset": "public_chat"}, &out)
	return out.RoomID, err
}

func (c *client) join(ctx context.Context, u *user, roomID string) error {
	_, err := c.do(ctx, "join", http.MethodPost, "/_matrix/client/v3/join/"+url.PathEscape(roomID), u.token, map[string]any{}, nil)
	return err
}

// sentMarker prefixes a message body with the time it was sent.
const sentMarker = "warp-loadtest sent="

// sendLoop sends at a steady rate until the context ends. Users start
// staggered across the first interval — phase is this user's place in it,
// from 0 to 1 — so that the load is a flow rather than a pulse.
func (c *client) sendLoop(ctx context.Context, u *user, rooms []string, rate, phase float64) {
	if rate <= 0 {
		return
	}
	interval := time.Duration(float64(time.Second) / rate)
	select {
	case <-ctx.Done():
		return
	case <-time.After(time.Duration(phase * float64(interval))):
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	n := 0
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
		roomID := rooms[n%len(rooms)]
		n++
		body := sentMarker + strconv.FormatInt(time.Now().UnixNano(), 10)
		path := fmt.Sprintf("/_matrix/client/v3/rooms/%s/send/m.room.message/lt%d",
			url.PathEscape(roomID), c.txn.Add(1))
		_, _ = c.do(ctx, "send", http.MethodPut, path, u.token, map[string]any{"msgtype": "m.text", "body": body}, nil)
	}
}

// syncLoop sits in /sync, measuring how long each message took to arrive.
func (c *client) syncLoop(ctx context.Context, u *user) {
	since := ""
	for ctx.Err() == nil {
		path := "/_matrix/client/v3/sync?timeout=30000"
		if since != "" {
			path += "&since=" + url.QueryEscape(since)
		}
		var resp struct {
			NextBatch string `json:"next_batch"`
			Rooms     struct {
				Join map[string]struct {
					Timeline struct {
						Events []timelineEvent `json:"events"`
					} `json:"timeline"`
				} `json:"join"`
			} `json:"rooms"`
		}
		op := "sync"
		if since == "" {
			op = "initial sync"
		}
		if _, err := c.do(ctx, op, http.MethodGet, path, u.token, nil, &resp); err != nil {
			if ctx.Err() == nil {
				time.Sleep(time.Second)
			}
			continue
		}
		received := time.Now()
		if since != "" {
			for _, room := range resp.Rooms.Join {
				c.deliveries(u, room.Timeline.Events, received)
			}
		}
		since = resp.NextBatch
	}
}

// slidingLoop is syncLoop on Simplified Sliding Sync.
func (c *client) slidingLoop(ctx context.Context, u *user, rooms int) {
	const path = "/_matrix/client/unstable/org.matrix.simplified_msc3575/sync"
	body := map[string]any{
		"lists": map[string]any{"all": map[string]any{
			"ranges": [][2]int{{0, max(rooms-1, 0)}}, "timeline_limit": 20,
			"required_state": [][]string{{"m.room.name", ""}},
		}},
	}
	pos := ""
	for ctx.Err() == nil {
		q := "?timeout=30000"
		if pos != "" {
			q += "&pos=" + url.QueryEscape(pos)
		}
		var resp struct {
			Pos   string `json:"pos"`
			Rooms map[string]struct {
				Timeline []timelineEvent `json:"timeline"`
			} `json:"rooms"`
		}
		op := "sliding sync"
		if pos == "" {
			op = "initial sliding sync"
		}
		status, err := c.do(ctx, op, http.MethodPost, path+q, u.token, body, &resp)
		if err != nil {
			if status == http.StatusBadRequest {
				pos = "" // M_UNKNOWN_POS: start again
			}
			if ctx.Err() == nil {
				time.Sleep(time.Second)
			}
			continue
		}
		received := time.Now()
		if pos != "" {
			for _, room := range resp.Rooms {
				c.deliveries(u, room.Timeline, received)
			}
		}
		pos = resp.Pos
	}
}

type timelineEvent struct {
	Type    string `json:"type"`
	Sender  string `json:"sender"`
	Content struct {
		Body string `json:"body"`
	} `json:"content"`
}

// deliveries records how long each message from somebody else took to
// reach this user.
func (c *client) deliveries(u *user, events []timelineEvent, received time.Time) {
	for _, e := range events {
		if e.Type != "m.room.message" || e.Sender == u.id {
			continue
		}
		raw, ok := strings.CutPrefix(e.Content.Body, sentMarker)
		if !ok {
			continue
		}
		sent, err := strconv.ParseInt(raw, 10, 64)
		if err != nil {
			continue
		}
		c.stats.observe("delivery", received.Sub(time.Unix(0, sent)), true)
	}
}

// stats collects latencies per operation.
type stats struct {
	mu     sync.Mutex
	ops    map[string][]time.Duration
	errors map[string]int
}

func newStats() *stats {
	return &stats{ops: map[string][]time.Duration{}, errors: map[string]int{}}
}

func (s *stats) observe(op string, d time.Duration, ok bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.ops[op] = append(s.ops[op], d)
	if !ok {
		s.errors[op]++
	}
}

func (s *stats) fail(op string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.errors[op]++
}

func (s *stats) report(w io.Writer, elapsed time.Duration) {
	s.mu.Lock()
	defer s.mu.Unlock()

	names := make([]string, 0, len(s.ops))
	for name := range s.ops {
		names = append(names, name)
	}
	slices.Sort(names)

	fmt.Fprintf(w, "\nran for %s\n\n", elapsed.Round(time.Millisecond))
	tw := tabwriter.NewWriter(w, 0, 0, 2, ' ', tabwriter.AlignRight)
	fmt.Fprintln(tw, "operation\tcount\terrors\trate/s\tp50\tp90\tp99\tmax\t")
	for _, name := range names {
		d := s.ops[name]
		slices.Sort(d)
		rate := float64(len(d)) / elapsed.Seconds()
		fmt.Fprintf(tw, "%s\t%d\t%d\t%.1f\t%s\t%s\t%s\t%s\t\n", name, len(d), s.errors[name], rate,
			pct(d, 0.5), pct(d, 0.9), pct(d, 0.99), pct(d, 1))
	}
	_ = tw.Flush()
	fmt.Fprintln(w, "\nsync latencies include the long-poll wait; delivery is send-to-receipt across users.")
}

func pct(sorted []time.Duration, q float64) string {
	if len(sorted) == 0 {
		return "-"
	}
	i := int(q*float64(len(sorted))+0.5) - 1
	i = max(0, min(i, len(sorted)-1))
	return sorted[i].Round(100 * time.Microsecond).String()
}