package keyserver

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"slices"
	"sync"

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

// The cache of other servers' users' keys.
//
// A local client asks for the keys of everybody in a room before it sends
// an encrypted message, and again whenever a sync says one of them moved.
// Without a cache each of those is a federation round trip, in front of
// the message, to a server that may be slow or down. With one, a remote
// user's keys are fetched once — the whole device list, through GET
// /user/devices/{userId}, which also says where in their server's device
// list stream that list stands — and then kept current by the updates
// their server pushes: an m.device_list_update per device change, naming
// the stream position it follows on from, and an m.signing_key_update per
// change of their cross-signing keys.
//
// Only users who share a room with somebody here are cached, because only
// for them is there anybody to send updates here: a server tells the
// servers its user shares rooms with, and nobody else. A copy of a
// stranger's keys would go stale without anybody knowing. So a query
// about a stranger is passed on to their server as before, and a cached
// user who stops sharing a room with anybody here is dropped.
//
// An update that does not follow on from the copy held — one was lost, or
// arrived out of order — marks the copy stale, and it is fetched again in
// the background. Until it is, a query does what it does on a miss: asks
// their server, synchronously. Serving a copy known to be missing a change
// would hand a sender a device list without the device just added, and
// the message would go out undecryptable on it.

// cacheable reports whether a remote user's keys are kept here: whether a
// local user shares a room with them. An error reading the rooms is "no",
// which costs a federation round trip and nothing else.
func (k *API) cacheable(ctx context.Context, userID string) bool {
	if k.rooms == nil || k.isLocal(userID) {
		return false
	}
	shares, err := k.sharesWithLocal(ctx, userID)
	if err != nil {
		k.log.Warn("could not read who shares a room with a remote user", "user_id", userID, "error", err)
		return false
	}
	return shares
}

// sharesWithLocal reports whether any local user shares a room with
// userID.
func (k *API) sharesWithLocal(ctx context.Context, userID string) (bool, error) {
	if k.rooms == nil {
		return false, nil
	}
	users, err := k.rooms.UsersSharingRooms(ctx, userID)
	if err != nil {
		return false, err
	}
	return slices.ContainsFunc(users, k.isLocal), nil
}

// forget drops what is cached for users. A failure only leaves a copy
// behind that the next sweep takes.
func (k *API) forget(ctx context.Context, userIDs ...string) {
	if err := k.store.DeleteRemoteKeys(ctx, userIDs); err != nil {
		k.log.Warn("could not drop cached keys", "user_ids", userIDs, "error", err)
	}
}

// freshCopy reads the cached keys of a remote user, if there is a copy
// that can be served.
func (k *API) freshCopy(ctx context.Context, userID string) (storage.RemoteKeys, bool) {
	keys, ok, err := k.store.RemoteKeys(ctx, userID)
	if err != nil {
		k.log.Warn("could not read cached keys", "user_id", userID, "error", err)
		return storage.RemoteKeys{}, false
	}
	return keys, ok && !keys.Stale
}

// fetchRemoteKeys asks a user's server for their whole device list and
// cross-signing keys, and caches the answer.
//
// What the server sent is checked as far as the server can check it: a
// device whose keys name another user or device is left out, and so is a
// cross-signing key that is not that user's key for that purpose. The
// signatures are the client's to check — a server that verified them could
// still not tell a device its user added from one their server invented.
func (k *API) fetchRemoteKeys(ctx context.Context, server, userID string) (storage.RemoteKeys, error) {
	if k.remote == nil {
		return storage.RemoteKeys{}, errors.New("federation is not available on this server")
	}
	answer, err := k.remote.UserDevices(ctx, server, userID)
	if err != nil {
		return storage.RemoteKeys{}, err
	}
	if answer.UserID != userID {
		return storage.RemoteKeys{}, fmt.Errorf("%s answered for %q when asked about %q", server, answer.UserID, userID)
	}

	keys := storage.RemoteKeys{UserID: userID, StreamID: answer.StreamID, CrossSigning: map[string][]byte{}}
	for _, d := range answer.Devices {
		if !deviceKeysOf(d.Keys, userID, d.DeviceID) {
			continue
		}
		keys.Devices = append(keys.Devices, storage.RemoteDevice{
			DeviceID: d.DeviceID, DisplayName: d.DeviceDisplayName, JSON: d.Keys,
		})
	}
	setCrossSigning(keys.CrossSigning, userID, storage.CrossSigningMaster, answer.MasterKey)
	setCrossSigning(keys.CrossSigning, userID, storage.CrossSigningSelfSigning, answer.SelfSigningKey)

	// Still a stranger's keys if the last shared room went while the
	// request was out; they are served this once and not kept.
	if !k.cacheable(ctx, userID) {
		return keys, nil
	}
	if _, err := k.store.SetRemoteKeys(ctx, keys); err != nil {
		k.log.Warn("could not cache keys", "user_id", userID, "error", err)
	}
	return keys, nil
}

