package keyserver_test

import (
	"context"
	"encoding/json"
	"errors"
	"sync"
	"testing"

	"github.com/dariuszjaniak/warp/server/internal/federation"
	"github.com/dariuszjaniak/warp/server/internal/keyserver"
)

// snapshotRemote is remote.example with Carol's device list in it, and a
// count of how often it was asked.
type snapshotRemote struct {
	mu       sync.Mutex
	snapshot federation.UserDevicesResponse
	fetches  int
	queries  int
	down     bool
}

func (r *snapshotRemote) QueryUserKeys(context.Context, string, federation.UserKeysQueryRequest) (federation.UserKeysQueryResponse, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.queries++
	return federation.UserKeysQueryResponse{}, nil
}

func (r *snapshotRemote) ClaimUserKeys(context.Context, string, federation.UserKeysClaimRequest) (federation.UserKeysClaimResponse, error) {
	return federation.UserKeysClaimResponse{}, errors.New("unreachable")
}

func (r *snapshotRemote) UserDevices(_ context.Context, _, userID string) (federation.UserDevicesResponse, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.fetches++
	if r.down {
		return federation.UserDevicesResponse{}, &federation.HTTPError{Status: 502}
	}
	out := r.snapshot
	out.UserID = userID
	return out, nil
}

func (r *snapshotRemote) set(stream int64, devices ...federation.UserDevice) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.snapshot.StreamID = stream
	r.snapshot.Devices = devices
}

func (r *snapshotRemote) counts() (fetches, queries int) {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.fetches, r.queries
}

func queryCarol(t *testing.T, h harness) keyserver.QueryResponse {
	t.Helper()
	resp, err := h.keys.QueryKeys(t.Context(), alice, keyserver.QueryRequest{DeviceKeys: map[string][]string{carol: nil}})
	if err != nil {
		t.Fatal(err)
	}
	return resp
}

func carolDevice(t *testing.T, deviceID string) federation.UserDevice {
	t.Helper()
	_, keys := device(t, carol, deviceID)
	return federation.UserDevice{DeviceID: deviceID, Keys: keys, DeviceDisplayName: "Carol's " + deviceID}
}

func TestRemoteKeysAreCached(t *testing.T) {
	remote := &snapshotRemote{}
	master := newKey(t)
	remote.snapshot.MasterKey = signed(t, carol, master, crossSigningKey(carol, "master", master))
	remote.set(5, carolDevice(t, "C1"),
		// A device whose keys name somebody else is not Carol's.
		federation.UserDevice{DeviceID: "FAKE", Keys: json.RawMessage(`{"user_id":"@mallory:remote.example","device_id":"FAKE"}`)})
	h := newHarness(t, remote)

	first := queryCarol(t, h)
	second := queryCarol(t, h)
	if fetches, queries := remote.counts(); fetches != 1 || queries != 0 {
		t.Errorf("two queries made %d fetches and %d key queries, want 1 and 0", fetches, queries)
	}
	for i, resp := range []keyserver.QueryResponse{first, second} {
		raw, ok := resp.DeviceKeys[carol]["C1"]
		if !ok {
			t.Fatalf("query %d: device keys = %v", i, resp.DeviceKeys)
		}
		var keys struct {
			Unsigned map[string]string `json:"unsigned"`
		}
		_ = json.Unmarshal(raw, &keys)
		if keys.Unsigned["device_display_name"] != "Carol's C1" {
			t.Errorf("query %d: C1 = %s, want its display name", i, raw)
		}
		if _, ok := resp.DeviceKeys[carol]["FAKE"]; ok {
			t.Errorf("query %d: a device with another user's keys was served", i)
		}
		if resp.MasterKeys[carol] == nil {
			t.Errorf("query %d: master keys = %v", i, resp.MasterKeys)
		}
	}

	// Asking for one device is answered from the same copy.
	resp, _ := h.keys.QueryKeys(t.Context(), alice, keyserver.QueryRequest{DeviceKeys: map[string][]string{carol: {"NOPE"}}})
	if len(resp.DeviceKeys[carol]) != 0 {
		t.Errorf("a query for a device Carol does not have = %v", resp.DeviceKeys[carol])
	}

	// A stranger is not cached: every query for them is a key query.
	for range 2 {
		_, _ = h.keys.QueryKeys(t.Context(), alice, keyserver.QueryRequest{DeviceKeys: map[string][]string{dave: nil}})
	}
	if fetches, queries := remote.counts(); fetches != 1 || queries != 2 {
		t.Errorf("after two queries for a stranger: %d fetches, %d key queries; want 1 and 2", fetches, queries)
	}
	if _, ok, _ := h.store.RemoteKeys(t.Context(), dave); ok {
		t.Error("a stranger's keys were cached")
	}
}

