package setup_test

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"sync/atomic"
	"testing"

	"github.com/dariuszjaniak/warp/server/internal/matrix/signing"
	"github.com/dariuszjaniak/warp/server/internal/setup"
)

// The Phase 7 tests: end-to-end encryption's server half, through the
// whole stack.

// sendJSON sends a JSON body with any method.
func sendJSON(t *testing.T, srv *setup.Server, method, path string, body any, token string) *httptest.ResponseRecorder {
	t.Helper()
	raw, err := json.Marshal(body)
	if err != nil {
		t.Fatal(err)
	}
	return send(t, srv, method, path, "application/json", raw, token, nil)
}

func mustOK(t *testing.T, rec *httptest.ResponseRecorder, what string) map[string]any {
	t.Helper()
	if rec.Code != http.StatusOK {
		t.Fatalf("%s: status %d: %s", what, rec.Code, rec.Body.String())
	}
	return decode[map[string]any](t, rec)
}

func deviceIDOf(t *testing.T, srv *setup.Server, token string) string {
	t.Helper()
	whoami := mustOK(t, getWithToken(t, srv, "/_matrix/client/v3/account/whoami", token), "whoami")
	deviceID, _ := whoami["device_id"].(string)
	return deviceID
}

// signedObject signs a JSON object as userID with an ed25519 key named
// keyName.
func signedObject(t *testing.T, userID, keyName string, priv ed25519.PrivateKey, obj map[string]any) map[string]any {
	t.Helper()
	raw, _ := json.Marshal(obj)
	out, err := signing.SignJSON(userID, &signing.KeyPair{ID: keyName, Private: priv}, raw)
	if err != nil {
		t.Fatal(err)
	}
	var signed map[string]any
	_ = json.Unmarshal(out, &signed)
	return signed
}

func newEd25519(t *testing.T) (string, ed25519.PrivateKey) {
	t.Helper()
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	return base64.RawStdEncoding.EncodeToString(pub), priv
}

// uploadDeviceKeys uploads a device's identity keys and some one-time
// keys, and returns the response.
func uploadDeviceKeys(t *testing.T, srv *setup.Server, token, userID, deviceID string, otks int) map[string]any {
	t.Helper()
	pub, priv := newEd25519(t)
	deviceKeys := signedObject(t, userID, deviceID, priv, map[string]any{
		"user_id":    userID,
		"device_id":  deviceID,
		"algorithms": []string{"m.olm.v1.curve25519-aes-sha2", "m.megolm.v1.aes-sha2"},
		"keys": map[string]string{
			"ed25519:" + deviceID:    pub,
			"curve25519:" + deviceID: "curve" + deviceID,
		},
	})
	oneTime := map[string]any{}
	for i := range otks {
		oneTime["signed_curve25519:"+deviceID+string(rune('a'+i))] = map[string]any{"key": "otk" + string(rune('a'+i))}
	}
	return mustOK(t, post(t, srv, "/_matrix/client/v3/keys/upload", map[string]any{
		"device_keys":   deviceKeys,
		"one_time_keys": oneTime,
		"fallback_keys": map[string]any{"signed_curve25519:fb" + deviceID: map[string]any{"key": "fb", "fallback": true}},
	}, token), "keys/upload")
}

// syncResponse is the part of a sync these tests read.
type syncResponse struct {
	NextBatch   string `json:"next_batch"`
	DeviceLists struct {
		Changed []string `json:"changed"`
		Left    []string `json:"left"`
	} `json:"device_lists"`
	OTKCounts      map[string]int `json:"device_one_time_keys_count"`
	UnusedFallback *[]string      `json:"device_unused_fallback_key_types"`
}

func syncOnce(t *testing.T, srv *setup.Server, token, since string) syncResponse {
	t.Helper()
	path := "/_matrix/client/v3/sync?timeout=0"
	if since != "" {
		path += "&since=" + url.QueryEscape(since)
	}
	rec := getWithToken(t, srv, path, token)
	if rec.Code != http.StatusOK {
		t.Fatalf("sync: %d %s", rec.Code, rec.Body.String())
	}
	return decode[syncResponse](t, rec)
}