// deviceKeysOf reports whether a device_keys object is the one of userID's
// device deviceID.
func deviceKeysOf(raw json.RawMessage, userID, deviceID string) bool {
	var keys deviceKeysObject
	return json.Unmarshal(raw, &keys) == nil && keys.UserID == userID && keys.DeviceID == deviceID && deviceID != ""
}

// setCrossSigning adds a remote user's cross-signing key to keys when it
// is one.
func setCrossSigning(keys map[string][]byte, userID, purpose string, raw json.RawMessage) {
	if len(raw) == 0 || string(raw) == "null" {
		return
	}
	if encoded, _, err := checkCrossSigningKey(userID, purpose, raw); err == nil {
		keys[purpose] = encoded
	}
}

// serveRemote writes a cached copy into a query response, limited to the
// devices asked for.
func (k *API) serveRemote(ctx context.Context, requester string, keys storage.RemoteKeys, deviceIDs []string, resp *QueryResponse) {
	devices := map[string]json.RawMessage{}
	for _, d := range keys.Devices {
		if len(deviceIDs) > 0 && !slices.Contains(deviceIDs, d.DeviceID) {
			continue
		}
		raw, err := withSignatures(d.JSON, nil, d.DisplayName)
		if err != nil {
			raw = d.JSON
		}
		devices[d.DeviceID] = raw
	}
	resp.DeviceKeys[keys.UserID] = devices
	if master, ok := keys.CrossSigning[storage.CrossSigningMaster]; ok {
		resp.MasterKeys[keys.UserID] = k.withRequesterSignature(ctx, keys.UserID, requester, master)
	}
	if selfSigning, ok := keys.CrossSigning[storage.CrossSigningSelfSigning]; ok {
		resp.SelfSigningKeys[keys.UserID] = selfSigning
	}
}

// --- updates from other servers -------------------------------------------------

// RemoteDeviceListUpdate takes in an m.device_list_update from another
// server: one of its users' devices was added, changed or deleted.
//
// The update is applied to the cached copy when it follows on from it,
// and otherwise marks the copy stale and has it fetched again. Either way
// the user's position in the device list stream moves, which is what
// tells local clients — who keep their own copies — to ask again.
//
// Nothing happens unless a local user shares a room with the user: an
// update about a stranger concerns nobody here, and recording it would let
// any server grow the tables with users of its own invention. A cached
// user who has become a stranger is dropped.
func (k *API) RemoteDeviceListUpdate(ctx context.Context, u federation.DeviceListUpdateContent) error {
	if k.isLocal(u.UserID) || k.rooms == nil {
		return nil
	}
	shares, err := k.sharesWithLocal(ctx, u.UserID)
	if err != nil {
		return err
	}
	if !shares {
		k.forget(ctx, u.UserID)
		return nil
	}

	update := storage.RemoteDeviceUpdate{
		UserID: u.UserID, DeviceID: u.DeviceID, DisplayName: u.DeviceDisplayName,
		Deleted: u.Deleted, StreamID: u.StreamID, PrevIDs: u.PrevID,
	}
	// Keys that are not the device's own are as good as none: the update
	// cannot be applied, and the copy is fetched again.
	if !u.Deleted && deviceKeysOf(u.Keys, u.UserID, u.DeviceID) {
		update.JSON = u.Keys
	}
	result, err := k.store.UpdateRemoteDevice(ctx, update)
	if err != nil {
		return err
	}
	switch result {
	case storage.RemoteUpdateDuplicate:
		return nil
	case storage.RemoteUpdateStale:
		k.resyncs.add(u.UserID)
	}
	return k.remoteChanged(ctx, u.UserID)
}

// RemoteSigningKeyUpdate takes in an m.signing_key_update: one of another
// server's users replaced their master or self-signing key. The update
// carries the keys as they now stand, and replaces the cached ones; a key
// it leaves out no longer exists.
func (k *API) RemoteSigningKeyUpdate(ctx context.Context, u federation.SigningKeyUpdateContent) error {
	if k.isLocal(u.UserID) || k.rooms == nil {
		return nil
	}
	shares, err := k.sharesWithLocal(ctx, u.UserID)
	if err != nil {
		return err
	}
	if !shares {
		k.forget(ctx, u.UserID)
		return nil
	}

	keys := map[string][]byte{}
	setCrossSigning(keys, u.UserID, storage.CrossSigningMaster, u.MasterKey)
	setCrossSigning(keys, u.UserID, storage.CrossSigningSelfSigning, u.SelfSigningKey)
	if _, err := k.store.SetRemoteCrossSigningKeys(ctx, u.UserID, keys); err != nil {
		return err
	}
	return k.remoteChanged(ctx, u.UserID)
}

