package roomserver_test

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"path/filepath"
	"slices"
	"sync"
	"testing"
	"time"

	"github.com/dariuszjaniak/warp/server/internal/federation"
	"github.com/dariuszjaniak/warp/server/internal/matrix/event"
	"github.com/dariuszjaniak/warp/server/internal/matrix/id"
	"github.com/dariuszjaniak/warp/server/internal/matrix/roomversion"
	"github.com/dariuszjaniak/warp/server/internal/matrix/signing"
	"github.com/dariuszjaniak/warp/server/internal/roomserver"
	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// The federation tests run several roomservers in one process, each on its
// own database, connected by a transport that calls the other server's
// federation methods directly. Nothing about HTTP or X-Matrix is under
// test here — the federation package has its own tests for that — only
// whether two roomservers that can reach each other end up agreeing about
// a room.

// mesh is a set of servers that can reach each other.
type mesh struct {
	t       *testing.T
	servers map[string]*peer

	// disabled turns federation off everywhere, for the switch tests.
	mu       sync.Mutex
	disabled bool
}

// peer is one server of the mesh.
type peer struct {
	name  string
	api   *roomserver.API
	store storage.RoomStore
	key   *signing.KeyPair
	mesh  *mesh

	mu sync.Mutex
	// outbox is every event this server persisted that it would send:
	// its own, and the ones it relays.
	outbox []roomserver.Notification
}

func newMesh(t *testing.T, names ...string) *mesh {
	t.Helper()
	m := &mesh{t: t, servers: map[string]*peer{}}
	for _, name := range names {
		key, err := signing.Generate("k_" + name[:1])
		if err != nil {
			t.Fatal(err)
		}
		m.servers[name] = &peer{name: name, key: key, mesh: m}
	}
	for _, name := range names {
		m.servers[name].start(t)
	}
	return m
}

func (p *peer) start(t *testing.T) {
	t.Helper()
	path := filepath.ToSlash(filepath.Join(t.TempDir(), p.name+".db"))
	db, err := storage.Open(t.Context(), storage.Config{
		Driver:           storage.DriverSQLite,
		ConnectionString: "file:" + path + "?_pragma=busy_timeout(5000)",
	}, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatalf("open storage: %v", err)
	}
	t.Cleanup(func() { _ = db.Close() })
	if err := db.Migrate(t.Context()); err != nil {
		t.Fatalf("migrate: %v", err)
	}

	api, err := roomserver.New(roomserver.Options{
		ServerName:        p.name,
		Store:             db.Rooms(),
		SigningKey:        p.key,
		Profiles:          anyProfile{},
		FederationEnabled: func() bool { return !p.mesh.isDisabled() },
		Notifier:          p,
		Remote:            transport{from: p},
		Keys:              keys{m: p.mesh},
		Logger:            slog.New(slog.DiscardHandler),
	})
	if err != nil {
		t.Fatalf("roomserver.New: %v", err)
	}
	p.api, p.store = api, db.Rooms()
}

func (m *mesh) isDisabled() bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	return m.disabled
}

// anyProfile gives every local user a display name.
type anyProfile struct{}

func (anyProfile) Account(_ context.Context, localpart string) (storage.Account, error) {
	return storage.Account{Localpart: localpart, DisplayName: "User " + localpart}, nil
}