func createPublicRoom(t *testing.T, srv *setup.Server, token string) string {
	t.Helper()
	resp := mustOK(t, post(t, srv, "/_matrix/client/v3/createRoom", map[string]any{"preset": "public_chat"}, token), "createRoom")
	roomID, _ := resp["room_id"].(string)
	return roomID
}

func TestEndToEndEncryptionThroughTheAPI(t *testing.T) {
	srv := newSigningServer(t, nil)
	aliceToken := registerThrough(t, srv, "alice")
	bobToken := registerThrough(t, srv, "bob")
	const alice, bob = "@alice:warp.example", "@bob:warp.example"
	aliceDevice := deviceIDOf(t, srv, aliceToken)

	roomID := createPublicRoom(t, srv, aliceToken)
	mustOK(t, post(t, srv, "/_matrix/client/v3/rooms/"+url.PathEscape(roomID)+"/join", map[string]any{}, bobToken), "join")
	bobSince := syncOnce(t, srv, bobToken, "").NextBatch

	// Alice uploads her keys, and is told how many one-time keys she has.
	counts, _ := uploadDeviceKeys(t, srv, aliceToken, alice, aliceDevice, 3)["one_time_key_counts"].(map[string]any)
	if counts["signed_curve25519"] != float64(3) {
		t.Errorf("one_time_key_counts = %v", counts)
	}

	// Bob's next sync says Alice's keys changed; Alice's carries her key
	// counts, and both carry the fallback key list even when it is empty.
	bobSync := syncOnce(t, srv, bobToken, bobSince)
	if !contains(bobSync.DeviceLists.Changed, alice) {
		t.Errorf("bob's device_lists.changed = %v, want alice", bobSync.DeviceLists.Changed)
	}
	if bobSync.UnusedFallback == nil || len(*bobSync.UnusedFallback) != 0 {
		t.Errorf("bob's device_unused_fallback_key_types = %v, want []", bobSync.UnusedFallback)
	}
	if bobSync.OTKCounts["signed_curve25519"] != 0 {
		t.Errorf("bob's one-time key count = %v", bobSync.OTKCounts)
	}
	aliceSync := syncOnce(t, srv, aliceToken, "")
	if aliceSync.OTKCounts["signed_curve25519"] != 3 || aliceSync.UnusedFallback == nil || len(*aliceSync.UnusedFallback) != 1 {
		t.Errorf("alice's key counts = %v, fallback %v", aliceSync.OTKCounts, aliceSync.UnusedFallback)
	}

	changes := mustOK(t, getWithToken(t, srv, "/_matrix/client/v3/keys/changes?from="+url.QueryEscape(bobSince), bobToken), "keys/changes")
	if changed, _ := changes["changed"].([]any); len(changed) != 1 || changed[0] != alice {
		t.Errorf("keys/changes = %v", changes)
	}

	// Bob fetches Alice's keys and claims one of her one-time keys.
	query := mustOK(t, post(t, srv, "/_matrix/client/v3/keys/query", map[string]any{
		"device_keys": map[string]any{alice: []string{}},
	}, bobToken), "keys/query")
	devices, _ := query["device_keys"].(map[string]any)[alice].(map[string]any)
	if _, ok := devices[aliceDevice]; !ok {
		t.Fatalf("keys/query = %v", query)
	}
	claim := mustOK(t, post(t, srv, "/_matrix/client/v3/keys/claim", map[string]any{
		"one_time_keys": map[string]any{alice: map[string]string{aliceDevice: "signed_curve25519"}},
	}, bobToken), "keys/claim")
	if keys, _ := claim["one_time_keys"].(map[string]any)[alice].(map[string]any)[aliceDevice].(map[string]any); len(keys) != 1 {
		t.Errorf("keys/claim = %v", claim)
	}

	// Cross-signing: the first upload needs no password, replacing the
	// master key does.
	masterPub, masterPriv := newEd25519(t)
	sskPub, _ := newEd25519(t)
	crossSigning := func(master string, priv ed25519.PrivateKey, ssk string) map[string]any {
		return map[string]any{
			"master_key": map[string]any{"user_id": alice, "usage": []string{"master"},
				"keys": map[string]string{"ed25519:" + master: master}},
			"self_signing_key": signedObject(t, alice, master, priv, map[string]any{
				"user_id": alice, "usage": []string{"self_signing"},
				"keys": map[string]string{"ed25519:" + ssk: ssk}}),
		}
	}
	mustOK(t, post(t, srv, "/_matrix/client/v3/keys/device_signing/upload", crossSigning(masterPub, masterPriv, sskPub), aliceToken),
		"first device_signing/upload")

	newMaster, newMasterPriv := newEd25519(t)
	body := crossSigning(newMaster, newMasterPriv, sskPub)
	rec := post(t, srv, "/_matrix/client/v3/keys/device_signing/upload", body, aliceToken)
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("replacing the master key without auth = %d, want 401: %s", rec.Code, rec.Body.String())
	}
	body["auth"] = map[string]any{
		"type": "m.login.password", "session": decode[map[string]any](t, rec)["session"],
		"identifier": map[string]string{"type": "m.id.user", "user": "alice"}, "password": "sUp3rs3kr1t",
	}
	mustOK(t, post(t, srv, "/_matrix/client/v3/keys/device_signing/upload", body, aliceToken), "device_signing/upload with auth")
	query = mustOK(t, post(t, srv, "/_matrix/client/v3/keys/query", map[string]any{
		"device_keys": map[string]any{alice: []string{}},
	}, bobToken), "keys/query after cross-signing")
	if master, _ := query["master_keys"].(map[string]any)[alice].(map[string]any); master == nil ||
		master["keys"].(map[string]any)["ed25519:"+newMaster] != newMaster {
		t.Errorf("master_keys = %v, want the replacement", query["master_keys"])
	}

	// Key backup, version and keys.
	version := mustOK(t, post(t, srv, "/_matrix/client/v3/room_keys/version", map[string]any{
		"algorithm": "m.megolm_backup.v1.curve25519-aes-sha2", "auth_data": map[string]any{"public_key": "abc"},
	}, aliceToken), "create backup")["version"].(string)
	put := mustOK(t, sendJSON(t, srv, http.MethodPut, "/_matrix/client/v3/room_keys/keys/"+url.PathEscape(roomID)+"/sess1?version="+version,
		map[string]any{"first_message_index": 0, "forwarded_count": 0, "is_verified": false,
			"session_data": map[string]any{"ciphertext": "c"}}, aliceToken), "put room key")
	if put["count"] != float64(1) {
		t.Errorf("put room key = %v", put)
	}
	backup := mustOK(t, getWithToken(t, srv, "/_matrix/client/v3/room_keys/version", aliceToken), "get backup")
	if backup["version"] != version || backup["count"] != float64(1) {
		t.Errorf("backup = %v", backup)
	}
	keys := mustOK(t, getWithToken(t, srv, "/_matrix/client/v3/room_keys/keys?version="+version, aliceToken), "get room keys")
	if rooms, _ := keys["rooms"].(map[string]any); rooms[roomID] == nil {
		t.Errorf("room keys = %v", keys)
	}
	rec = sendJSON(t, srv, http.MethodPut, "/_matrix/client/v3/room_keys/keys?version=999", map[string]any{"rooms": map[string]any{}}, aliceToken)
	if rec.Code != http.StatusNotFound {
		t.Errorf("writing to an unknown backup version = %d: %s", rec.Code, rec.Body.String())
	}
	if rec := getWithToken(t, srv, "/_matrix/client/v3/room_keys/version", bobToken); rec.Code != http.StatusNotFound {
		t.Errorf("a user with no backup = %d, want 404", rec.Code)
	}

	// Logging in a second device, uploading its keys and deleting it: the
	// keys go with the device, and Bob is told.
	login := mustOK(t, post(t, srv, "/_matrix/client/v3/login", map[string]any{
		"type": "m.login.password", "identifier": map[string]string{"type": "m.id.user", "user": "alice"},
		"password": "sUp3rs3kr1t", "device_id": "PHONE",
	}, ""), "login")
	uploadDeviceKeys(t, srv, login["access_token"].(string), alice, "PHONE", 1)
	bobSince = syncOnce(t, srv, bobToken, "").NextBatch

	rec = sendJSON(t, srv, http.MethodDelete, "/_matrix/client/v3/devices/PHONE", map[string]any{}, aliceToken)
	if rec.Code != http.StatusUnauthorized {
		t.Fatalf("delete device = %d, want 401", rec.Code)
	}
	mustOK(t, sendJSON(t, srv, http.MethodDelete, "/_matrix/client/v3/devices/PHONE", map[string]any{"auth": map[string]any{
		"type": "m.login.password", "session": decode[map[string]any](t, rec)["session"],
		"identifier": map[string]string{"type": "m.id.user", "user": "alice"}, "password": "sUp3rs3kr1t",
	}}, aliceToken), "delete device with auth")

	if got := syncOnce(t, srv, bobToken, bobSince); !contains(got.DeviceLists.Changed, alice) {
		t.Errorf("after the device was deleted, bob's changed = %v", got.DeviceLists.Changed)
	}
	query = mustOK(t, post(t, srv, "/_matrix/client/v3/keys/query", map[string]any{
		"device_keys": map[string]any{alice: []string{}},
	}, bobToken), "keys/query after deletion")
	if _, ok := query["device_keys"].(map[string]any)[alice].(map[string]any)["PHONE"]; ok {
		t.Error("a deleted device's keys are still served")
	}

	// Bob leaving the only room they share puts him in Alice's left list.
	aliceSince := syncOnce(t, srv, aliceToken, "").NextBatch
	mustOK(t, post(t, srv, "/_matrix/client/v3/rooms/"+url.PathEscape(roomID)+"/leave", map[string]any{}, bobToken), "leave")
	if got := syncOnce(t, srv, aliceToken, aliceSince); !contains(got.DeviceLists.Left, bob) {
		t.Errorf("after bob left, alice's device_lists = %+v", got.DeviceLists)
	}
}