// remoteChanged moves a remote user's position in the device list stream
// and wakes the local syncs that share a room with them.
func (k *API) remoteChanged(ctx context.Context, userID string) error {
	if _, err := k.store.MarkDeviceListsChanged(ctx, []string{userID}); err != nil {
		return err
	}
	k.changed(ctx, userID)
	return nil
}

// resync fetches a stale copy again. It runs in the background: the
// update that found the copy stale arrived in a transaction from the
// user's server, which should not wait on a request back to itself.
func (k *API) resync(ctx context.Context, userID string) {
	if !k.cacheable(ctx, userID) {
		k.forget(ctx, userID)
		return
	}
	server := serverOf(userID)
	if server == "" {
		return
	}
	ctx, cancel := context.WithTimeout(ctx, defaultRemoteTimeout)
	defer cancel()
	if _, err := k.fetchRemoteKeys(ctx, server, userID); err != nil {
		// The copy stays stale, and the next query asks again.
		k.log.Debug("could not fetch stale keys again", "user_id", userID, "error", err)
	}
}

// --- dropping strangers ---------------------------------------------------------

// MembershipLeft is told that a user's membership of a room ended: they
// left, were kicked or were banned. If the user is remote, they may no
// longer share a room with anybody here, and their cached keys go; if the
// user is local, any of the remote users they shared that room with may
// have stopped sharing one with anybody, and every cached user is
// checked.
//
// It does not block. The check reads room memberships, and is called from
// the roomserver's emit stage with the room's lock held.
func (k *API) MembershipLeft(userID string) {
	if k.isLocal(userID) {
		k.sweeps.add(sweepEveryone)
		return
	}
	k.sweeps.add(userID)
}

// sweepEveryone is the sweep queue's key for "check every cached user".
const sweepEveryone = ""

// sweep drops the cached keys of users nobody here shares a room with any
// more: one user, or with sweepEveryone all of them.
func (k *API) sweep(ctx context.Context, userID string) {
	users := []string{userID}
	if userID == sweepEveryone {
		var err error
		if users, err = k.store.RemoteKeyUsers(ctx); err != nil {
			k.log.Warn("could not list the users with cached keys", "error", err)
			return
		}
	}
	var strangers []string
	for _, u := range users {
		shares, err := k.sharesWithLocal(ctx, u)
		if err != nil {
			k.log.Warn("could not read who shares a room with a remote user", "user_id", u, "error", err)
			continue
		}
		if !shares {
			strangers = append(strangers, u)
		}
	}
	if len(strangers) > 0 {
		k.forget(ctx, strangers...)
	}
}

// --- background work ------------------------------------------------------------

// workQueue runs a job per key in one background goroutine, one key at a
// time, and folds a key added while it is already waiting into the wait.
// A burst of updates about one user is one resync, and a burst of leaves
// one sweep.
type workQueue struct {
	k   *API
	run func(ctx context.Context, key string)

	mu      sync.Mutex
	pending []string
	running bool
}

func newWorkQueue(k *API, run func(context.Context, string)) *workQueue {
	return &workQueue{k: k, run: run}
}

// add queues a key, and starts the queue's goroutine if it is not running.
// After Close it does nothing.
func (q *workQueue) add(key string) {
	q.mu.Lock()
	defer q.mu.Unlock()
	if !slices.Contains(q.pending, key) {
		q.pending = append(q.pending, key)
	}
	if q.running || !q.k.background.start() {
		return
	}
	q.running = true
	go func() {
		defer q.k.background.done()
		for {
			q.mu.Lock()
			if len(q.pending) == 0 {
				q.running = false
				q.mu.Unlock()
				return
			}
			key := q.pending[0]
			q.pending = q.pending[1:]
			q.mu.Unlock()

			q.run(q.k.background.ctx, key)
		}
	}()
}

// background is the lifetime of the keyserver's goroutines.
type background struct {
	ctx    context.Context
	cancel context.CancelFunc

	mu     sync.Mutex
	closed bool
	wg     sync.WaitGroup
}

func newBackground() *background {
	ctx, cancel := context.WithCancel(context.Background())
	return &background{ctx: ctx, cancel: cancel}
}

// start counts a goroutine about to start, unless the keyserver is closed.
func (b *background) start() bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	if b.closed {
		return false
	}
	b.wg.Add(1)
	return true
}

func (b *background) done() { b.wg.Done() }

// wait waits for the goroutines running now to finish.
func (b *background) wait() { b.wg.Wait() }

// close stops the goroutines and waits for them.
func (b *background) close() {
	b.mu.Lock()
	b.closed = true
	b.mu.Unlock()
	b.cancel()
	b.wg.Wait()
}

// Close stops the keyserver's background work — fetching stale copies of
// remote keys again, and dropping those of users nobody here shares a room
// with — and waits for it to stop.
func (k *API) Close() { k.background.close() }