func TestRemoteDeviceListUpdatesKeepTheCacheCurrent(t *testing.T) {
	ctx := t.Context()
	remote := &snapshotRemote{}
	remote.set(5, carolDevice(t, "C1"))
	h := newHarness(t, remote)
	queryCarol(t, h)

	// An update that follows on from the copy is applied to it.
	c2 := carolDevice(t, "C2")
	if err := h.keys.RemoteDeviceListUpdate(ctx, federation.DeviceListUpdateContent{
		UserID: carol, DeviceID: "C2", StreamID: 6, PrevID: []int64{5}, Keys: c2.Keys,
	}); err != nil {
		t.Fatal(err)
	}
	if resp := queryCarol(t, h); resp.DeviceKeys[carol]["C2"] == nil || resp.DeviceKeys[carol]["C1"] == nil {
		t.Errorf("after the update: %v, want C1 and C2", resp.DeviceKeys[carol])
	}
	if fetches, _ := remote.counts(); fetches != 1 {
		t.Errorf("an update that followed on caused %d fetches, want 1", fetches)
	}
	pos, _ := h.store.DeviceListPosition(ctx, carol)
	if pos == 0 {
		t.Error("the update did not move Carol's device list")
	}

	// The same update again changes nothing, not even the stream.
	if err := h.keys.RemoteDeviceListUpdate(ctx, federation.DeviceListUpdateContent{
		UserID: carol, DeviceID: "C2", StreamID: 6, PrevID: []int64{5}, Keys: c2.Keys,
	}); err != nil {
		t.Fatal(err)
	}
	if again, _ := h.store.DeviceListPosition(ctx, carol); again != pos {
		t.Errorf("a repeated update moved the device list from %d to %d", pos, again)
	}

	// A deletion that follows on.
	if err := h.keys.RemoteDeviceListUpdate(ctx, federation.DeviceListUpdateContent{
		UserID: carol, DeviceID: "C1", StreamID: 7, PrevID: []int64{6}, Deleted: true,
	}); err != nil {
		t.Fatal(err)
	}
	if resp := queryCarol(t, h); resp.DeviceKeys[carol]["C1"] != nil {
		t.Errorf("a deleted device is still served: %v", resp.DeviceKeys[carol])
	}

	// An update after a gap marks the copy stale, and it is fetched again
	// in the background.
	remote.set(10, carolDevice(t, "C3"))
	if err := h.keys.RemoteDeviceListUpdate(ctx, federation.DeviceListUpdateContent{
		UserID: carol, DeviceID: "C3", StreamID: 10, PrevID: []int64{9}, Keys: json.RawMessage(`{}`),
	}); err != nil {
		t.Fatal(err)
	}
	h.keys.WaitIdle()
	if fetches, _ := remote.counts(); fetches != 2 {
		t.Errorf("after a gap there were %d fetches, want 2", fetches)
	}
	cached, ok, _ := h.store.RemoteKeys(ctx, carol)
	if !ok || cached.Stale || cached.StreamID != 10 || len(cached.Devices) != 1 || cached.Devices[0].DeviceID != "C3" {
		t.Errorf("after the resync the cache holds %+v", cached)
	}
	resp := queryCarol(t, h)
	if resp.DeviceKeys[carol]["C3"] == nil || len(resp.DeviceKeys[carol]) != 1 {
		t.Errorf("after the resync: %v, want C3 alone", resp.DeviceKeys[carol])
	}
	if fetches, _ := remote.counts(); fetches != 2 {
		t.Errorf("a query after the resync fetched again: %d fetches", fetches)
	}
}