// EventPersisted implements roomserver.Notifier: it fills the outbox the
// way the real outbound queue picks its events.
func (p *peer) EventPersisted(_ context.Context, n roomserver.Notification) {
	if n.Origin != "" && !n.Relay {
		return
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	p.outbox = append(p.outbox, n)
}

// takeOutbox empties the outbox.
func (p *peer) takeOutbox() []roomserver.Notification {
	p.mu.Lock()
	defer p.mu.Unlock()
	out := p.outbox
	p.outbox = nil
	return out
}

// flush delivers every outbox to the servers in each event's room, the
// way the queue would, until nothing is left to send. It returns the
// errors the receivers reported, keyed by event ID.
func (m *mesh) flush() map[string]error {
	m.t.Helper()
	failures := map[string]error{}
	for range 20 {
		sent := false
		for _, name := range slices.Sorted(mapsKeys(m.servers)) {
			p := m.servers[name]
			for _, n := range p.takeOutbox() {
				sent = true
				m.deliver(p, n, failures)
			}
		}
		if !sent {
			return failures
		}
	}
	m.t.Fatal("flush did not settle")
	return nil
}

// deliver sends one event from a server to the others in its room.
func (m *mesh) deliver(from *peer, n roomserver.Notification, failures map[string]error) {
	m.t.Helper()
	servers, err := from.api.JoinedServers(m.t.Context(), n.RoomID)
	if err != nil {
		m.t.Fatalf("JoinedServers: %v", err)
	}
	// A membership event also goes to the server of the user it is
	// about, so that a kicked or banned server hears of it.
	if key, ok := n.Event.StateKey(); ok && n.Event.Type() == event.TypeMember {
		if server, err := id.ServerNameOf(key); err == nil && !slices.Contains(servers, server) {
			servers = append(servers, server)
		}
	}
	for _, server := range servers {
		if server == from.name || server == n.Origin {
			continue
		}
		to, ok := m.servers[server]
		if !ok {
			continue
		}
		if _, err := to.api.ReceivePDU(m.t.Context(), from.name, n.Event.JSON()); err != nil {
			failures[n.Event.EventID()] = err
		}
	}
}

// deliverOnly sends one specific event, bypassing the outbox: the tests
// that leave gaps use it.
func (m *mesh) deliverOnly(from, to string, e *event.Event) error {
	_, err := m.servers[to].api.ReceivePDU(m.t.Context(), from, e.JSON())
	return err
}

func mapsKeys[V any](m map[string]V) func(func(string) bool) {
	return func(yield func(string) bool) {
		for k := range m {
			if !yield(k) {
				return
			}
		}
	}
}

// keys verifies signatures with the mesh's keys, standing in for the key
// ring.
type keys struct{ m *mesh }

func (k keys) VerifyEventSignature(_ context.Context, serverName string, e *event.Event) error {
	p, ok := k.m.servers[serverName]
	if !ok {
		return fmt.Errorf("no key for %s", serverName)
	}
	return e.VerifySignature(serverName, p.key.KeyID(), p.key.Public)
}

// transport is RemoteServers for one peer: every call goes straight to
// the destination's federation-side method, with this peer as the origin.
type transport struct{ from *peer }

func (t transport) to(destination string) (*roomserver.API, error) {
	if t.from.mesh.isDisabled() {
		return nil, federation.ErrDisabled
	}
	p, ok := t.from.mesh.servers[destination]
	if !ok {
		return nil, fmt.Errorf("%s is unreachable", destination)
	}
	return p.api, nil
}

// httpError turns a roomserver error into what the HTTP layer would have
// sent back, so that the caller's error mapping is exercised.
func httpError(err error) error {
	switch {
	case err == nil:
		return nil
	case errors.Is(err, roomserver.ErrIncompatibleRoomVersion):
		return &federation.HTTPError{Status: http.StatusBadRequest, Code: "M_INCOMPATIBLE_ROOM_VERSION", Message: err.Error()}
	case errors.Is(err, roomserver.ErrForbidden), errors.Is(err, roomserver.ErrRejected), errors.Is(err, roomserver.ErrNotInRoom):
		return &federation.HTTPError{Status: http.StatusForbidden, Code: "M_FORBIDDEN", Message: err.Error()}
	case errors.Is(err, roomserver.ErrNoSuchRoom), errors.Is(err, roomserver.ErrNoSuchEvent), errors.Is(err, roomserver.ErrNoSuchAlias):
		return &federation.HTTPError{Status: http.StatusNotFound, Code: "M_NOT_FOUND", Message: err.Error()}
	}
	return &federation.HTTPError{Status: http.StatusInternalServerError, Code: "M_UNKNOWN", Message: err.Error()}
}

func (t transport) MakeJoin(ctx context.Context, dest, roomID, userID string, versions []string) (federation.MakeJoinResponse, error) {
	api, err := t.to(dest)
	if err != nil {
		return federation.MakeJoinResponse{}, err
	}
	version, tmpl, err := api.MakeJoin(ctx, t.from.name, roomID, userID, versions)
	return federation.MakeJoinResponse{RoomVersion: version, Event: tmpl}, httpError(err)
}

func (t transport) SendJoin(ctx context.Context, dest, roomID, eventID string, raw json.RawMessage) (federation.SendJoinResponse, error) {
	api, err := t.to(dest)
	if err != nil {
		return federation.SendJoinResponse{}, err
	}
	res, err := api.SendJoin(ctx, t.from.name, roomID, eventID, raw)
	return federation.SendJoinResponse{Origin: dest, State: res.State, AuthChain: res.AuthChain, Event: res.Event}, httpError(err)
}

func (t transport) MakeLeave(ctx context.Context, dest, roomID, userID string) (federation.MakeJoinResponse, error) {
	api, err := t.to(dest)
	if err != nil {
		return federation.MakeJoinResponse{}, err
	}
	version, tmpl, err := api.MakeLeave(ctx, t.from.name, roomID, userID)
	return federation.MakeJoinResponse{RoomVersion: version, Event: tmpl}, httpError(err)
}

func (t transport) SendLeave(ctx context.Context, dest, roomID, eventID string, raw json.RawMessage) error {
	api, err := t.to(dest)
	if err != nil {
		return err
	}
	return httpError(api.SendLeave(ctx, t.from.name, roomID, eventID, raw))
}

func (t transport) Invite(ctx context.Context, dest, _, _ string, req federation.InviteRequest) (federation.InviteResponse, error) {
	api, err := t.to(dest)
	if err != nil {
		return federation.InviteResponse{}, err
	}
	signed, err := api.ReceiveInvite(ctx, t.from.name, req.RoomVersion, req.Event, req.InviteRoomState)
	return federation.InviteResponse{Event: signed}, httpError(err)
}

func (t transport) Event(ctx context.Context, dest, eventID string) (federation.Transaction, error) {
	api, err := t.to(dest)
	if err != nil {
		return federation.Transaction{}, err
	}
	raw, err := api.EventForServer(ctx, t.from.name, eventID)
	if err != nil {
		return federation.Transaction{}, httpError(err)
	}
	return federation.Transaction{Origin: dest, PDUs: []json.RawMessage{raw}}, nil
}

func (t transport) EventAuth(ctx context.Context, dest, roomID, eventID string) (federation.EventAuthResponse, error) {
	api, err := t.to(dest)
	if err != nil {
		return federation.EventAuthResponse{}, err
	}
	chain, err := api.EventAuthForServer(ctx, t.from.name, roomID, eventID)
	return federation.EventAuthResponse{AuthChain: chain}, httpError(err)
}

func (t transport) StateIDs(ctx context.Context, dest, roomID, eventID string) (federation.StateIDsResponse, error) {
	api, err := t.to(dest)
	if err != nil {
		return federation.StateIDsResponse{}, err
	}
	state, chain, err := api.StateIDsForServer(ctx, t.from.name, roomID, eventID)
	return federation.StateIDsResponse{PDUIDs: state, AuthChainIDs: chain}, httpError(err)
}

func (t transport) State(ctx context.Context, dest, roomID, eventID string) (federation.StateResponse, error) {
	api, err := t.to(dest)
	if err != nil {
		return federation.StateResponse{}, err
	}
	state, chain, err := api.StateForServer(ctx, t.from.name, roomID, eventID)
	return federation.StateResponse{PDUs: state, AuthChain: chain}, httpError(err)
}

func (t transport) MissingEvents(ctx context.Context, dest, roomID string, req federation.MissingEventsRequest) (federation.MissingEventsResponse, error) {
	api, err := t.to(dest)
	if err != nil {
		return federation.MissingEventsResponse{}, err
	}
	events, err := api.MissingEventsForServer(ctx, t.from.name, roomID, req.EarliestEvents, req.LatestEvents, req.Limit, req.MinDepth)
	return federation.MissingEventsResponse{Events: events}, httpError(err)
}

func (t transport) QueryDirectory(ctx context.Context, dest, alias string) (federation.DirectoryResponse, error) {
	api, err := t.to(dest)
	if err != nil {
		return federation.DirectoryResponse{}, err
	}
	resolved, err := api.QueryAliasForServer(ctx, alias)
	return federation.DirectoryResponse{RoomID: resolved.RoomID, Servers: resolved.Servers}, httpError(err)
}

func (t transport) Hierarchy(ctx context.Context, dest, roomID string, suggestedOnly bool) (json.RawMessage, error) {
	api, err := t.to(dest)
	if err != nil {
		return nil, err
	}
	resp, err := api.FederatedHierarchy(ctx, t.from.name, roomID, suggestedOnly)
	if err != nil {
		return nil, httpError(err)
	}
	return json.Marshal(resp)
}

func (t transport) PublicRooms(ctx context.Context, dest string, limit int, since, search string) (json.RawMessage, error) {
	api, err := t.to(dest)
	if err != nil {
		return nil, err
	}
	resp, err := api.PublicRooms(ctx, roomserver.PublicRoomsRequest{Limit: limit, Since: since, Search: search})
	if err != nil {
		return nil, httpError(err)
	}
	return json.Marshal(resp)
}

// ---------------------------------------------------------------------------

const (
	serverA = "a.example"
	serverB = "b.example"
	serverC = "c.example"

	userA  = "@alice:" + serverA
	userA2 = "@anne:" + serverA
	userB  = "@bob:" + serverB
	userC  = "@cath:" + serverC
)

// send sends a message and returns its event ID.
func (p *peer) send(t *testing.T, roomID, sender, body string) string {
	t.Helper()
	eventID, err := p.api.SendEvent(t.Context(), roomserver.SendEventRequest{
		RoomID: roomID, Sender: sender, Type: "m.room.message",
		Content: json.RawMessage(`{"msgtype":"m.text","body":"` + body + `"}`),
	})
	if err != nil {
		t.Fatalf("%s: SendEvent: %v", p.name, err)
	}
	return eventID
}

// timelineIDs returns the event IDs of a room's timeline on a server,
// oldest first.
func (p *peer) timelineIDs(t *testing.T, roomID string) []string {
	t.Helper()
	events, err := p.store.Timeline(t.Context(), storage.TimelineRequest{
		RoomID: roomID, Limit: 1000, Direction: storage.DirectionForward,
	})
	if err != nil {
		t.Fatalf("Timeline: %v", err)
	}
	ids := make([]string, 0, len(events))
	for _, e := range events {
		ids = append(ids, e.EventID)
	}
	return ids
}

// membership reads a user's membership on a server.
func (p *peer) membership(t *testing.T, roomID, userID string) string {
	t.Helper()
	m, err := p.api.Membership(t.Context(), roomID, userID)
	if err != nil && !errors.Is(err, roomserver.ErrNoSuchRoom) {
		t.Fatalf("Membership: %v", err)
	}
	return m
}

// currentState returns a server's current state of a room as
// type/state_key -> event ID.
func (p *peer) currentState(t *testing.T, roomID string) map[string]string {
	t.Helper()
	entries, err := p.store.CurrentState(t.Context(), roomID)
	if err != nil {
		t.Fatalf("CurrentState: %v", err)
	}
	out := map[string]string{}
	for _, e := range entries {
		out[e.Type+"|"+e.StateKey] = e.EventID
	}
	return out
}

// publicRoom creates a public room on a server.
func (p *peer) publicRoom(t *testing.T, creator string) string {
	t.Helper()
	resp, err := p.api.CreateRoom(t.Context(), roomserver.CreateRoomRequest{
		Creator: creator, Preset: roomserver.PresetPublicChat, Name: "Federated",
	})
	if err != nil {
		t.Fatalf("CreateRoom: %v", err)
	}
	return resp.RoomID
}

func TestFederatedJoinAndConversation(t *testing.T) {
	for _, version := range []string{"10", "11", "12"} {
		t.Run("v"+version, func(t *testing.T) {
			m := newMesh(t, serverA, serverB)
			a, b := m.servers[serverA], m.servers[serverB]

			resp, err := a.api.CreateRoom(t.Context(), roomserver.CreateRoomRequest{
				Creator: userA, Preset: roomserver.PresetPublicChat, Name: "Federated", Version: version,
			})
			if err != nil {
				t.Fatalf("CreateRoom: %v", err)
			}
			roomID := resp.RoomID
			a.takeOutbox() // nobody else is in the room yet

			if _, err := b.api.JoinRoom(t.Context(), roomserver.JoinRequest{
				RoomID: roomID, User: userB, ServerNames: []string{serverA},
			}); err != nil {
				t.Fatalf("remote join: %v", err)
			}
			if got := b.membership(t, roomID, userB); got != event.MembershipJoin {
				t.Fatalf("B sees its user as %q", got)
			}
			if got := a.membership(t, roomID, userB); got != event.MembershipJoin {
				t.Fatalf("A sees B's user as %q", got)
			}
			if failures := m.flush(); len(failures) > 0 {
				t.Fatalf("delivery failures: %v", failures)
			}

			// Both servers hold the same room.
			if sa, sb := a.currentState(t, roomID), b.currentState(t, roomID); !mapsEqual(sa, sb) {
				t.Fatalf("current state differs:\nA: %v\nB: %v", sa, sb)
			}
			summary, err := b.store.RoomSummary(t.Context(), roomID)
			if err != nil {
				t.Fatalf("RoomSummary: %v", err)
			}
			if summary.Name != "Federated" || summary.JoinedMembers != 2 {
				t.Errorf("B's listing of the room = %+v", summary.RoomProfile)
			}

			// A conversation in both directions.
			m1 := a.send(t, roomID, userA, "hello from A")
			m2 := b.send(t, roomID, userB, "hello from B")
			if failures := m.flush(); len(failures) > 0 {
				t.Fatalf("delivery failures: %v", failures)
			}
			for _, p := range []*peer{a, b} {
				ids := p.timelineIDs(t, roomID)
				if !slices.Contains(ids, m1) || !slices.Contains(ids, m2) {
					t.Errorf("%s's timeline lacks a message: %v", p.name, ids)
				}
			}

			// The two messages were sent concurrently, so each server now
			// has two forward extremities; the next event joins them.
			m3 := a.send(t, roomID, userA, "merge")
			if failures := m.flush(); len(failures) > 0 {
				t.Fatalf("delivery failures: %v", failures)
			}
			for _, p := range []*peer{a, b} {
				exts, err := p.store.Extremities(t.Context(), roomID)
				if err != nil {
					t.Fatal(err)
				}
				if len(exts) != 1 || exts[0].EventID != m3 {
					t.Errorf("%s's extremities after the merge = %v", p.name, eventIDsOf(exts))
				}
			}
			if sa, sb := a.currentState(t, roomID), b.currentState(t, roomID); !mapsEqual(sa, sb) {
				t.Errorf("current state differs after the conversation:\nA: %v\nB: %v", sa, sb)
			}
		})
	}
}

func TestRemoteJoinFailsOverToTheNextServer(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a, b := m.servers[serverA], m.servers[serverB]
	roomID := a.publicRoom(t, userA)

	if _, err := b.api.JoinRoom(t.Context(), roomserver.JoinRequest{
		RoomID: roomID, User: userB, ServerNames: []string{"gone.example", serverA},
	}); err != nil {
		t.Fatalf("join with an unreachable first server: %v", err)
	}
	if got := b.membership(t, roomID, userB); got != event.MembershipJoin {
		t.Errorf("membership = %q", got)
	}
}

func TestRemoteJoinToAnInviteOnlyRoomIsRefused(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a, b := m.servers[serverA], m.servers[serverB]
	resp, err := a.api.CreateRoom(t.Context(), roomserver.CreateRoomRequest{Creator: userA, Preset: roomserver.PresetPrivateChat})
	if err != nil {
		t.Fatal(err)
	}

	_, err = b.api.JoinRoom(t.Context(), roomserver.JoinRequest{
		RoomID: resp.RoomID, User: userB, ServerNames: []string{serverA},
	})
	if !errors.Is(err, roomserver.ErrForbidden) {
		t.Fatalf("join to an invite-only room = %v, want ErrForbidden", err)
	}
	if _, err := b.store.Room(t.Context(), resp.RoomID); !errors.Is(err, storage.ErrNoSuchRoom) {
		t.Errorf("a refused join left the room behind on B: %v", err)
	}
}

func TestRemoteInviteJoinAndReject(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a, b := m.servers[serverA], m.servers[serverB]
	resp, err := a.api.CreateRoom(t.Context(), roomserver.CreateRoomRequest{
		Creator: userA, Preset: roomserver.PresetPrivateChat, Name: "Secret",
	})
	if err != nil {
		t.Fatal(err)
	}
	roomID := resp.RoomID

	if err := a.api.InviteUser(t.Context(), roomserver.InviteRequest{RoomID: roomID, Sender: userA, Target: userB}); err != nil {
		t.Fatalf("remote invite: %v", err)
	}
	if got := b.membership(t, roomID, userB); got != event.MembershipInvite {
		t.Fatalf("B sees its user as %q after the invite", got)
	}
	// The invite on A carries both servers' signatures.
	stored, err := a.store.CurrentStateEvent(t.Context(), roomID, event.TypeMember, userB)
	if err != nil {
		t.Fatal(err)
	}
	var sigs struct {
		Signatures map[string]any `json:"signatures"`
	}
	_ = json.Unmarshal(stored.JSON, &sigs)
	if _, ok := sigs.Signatures[serverB]; !ok {
		t.Error("the invite A holds is not countersigned by B")
	}

	// B shows the invite with the state A sent along.
	stripped, err := b.api.StrippedState(t.Context(), roomID, []string{userB})
	if err != nil {
		t.Fatalf("StrippedState: %v", err)
	}
	var sawName, sawInvite bool
	for _, s := range stripped {
		if s.Type == "m.room.name" {
			sawName = true
		}
		if s.Type == event.TypeMember && s.StateKey == userB {
			sawInvite = true
		}
	}
	if !sawName || !sawInvite {
		t.Errorf("B's invite state = %+v", stripped)
	}

	// Accepting goes through A, which is the inviter.
	if _, err := b.api.JoinRoom(t.Context(), roomserver.JoinRequest{RoomID: roomID, User: userB}); err != nil {
		t.Fatalf("accepting the invite: %v", err)
	}
	if got := a.membership(t, roomID, userB); got != event.MembershipJoin {
		t.Errorf("A sees B's user as %q after the join", got)
	}
	if got, _ := b.store.InviteState(t.Context(), roomID, userB); got != nil {
		t.Error("the invite state outlived the invite")
	}

	// B is in the room now, so a second invite to one of its users is
	// taken in as an ordinary event of the room rather than out of band.
	m.flush()
	if err := a.api.InviteUser(t.Context(), roomserver.InviteRequest{RoomID: roomID, Sender: userA, Target: "@zed:" + serverB}); err != nil {
		t.Fatalf("second invite: %v", err)
	}
	m.flush()
	if got := b.membership(t, roomID, "@zed:"+serverB); got != event.MembershipInvite {
		t.Fatalf("B sees zed as %q", got)
	}
}

func TestRejectingARemoteInvite(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a, b := m.servers[serverA], m.servers[serverB]
	resp, err := a.api.CreateRoom(t.Context(), roomserver.CreateRoomRequest{Creator: userA, Preset: roomserver.PresetPrivateChat})
	if err != nil {
		t.Fatal(err)
	}
	roomID := resp.RoomID
	if err := a.api.InviteUser(t.Context(), roomserver.InviteRequest{RoomID: roomID, Sender: userA, Target: userB}); err != nil {
		t.Fatal(err)
	}

	if err := b.api.LeaveRoom(t.Context(), roomID, userB, "no thanks"); err != nil {
		t.Fatalf("rejecting: %v", err)
	}
	if got := b.membership(t, roomID, userB); got != event.MembershipLeave {
		t.Errorf("B sees its user as %q after rejecting", got)
	}
	if got := a.membership(t, roomID, userB); got != event.MembershipLeave {
		t.Errorf("A sees B's user as %q after the rejection", got)
	}
}

func TestRejectingAnInviteFromAnUnreachableServer(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a, b := m.servers[serverA], m.servers[serverB]
	resp, err := a.api.CreateRoom(t.Context(), roomserver.CreateRoomRequest{Creator: userA, Preset: roomserver.PresetPrivateChat})
	if err != nil {
		t.Fatal(err)
	}
	if err := a.api.InviteUser(t.Context(), roomserver.InviteRequest{RoomID: resp.RoomID, Sender: userA, Target: userB}); err != nil {
		t.Fatal(err)
	}
	delete(m.servers, serverA)

	if err := b.api.LeaveRoom(t.Context(), resp.RoomID, userB, ""); err != nil {
		t.Fatalf("rejecting locally: %v", err)
	}
	if got := b.membership(t, resp.RoomID, userB); got != event.MembershipLeave {
		t.Errorf("membership after a local rejection = %q", got)
	}
}

func TestMissingEventsAreFetched(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a, b := m.servers[serverA], m.servers[serverB]
	roomID := a.publicRoom(t, userA)
	if _, err := b.api.JoinRoom(t.Context(), roomserver.JoinRequest{RoomID: roomID, User: userB, ServerNames: []string{serverA}}); err != nil {
		t.Fatal(err)
	}
	m.flush()

	var sent []string
	for i := range 4 {
		sent = append(sent, a.send(t, roomID, userA, fmt.Sprint("m", i)))
	}
	outbox := a.takeOutbox()
	// Only the last one gets through; B has to ask for the other three.
	last := outbox[len(outbox)-1]
	if err := m.deliverOnly(serverA, serverB, last.Event); err != nil {
		t.Fatalf("delivering the last message: %v", err)
	}
	ids := b.timelineIDs(t, roomID)
	for _, id := range sent {
		if !slices.Contains(ids, id) {
			t.Errorf("B's timeline lacks %s", id)
		}
	}
	// In order.
	var positions []int
	for _, id := range sent {
		positions = append(positions, slices.Index(ids, id))
	}
	if !slices.IsSorted(positions) {
		t.Errorf("fetched events are out of order: %v", positions)
	}
}

func TestAWideGapIsClosedWithState(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a, b := m.servers[serverA], m.servers[serverB]
	roomID := a.publicRoom(t, userA)
	if _, err := b.api.JoinRoom(t.Context(), roomserver.JoinRequest{RoomID: roomID, User: userB, ServerNames: []string{serverA}}); err != nil {
		t.Fatal(err)
	}
	m.flush()

	// More than get_missing_events will be asked for, with a state change
	// in the part that is not fetched.
	for i := range 6 {
		a.send(t, roomID, userA, fmt.Sprint("early", i))
	}
	if _, err := a.api.SendEvent(t.Context(), roomserver.SendEventRequest{
		RoomID: roomID, Sender: userA, Type: "m.room.topic", StateKey: ptrTo(""),
		Content: json.RawMessage(`{"topic":"changed in the gap"}`),
	}); err != nil {
		t.Fatal(err)
	}
	for i := range 15 {
		a.send(t, roomID, userA, fmt.Sprint("late", i))
	}
	outbox := a.takeOutbox()
	last := outbox[len(outbox)-1]
	if err := m.deliverOnly(serverA, serverB, last.Event); err != nil {
		t.Fatalf("delivering across the gap: %v", err)
	}

	if !slices.Contains(b.timelineIDs(t, roomID), last.Event.EventID()) {
		t.Error("the event across the gap is not in B's timeline")
	}
	if sa, sb := a.currentState(t, roomID), b.currentState(t, roomID); !mapsEqual(sa, sb) {
		t.Errorf("current state differs after the gap:\nA: %v\nB: %v", sa, sb)
	}
}

func TestForgedEventIsRejected(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a, b := m.servers[serverA], m.servers[serverB]
	roomID := a.publicRoom(t, userA)
	if _, err := b.api.JoinRoom(t.Context(), roomserver.JoinRequest{RoomID: roomID, User: userB, ServerNames: []string{serverA}}); err != nil {
		t.Fatal(err)
	}
	m.flush()

	// An event in alice's name, signed by a key that is not A's.
	exts, err := b.store.Extremities(t.Context(), roomID)
	if err != nil {
		t.Fatal(err)
	}
	impostor, _ := signing.Generate("k_a")
	room, _ := b.store.Room(t.Context(), roomID)
	rules := mustRules(t, room.Version)
	state := b.currentState(t, roomID)
	authEvents := []string{state[event.TypePowerLevels+"|"], state[event.TypeMember+"|"+userA]}
	if !rules.CreateEventExcludedFromAuthEvents {
		authEvents = append(authEvents, state[event.TypeCreate+"|"])
	}
	forged, err := (&event.Builder{
		RoomID: roomID, Type: "m.room.message", Sender: userA,
		Content:    json.RawMessage(`{"body":"not really alice"}`),
		PrevEvents: []string{exts[0].EventID}, AuthEvents: authEvents, Depth: exts[0].Depth + 1,
	}).Build(fixedNow(), impostor, rules)
	if err != nil {
		t.Fatal(err)
	}
	if err := m.deliverOnly(serverA, serverB, forged); !errors.Is(err, roomserver.ErrBadEvent) {
		t.Errorf("a forged event = %v, want ErrBadEvent", err)
	}
	if slices.Contains(b.timelineIDs(t, roomID), forged.EventID()) {
		t.Error("the forged event is in the timeline")
	}
}

func TestEventFailingOnlyAgainstCurrentStateIsSoftFailed(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a, b := m.servers[serverA], m.servers[serverB]
	roomID := a.publicRoom(t, userA)
	if _, err := b.api.JoinRoom(t.Context(), roomserver.JoinRequest{RoomID: roomID, User: userB, ServerNames: []string{serverA}}); err != nil {
		t.Fatal(err)
	}
	m.flush()

	// B's user speaks, and the message is held back; meanwhile A bans
	// them, and the ban reaches B.
	b.send(t, roomID, userB, "last words")
	late := b.takeOutbox()[0].Event
	if err := a.api.BanUser(t.Context(), roomID, userA, userB, "spam"); err != nil {
		t.Fatalf("BanUser: %v", err)
	}
	if failures := m.flush(); len(failures) > 0 {
		t.Fatalf("delivery failures: %v", failures)
	}

	// The message was valid when it was sent, and is not now.
	if err := m.deliverOnly(serverB, serverA, late); err != nil {
		t.Fatalf("a soft-failed event is not an error to its sender: %v", err)
	}
	stored, err := a.store.Event(t.Context(), roomID, late.EventID())
	if err != nil {
		t.Fatalf("the soft-failed event was not kept: %v", err)
	}
	if !stored.Outlier {
		t.Error("the soft-failed event is not an outlier")
	}
	if slices.Contains(a.timelineIDs(t, roomID), late.EventID()) {
		t.Error("the soft-failed event is in A's timeline")
	}
}

func TestFederationSwitchRefusesRemoteWork(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a, b := m.servers[serverA], m.servers[serverB]
	roomID := a.publicRoom(t, userA)

	m.setDisabled(true)
	if _, err := b.api.JoinRoom(t.Context(), roomserver.JoinRequest{RoomID: roomID, User: userB, ServerNames: []string{serverA}}); !errors.Is(err, roomserver.ErrFederationDisabled) {
		t.Errorf("remote join while disabled = %v, want ErrFederationDisabled", err)
	}
	if err := a.api.InviteUser(t.Context(), roomserver.InviteRequest{RoomID: roomID, Sender: userA, Target: userB}); !errors.Is(err, roomserver.ErrFederationDisabled) {
		t.Errorf("remote invite while disabled = %v, want ErrFederationDisabled", err)
	}
	if _, _, err := a.api.MakeJoin(t.Context(), serverB, roomID, userB, []string{"11"}); !errors.Is(err, roomserver.ErrFederationDisabled) {
		t.Errorf("make_join while disabled = %v, want ErrFederationDisabled", err)
	}

	m.setDisabled(false)
	if _, err := b.api.JoinRoom(t.Context(), roomserver.JoinRequest{RoomID: roomID, User: userB, ServerNames: []string{serverA}}); err != nil {
		t.Errorf("remote join after re-enabling: %v", err)
	}
}

func TestThreeServersAgree(t *testing.T) {
	m := newMesh(t, serverA, serverB, serverC)
	a, b, c := m.servers[serverA], m.servers[serverB], m.servers[serverC]
	roomID := a.publicRoom(t, userA)

	if _, err := b.api.JoinRoom(t.Context(), roomserver.JoinRequest{RoomID: roomID, User: userB, ServerNames: []string{serverA}}); err != nil {
		t.Fatal(err)
	}
	m.flush()
	// C joins through B, which is not the room's creator: B relays the
	// join to A.
	if _, err := c.api.JoinRoom(t.Context(), roomserver.JoinRequest{RoomID: roomID, User: userC, ServerNames: []string{serverB}}); err != nil {
		t.Fatalf("C joining through B: %v", err)
	}
	if failures := m.flush(); len(failures) > 0 {
		t.Fatalf("delivery failures: %v", failures)
	}
	if got := a.membership(t, roomID, userC); got != event.MembershipJoin {
		t.Fatalf("A did not hear of C's join through B's relay: %q", got)
	}

	a.send(t, roomID, userA2Join(t, a, roomID), "from A")
	b.send(t, roomID, userB, "from B")
	c.send(t, roomID, userC, "from C")
	if failures := m.flush(); len(failures) > 0 {
		t.Fatalf("delivery failures: %v", failures)
	}
	final := a.send(t, roomID, userA, "settle")
	m.flush()

	sa, sb, sc := a.currentState(t, roomID), b.currentState(t, roomID), c.currentState(t, roomID)
	if !mapsEqual(sa, sb) || !mapsEqual(sa, sc) {
		t.Errorf("the three servers disagree:\nA: %v\nB: %v\nC: %v", sa, sb, sc)
	}
	for _, p := range []*peer{a, b, c} {
		if !slices.Contains(p.timelineIDs(t, roomID), final) {
			t.Errorf("%s lacks the final message", p.name)
		}
	}
}

// userA2Join joins a second local user on A, so that local joins in a
// federated room are covered too, and returns them.
func userA2Join(t *testing.T, a *peer, roomID string) string {
	t.Helper()
	if _, err := a.api.JoinRoom(t.Context(), roomserver.JoinRequest{RoomID: roomID, User: userA2}); err != nil {
		t.Fatalf("local join in a federated room: %v", err)
	}
	return userA2
}

func TestRemoteAliasAndDirectory(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a, b := m.servers[serverA], m.servers[serverB]
	resp, err := a.api.CreateRoom(t.Context(), roomserver.CreateRoomRequest{
		Creator: userA, Preset: roomserver.PresetPublicChat, Visibility: roomserver.VisibilityPublic,
		RoomAliasName: "lobby", Name: "Lobby",
	})
	if err != nil {
		t.Fatal(err)
	}

	resolved, err := b.api.ResolveAlias(t.Context(), "#lobby:"+serverA)
	if err != nil {
		t.Fatalf("ResolveAlias: %v", err)
	}
	if resolved.RoomID != resp.RoomID || !slices.Contains(resolved.Servers, serverA) {
		t.Errorf("ResolveAlias = %+v", resolved)
	}
	if _, err := b.api.ResolveAlias(t.Context(), "#nothing:"+serverA); !errors.Is(err, roomserver.ErrNoSuchAlias) {
		t.Errorf("unknown remote alias = %v, want ErrNoSuchAlias", err)
	}

	dir, err := b.api.PublicRooms(t.Context(), roomserver.PublicRoomsRequest{Server: serverA})
	if err != nil {
		t.Fatalf("PublicRooms: %v", err)
	}
	if len(dir.Chunk) != 1 || dir.Chunk[0].RoomID != resp.RoomID || dir.Chunk[0].Name != "Lobby" {
		t.Errorf("remote directory = %+v", dir)
	}
}

func TestServersOutsideARoomCannotReadIt(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a := m.servers[serverA]
	resp, err := a.api.CreateRoom(t.Context(), roomserver.CreateRoomRequest{Creator: userA, Preset: roomserver.PresetPrivateChat})
	if err != nil {
		t.Fatal(err)
	}
	exts, err := a.store.Extremities(t.Context(), resp.RoomID)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := a.api.EventForServer(t.Context(), serverB, exts[0].EventID); !errors.Is(err, roomserver.ErrForbidden) {
		t.Errorf("/event from outside = %v, want ErrForbidden", err)
	}
	if _, _, err := a.api.StateIDsForServer(t.Context(), serverB, resp.RoomID, exts[0].EventID); !errors.Is(err, roomserver.ErrForbidden) {
		t.Errorf("/state_ids from outside = %v, want ErrForbidden", err)
	}
	if _, err := a.api.BackfillForServer(t.Context(), serverB, resp.RoomID, []string{exts[0].EventID}, 10); !errors.Is(err, roomserver.ErrForbidden) {
		t.Errorf("/backfill from outside = %v, want ErrForbidden", err)
	}
}

func TestBackfillWalksBackwards(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a, b := m.servers[serverA], m.servers[serverB]
	roomID := a.publicRoom(t, userA)
	if _, err := b.api.JoinRoom(t.Context(), roomserver.JoinRequest{RoomID: roomID, User: userB, ServerNames: []string{serverA}}); err != nil {
		t.Fatal(err)
	}
	var sent []string
	for i := range 5 {
		sent = append(sent, a.send(t, roomID, userA, fmt.Sprint(i)))
	}
	events, err := a.api.BackfillForServer(t.Context(), serverB, roomID, []string{sent[4]}, 3)
	if err != nil {
		t.Fatalf("BackfillForServer: %v", err)
	}
	if len(events) != 3 {
		t.Fatalf("backfill returned %d events, want 3", len(events))
	}
	var got []string
	for _, raw := range events {
		var probe struct {
			Content struct {
				Body string `json:"body"`
			} `json:"content"`
		}
		_ = json.Unmarshal(raw, &probe)
		got = append(got, probe.Content.Body)
	}
	if !slices.Equal(got, []string{"4", "3", "2"}) {
		t.Errorf("backfill bodies = %v, want newest first from the starting event", got)
	}
}

func mapsEqual(a, b map[string]string) bool {
	if len(a) != len(b) {
		return false
	}
	for k, v := range a {
		if b[k] != v {
			return false
		}
	}
	return true
}

func eventIDsOf(events []storage.StoredEvent) []string {
	out := make([]string, 0, len(events))
	for _, e := range events {
		out = append(out, e.EventID)
	}
	return out
}

func ptrTo[T any](v T) *T { return &v }

func mustRules(t *testing.T, version string) roomversion.Rules {
	t.Helper()
	rules, err := roomversion.Get(roomversion.Version(version))
	if err != nil {
		t.Fatal(err)
	}
	return rules
}

func fixedNow() time.Time { return time.Unix(1_700_000_000, 0) }

func TestServerACLKeepsAServerOut(t *testing.T) {
	m := newMesh(t, serverA, serverB, serverC)
	a, b, c := m.servers[serverA], m.servers[serverB], m.servers[serverC]
	roomID := a.publicRoom(t, userA)
	for _, p := range []*peer{b, c} {
		user := userB
		if p == c {
			user = userC
		}
		if _, err := p.api.JoinRoom(t.Context(), roomserver.JoinRequest{RoomID: roomID, User: user, ServerNames: []string{serverA}}); err != nil {
			t.Fatal(err)
		}
	}
	m.flush()

	if _, err := a.api.SendEvent(t.Context(), roomserver.SendEventRequest{
		RoomID: roomID, Sender: userA, Type: roomserver.TypeServerACL, StateKey: ptrTo(""),
		Content: json.RawMessage(`{"allow":["*"],"deny":["` + serverC + `"]}`),
	}); err != nil {
		t.Fatalf("setting the ACL: %v", err)
	}
	if failures := m.flush(); len(failures) > 0 {
		t.Fatalf("delivery failures: %v", failures)
	}

	c.send(t, roomID, userC, "from a denied server")
	failures := m.flush()
	if len(failures) == 0 {
		t.Fatal("the denied server's event was accepted")
	}
	for _, err := range failures {
		if !errors.Is(err, roomserver.ErrForbidden) {
			t.Errorf("refusal = %v, want ErrForbidden", err)
		}
	}
	if _, _, err := a.api.MakeJoin(t.Context(), serverC, roomID, "@other:"+serverC, []string{"11", "12", "10"}); !errors.Is(err, roomserver.ErrForbidden) {
		t.Errorf("make_join from a denied server = %v, want ErrForbidden", err)
	}

	// B, which the ACL allows, is unaffected.
	b.send(t, roomID, userB, "from an allowed server")
	if failures := m.flush(); len(failures) > 0 {
		t.Errorf("an allowed server's event was refused: %v", failures)
	}
}

func (m *mesh) setDisabled(disabled bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.disabled = disabled
}

func TestFederatedSpaceHierarchy(t *testing.T) {
	m := newMesh(t, serverA, serverB)
	a, b := m.servers[serverA], m.servers[serverB]
	ctx := t.Context()

	space, err := a.api.CreateRoom(ctx, roomserver.CreateRoomRequest{
		Creator: userA, Preset: roomserver.PresetPublicChat, Name: "Space",
		CreationContent: json.RawMessage(`{"type":"m.space"}`),
	})
	if err != nil {
		t.Fatalf("CreateRoom space: %v", err)
	}
	public := a.publicRoom(t, userA)
	private, err := a.api.CreateRoom(ctx, roomserver.CreateRoomRequest{Creator: userA, Name: "Private"})
	if err != nil {
		t.Fatalf("CreateRoom private: %v", err)
	}
	for _, child := range []string{public, private.RoomID} {
		key := child
		if _, err := a.api.SendEvent(ctx, roomserver.SendEventRequest{
			RoomID: space.RoomID, Sender: userA, Type: roomserver.TypeSpaceChild, StateKey: &key,
			Content: json.RawMessage(`{"via":["` + serverA + `"]}`),
		}); err != nil {
			t.Fatalf("add child: %v", err)
		}
	}
	a.takeOutbox()

	if _, err := b.api.JoinRoom(ctx, roomserver.JoinRequest{
		RoomID: space.RoomID, User: userB, ServerNames: []string{serverA},
	}); err != nil {
		t.Fatalf("remote join: %v", err)
	}
	if failures := m.flush(); len(failures) > 0 {
		t.Fatalf("delivery failures: %v", failures)
	}

	// B holds the space itself and asks A about the rooms it is not in:
	// the public one comes back, the private one is A's to keep.
	page, err := b.api.Hierarchy(ctx, roomserver.HierarchyRequest{RoomID: space.RoomID, User: userB, MaxDepth: -1})
	if err != nil {
		t.Fatalf("Hierarchy: %v", err)
	}
	var got []string
	for _, room := range page.Rooms {
		got = append(got, room.RoomID)
	}
	if len(got) != 2 || got[0] != space.RoomID || got[1] != public {
		t.Fatalf("B's walk of the space = %v, want the space and the public room", got)
	}
	if page.Rooms[1].Name != "Federated" || page.Rooms[1].NumJoinedMembers != 1 {
		t.Errorf("the remote child = %+v", page.Rooms[1])
	}

	answer, err := a.api.FederatedHierarchy(ctx, serverB, space.RoomID, false)
	if err != nil {
		t.Fatalf("FederatedHierarchy: %v", err)
	}
	if len(answer.InaccessibleChildren) != 1 || answer.InaccessibleChildren[0] != private.RoomID {
		t.Errorf("inaccessible_children = %v, want the private room", answer.InaccessibleChildren)
	}
	if _, err := a.api.FederatedHierarchy(ctx, serverB, private.RoomID, false); !errors.Is(err, roomserver.ErrNoSuchRoom) {
		t.Errorf("FederatedHierarchy of the private room = %v, want ErrNoSuchRoom", err)
	}
}