package keyserver

import (
	"cmp"
	"context"
	"encoding/json"
	"slices"

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

// Uploading keys: POST /keys/upload, and what deleting a device does to
// them.

// Algorithm names the server has to know about.
const (
	// AlgorithmSignedCurve25519 is the one-time key algorithm of Olm. A
	// client counts on its count being reported even at zero: it is how
	// the client learns it has to upload more.
	AlgorithmSignedCurve25519 = "signed_curve25519"
)

// maxKeysPerUpload caps one upload's one-time keys. A client uploads up to
// half of what its Olm account holds, which is fifty for libolm and
// vodozemac alike; a thousand is a client that has gone wrong.
const maxKeysPerUpload = 1000

// UploadRequest is the body of POST /keys/upload.
type UploadRequest struct {
	DeviceKeys   json.RawMessage            `json:"device_keys,omitempty"`
	OneTimeKeys  map[string]json.RawMessage `json:"one_time_keys,omitempty"`
	FallbackKeys map[string]json.RawMessage `json:"fallback_keys,omitempty"`

	// UnstableFallbackKeys is MSC2732's name for fallback_keys, which
	// clients older than Matrix 1.2 still send.
	UnstableFallbackKeys map[string]json.RawMessage `json:"org.matrix.msc2732.fallback_keys,omitempty"`
}

// UploadKeys stores what a device uploaded and returns how many one-time
// keys it now has.
//
// The device keys are checked to be the caller's own device, and nothing
// more. Whether they are signed by the key they say they are is checked by
// every client that reads them, which is the check that matters: a server
// that verified them would only be adding a point at which a key could be
// refused and none at which it could be trusted.
func (k *API) UploadKeys(ctx context.Context, userID, deviceID string, req UploadRequest) (map[string]int, error) {
	if len(req.OneTimeKeys) > maxKeysPerUpload {
		return nil, invalidParamf("at most %d one-time keys can be uploaded at once", maxKeysPerUpload)
	}

	if len(req.DeviceKeys) > 0 && string(req.DeviceKeys) != "null" {
		if err := k.uploadDeviceKeys(ctx, userID, deviceID, req.DeviceKeys); err != nil {
			return nil, err
		}
	}

	otks, err := oneTimeKeys(req.OneTimeKeys)
	if err != nil {
		return nil, err
	}
	if err := k.store.AddOneTimeKeys(ctx, userID, deviceID, otks); err != nil {
		return nil, err
	}

	fallback := req.FallbackKeys
	if len(fallback) == 0 {
		fallback = req.UnstableFallbackKeys
	}
	fallbackKeys, err := oneTimeKeys(fallback)
	if err != nil {
		return nil, err
	}
	if err := k.store.SetFallbackKeys(ctx, userID, deviceID, fallbackKeys); err != nil {
		return nil, err
	}

	return k.OneTimeKeyCounts(ctx, userID, deviceID)
}

// uploadDeviceKeys checks and stores a device's identity keys.
func (k *API) uploadDeviceKeys(ctx context.Context, userID, deviceID string, raw json.RawMessage) error {
	var keys deviceKeysObject
	if err := json.Unmarshal(raw, &keys); err != nil {
		return invalidParamf("device_keys is malformed: %v", err)
	}
	switch {
	case keys.UserID != userID:
		return invalidParamf("device_keys.user_id must be %s", userID)
	case keys.DeviceID != deviceID:
		return invalidParamf("device_keys.device_id must be %s, the device the access token belongs to", deviceID)
	case len(keys.Keys) == 0:
		return missingParamf("device_keys.keys must not be empty")
	}

	encoded, err := canonical(raw)
	if err != nil {
		return err
	}
	// A client has no business writing unsigned: the server fills it in
	// on the way out, and nothing in it is covered by the signature.
	if encoded, err = withoutField(encoded, "unsigned"); err != nil {
		return err
	}

	update, err := k.store.SetDeviceKeys(ctx, storage.DeviceKeys{UserID: userID, DeviceID: deviceID, JSON: encoded})
	if err != nil {
		return err
	}
	if update.StreamPos == 0 {
		return nil
	}
	published, err := k.publishedDeviceKeys(ctx, userID, []string{deviceID})
	if err != nil {
		k.log.Warn("could not read back uploaded device keys", "user_id", userID, "device_id", deviceID, "error", err)
	}
	k.announceDevices(ctx, userID, []storage.DeviceListUpdate{update}, published)
	return nil
}

// oneTimeKeys checks and canonicalises a map of one-time or fallback keys.
func oneTimeKeys(in map[string]json.RawMessage) ([]storage.OneTimeKey, error) {
	out := make([]storage.OneTimeKey, 0, len(in))
	for keyID, raw := range in {
		if !validKeyID(keyID) {
			return nil, invalidParamf("key ID %q must have the form algorithm:id", keyID)
		}
		encoded, err := canonical(raw)
		if err != nil {
			return nil, err
		}
		out = append(out, storage.OneTimeKey{KeyID: keyID, JSON: encoded})
	}
	// In order, so that which of two conflicting keys is reported does
	// not depend on map iteration.
	slices.SortFunc(out, func(a, b storage.OneTimeKey) int { return cmp.Compare(a.KeyID, b.KeyID) })
	return out, nil
}

// OneTimeKeyCounts reports a device's unclaimed one-time keys by
// algorithm, with signed_curve25519 always present.
func (k *API) OneTimeKeyCounts(ctx context.Context, userID, deviceID string) (map[string]int, error) {
	counts, err := k.store.OneTimeKeyCounts(ctx, userID, deviceID)
	if err != nil {
		return nil, err
	}
	if _, ok := counts[AlgorithmSignedCurve25519]; !ok {
		counts[AlgorithmSignedCurve25519] = 0
	}
	return counts, nil
}

// UnusedFallbackKeyTypes lists the algorithms a device has an unclaimed
// fallback key for.
func (k *API) UnusedFallbackKeyTypes(ctx context.Context, userID, deviceID string) ([]string, error) {
	types, err := k.store.UnusedFallbackKeyTypes(ctx, userID, deviceID)
	if types == nil && err == nil {
		types = []string{}
	}
	return types, err
}

// DevicesDeleted drops the keys of deleted devices and tells everyone. It
// is the userapi's hook, called after the devices are gone; an empty
// deviceIDs means every device of the user.
//
// It runs after the deletion rather than inside it, so a crash between
// the two leaves keys behind for a device that no longer exists. Those
// keys are never served: a query reads the device list first and skips
// keys for a device that is not on it.
func (k *API) DevicesDeleted(ctx context.Context, localpart string, deviceIDs []string) {
	userID := "@" + localpart + ":" + k.serverName
	updates, err := k.store.DeleteDeviceKeys(ctx, userID, deviceIDs)
	if err != nil {
		k.log.Warn("could not delete the keys of deleted devices", "user_id", userID, "error", err)
		return
	}
	k.announceDevices(ctx, userID, updates, nil)
}

// DeviceRenamed moves the device list of a user whose device has a new
// display name, which /keys/query reports to local users. Nothing goes
// over federation: Warp does not tell other servers what a user's devices
// are called.
func (k *API) DeviceRenamed(ctx context.Context, localpart, deviceID string) {
	userID := "@" + localpart + ":" + k.serverName
	keys, err := k.store.DeviceKeys(ctx, userID, []string{deviceID})
	if err != nil || len(keys) == 0 {
		return
	}
	if _, err := k.store.MarkDeviceListsChanged(ctx, []string{userID}); err != nil {
		k.log.Warn("could not record a device rename", "user_id", userID, "error", err)
		return
	}
	k.changed(ctx, userID)
}

// DeviceListChanges returns which of the given users' device lists moved
// in a window of the device list stream. It is what a sync reads.
func (k *API) DeviceListChanges(ctx context.Context, userIDs []string, from, to int64) ([]storage.DeviceListChange, error) {
	return k.store.DeviceListChanges(ctx, userIDs, from, to)
}