func TestKeysOverFederation(t *testing.T) {
	if testing.Short() {
		t.Skip("starts two servers")
	}
	nodes := newFederation(t, "hs-a.example", "hs-b.example")
	a, b := nodes["hs-a.example"].srv, nodes["hs-b.example"].srv
	const alice, bob = "@alice:hs-a.example", "@bob:hs-b.example"

	// Count the key requests that reach A, which is how the test tells a
	// query B answered from its cache from one it passed on.
	var keyRequests atomic.Int64
	inner := *nodes["hs-a.example"].handler.Load()
	counting := http.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if strings.HasPrefix(r.URL.Path, "/_matrix/federation/v1/user/devices/") ||
			r.URL.Path == "/_matrix/federation/v1/user/keys/query" {
			keyRequests.Add(1)
		}
		inner.ServeHTTP(w, r)
	}))
	nodes["hs-a.example"].handler.Store(&counting)

	aliceToken := registerThrough(t, a, "alice")
	bobToken := registerThrough(t, b, "bob")
	aliceDevice := deviceIDOf(t, a, aliceToken)

	roomID := createPublicRoom(t, a, aliceToken)
	mustOK(t, post(t, b, "/_matrix/client/v3/join/"+url.PathEscape(roomID)+"?server_name=hs-a.example",
		map[string]any{}, bobToken), "remote join")
	bobSince := syncOnce(t, b, bobToken, "").NextBatch

	// Alice's upload reaches Bob's server as an m.device_list_update, and
	// Bob's next sync says so.
	uploadDeviceKeys(t, a, aliceToken, alice, aliceDevice, 2)
	eventually(t, "alice in bob's device_lists.changed", func() bool {
		got := syncOnce(t, b, bobToken, bobSince)
		return contains(got.DeviceLists.Changed, alice)
	})

	// Bob's first query is answered by Alice's server, and cached: the
	// second does not leave B.
	queryAlice := func(what string) map[string]any {
		t.Helper()
		query := mustOK(t, post(t, b, "/_matrix/client/v3/keys/query", map[string]any{
			"device_keys": map[string]any{alice: []string{}},
		}, bobToken), what)
		devices, _ := query["device_keys"].(map[string]any)[alice].(map[string]any)
		return devices
	}
	if devices := queryAlice("remote keys/query"); devices[aliceDevice] == nil {
		t.Fatalf("remote keys/query = %v", devices)
	}
	if n := keyRequests.Load(); n != 1 {
		t.Errorf("the first query made %d key requests to A, want 1", n)
	}
	if devices := queryAlice("second remote keys/query"); devices[aliceDevice] == nil {
		t.Fatalf("second remote keys/query = %v", devices)
	}
	if n := keyRequests.Load(); n != 1 {
		t.Errorf("the second query reached A: %d key requests, want 1", n)
	}

	// A new device of Alice's reaches B as an m.device_list_update that
	// follows on from the cached copy, which B applies: Bob sees the new
	// device, and A is still not asked.
	login := mustOK(t, post(t, a, "/_matrix/client/v3/login", map[string]any{
		"type": "m.login.password", "identifier": map[string]string{"type": "m.id.user", "user": "alice"},
		"password": "sUp3rs3kr1t", "device_id": "PHONE",
	}, ""), "login")
	uploadDeviceKeys(t, a, login["access_token"].(string), alice, "PHONE", 1)
	eventually(t, "alice's new device in bob's query", func() bool {
		devices := queryAlice("keys/query after the update")
		return devices["PHONE"] != nil && devices[aliceDevice] != nil
	})
	if n := keyRequests.Load(); n != 1 {
		t.Errorf("the update was not applied from the EDU: %d key requests to A, want 1", n)
	}
	claim := mustOK(t, post(t, b, "/_matrix/client/v3/keys/claim", map[string]any{
		"one_time_keys": map[string]any{alice: map[string]string{aliceDevice: "signed_curve25519"}},
	}, bobToken), "remote keys/claim")
	if keys, _ := claim["one_time_keys"].(map[string]any)[alice].(map[string]any)[aliceDevice].(map[string]any); len(keys) != 1 {
		t.Errorf("remote keys/claim = %v", claim)
	}
	counts := syncOnce(t, a, aliceToken, "").OTKCounts
	if counts["signed_curve25519"] != 1 {
		t.Errorf("after a remote claim alice has %v one-time keys, want 1", counts)
	}

	// With A's federation off, B reports A as a failure instead of
	// hanging.
	if err := a.Settings().SetFederationEnabled(t.Context(), false); err != nil {
		t.Fatal(err)
	}
	// The cached copy is still B's to serve.
	if devices := queryAlice("keys/query with A's federation off"); devices["PHONE"] == nil {
		t.Errorf("with A's federation off, B no longer serves its cached copy: %v", devices)
	}
	// A user B has nothing cached for is asked of A, and A is reported as
	// a failure instead of hanging.
	query := mustOK(t, post(t, b, "/_matrix/client/v3/keys/query", map[string]any{
		"device_keys": map[string]any{"@stranger:hs-a.example": []string{}}, "timeout": 5000,
	}, bobToken), "keys/query for a stranger with A's federation off")
	if failures, _ := query["failures"].(map[string]any); failures["hs-a.example"] == nil {
		t.Errorf("failures = %v, want hs-a.example", query["failures"])
	}
}