func TestAStaleCopyIsNotServed(t *testing.T) {
	ctx := t.Context()
	remote := &snapshotRemote{}
	remote.set(5, carolDevice(t, "C1"))
	h := newHarness(t, remote)
	queryCarol(t, h)

	// With Carol's server down, the background fetch fails and the copy
	// stays stale; a query then asks again rather than serve it.
	remote.mu.Lock()
	remote.down = true
	remote.mu.Unlock()
	if err := h.keys.RemoteDeviceListUpdate(ctx, federation.DeviceListUpdateContent{
		UserID: carol, DeviceID: "C2", StreamID: 8, PrevID: []int64{7},
	}); err != nil {
		t.Fatal(err)
	}
	h.keys.WaitIdle()
	resp := queryCarol(t, h)
	if failure, ok := resp.Failures["remote.example"].(map[string]any); !ok || failure["status"] != 502 {
		t.Errorf("failures = %v, want a 502 for remote.example", resp.Failures)
	}
	if fetches, _ := remote.counts(); fetches != 3 {
		t.Errorf("%d fetches, want the first, the background one and the query's", fetches)
	}
}

func TestRemoteSigningKeyUpdatesReplaceTheCrossSigningKeys(t *testing.T) {
	ctx := t.Context()
	remote := &snapshotRemote{}
	oldMaster, newMaster, selfSigning := newKey(t), newKey(t), newKey(t)
	remote.snapshot.MasterKey = signed(t, carol, oldMaster, crossSigningKey(carol, "master", oldMaster))
	remote.snapshot.SelfSigningKey = signed(t, carol, oldMaster, crossSigningKey(carol, "self_signing", selfSigning))
	remote.set(1, carolDevice(t, "C1"))
	h := newHarness(t, remote)
	queryCarol(t, h)

	if err := h.keys.RemoteSigningKeyUpdate(ctx, federation.SigningKeyUpdateContent{
		UserID:    carol,
		MasterKey: signed(t, carol, newMaster, crossSigningKey(carol, "master", newMaster)),
	}); err != nil {
		t.Fatal(err)
	}
	resp := queryCarol(t, h)
	var master struct {
		Keys map[string]string `json:"keys"`
	}
	_ = json.Unmarshal(resp.MasterKeys[carol], &master)
	if master.Keys["ed25519:"+newMaster.PublicBase64()] == "" {
		t.Errorf("master key after the update = %s, want the new one", resp.MasterKeys[carol])
	}
	if resp.SelfSigningKeys[carol] != nil {
		t.Errorf("a self-signing key the update left out is still served: %s", resp.SelfSigningKeys[carol])
	}
	if fetches, _ := remote.counts(); fetches != 1 {
		t.Errorf("a signing key update caused %d fetches, want 1", fetches)
	}
	if !h.notifier.wasWoken(alice) {
		t.Error("alice was not told Carol's keys changed")
	}
}

func TestCachedKeysOfStrangersAreDropped(t *testing.T) {
	ctx := t.Context()
	remote := &snapshotRemote{}
	remote.set(1, carolDevice(t, "C1"))
	h := newHarness(t, remote)
	queryCarol(t, h)

	// Alice leaving a room is not enough while Carol still shares another
	// one with her.
	h.keys.MembershipLeft(alice)
	h.keys.WaitIdle()
	if _, ok, _ := h.store.RemoteKeys(ctx, carol); !ok {
		t.Fatal("the cache was dropped while Carol still shares a room with alice")
	}

	h.rooms.set(carol, carol)
	h.keys.MembershipLeft(carol)
	h.keys.WaitIdle()
	if _, ok, _ := h.store.RemoteKeys(ctx, carol); ok {
		t.Error("Carol's keys are still cached after she stopped sharing a room with anybody here")
	}
	if users, _ := h.store.RemoteKeyUsers(ctx); len(users) != 0 {
		t.Errorf("users with cached keys = %v", users)
	}
}

func TestCachedKeysAreDroppedWhenALocalUserLeaves(t *testing.T) {
	ctx := t.Context()
	remote := &snapshotRemote{}
	remote.set(1, carolDevice(t, "C1"))
	h := newHarness(t, remote)
	queryCarol(t, h)

	// Alice left the room she shared with Carol: everybody cached is
	// checked, since it is not known who else was in it.
	h.rooms.set(carol, carol)
	h.keys.MembershipLeft(alice)
	h.keys.WaitIdle()
	if _, ok, _ := h.store.RemoteKeys(ctx, carol); ok {
		t.Error("Carol's keys are still cached after the last local user left")
	}

	// And a query for her now passes through.
	queryCarol(t, h)
	if fetches, queries := remote.counts(); fetches != 1 || queries != 1 {
		t.Errorf("%d fetches and %d key queries, want 1 and 1", fetches, queries)
	}
	if _, ok, _ := h.store.RemoteKeys(ctx, carol); ok {
		t.Error("a stranger's keys were cached")
	}
}