package storage

import (
	"context"
	"errors"
)

// Errors the key store returns for cases a caller has to tell apart.
var (
	// ErrKeyConflict is a one-time key uploaded under an ID that already
	// holds a different key. The spec has a key ID name one key for its
	// whole life: a client that reuses one has lost track of what it
	// uploaded, and silently replacing the key would leave a peer that
	// already claimed the old one unable to decrypt.
	ErrKeyConflict = errors.New("a different key is already stored under this key ID")

	// ErrNoSuchBackup is a key backup version that does not exist, or has
	// been deleted.
	ErrNoSuchBackup = errors.New("no such key backup version")

	// ErrNoSuchBackupKey is a room key that is not in a backup.
	ErrNoSuchBackupKey = errors.New("no such key in the backup")
)

// The purposes a cross-signing key can have, as the spec spells its usage.
const (
	CrossSigningMaster      = "master"
	CrossSigningSelfSigning = "self_signing"
	CrossSigningUserSigning = "user_signing"
)

// DeviceKeys is the identity keys of one device, as the device uploaded and
// signed them.
type DeviceKeys struct {
	UserID   string
	DeviceID string

	// JSON is the device_keys object as uploaded, without `unsigned`: the
	// server writes that itself on the way out, and anything a client put
	// there is not something it signed.
	JSON []byte
}

// OneTimeKey is one one-time or fallback key of a device.
type OneTimeKey struct {
	// KeyID is "algorithm:id", the key's name in the map it was uploaded
	// in.
	KeyID string

	// JSON is the key as uploaded: a bare string for an unsigned
	// curve25519 key, an object with signatures for a signed one.
	JSON []byte
}

// Algorithm is the part of the key ID before the colon.
func (k OneTimeKey) Algorithm() string {
	for i := 0; i < len(k.KeyID); i++ {
		if k.KeyID[i] == ':' {
			return k.KeyID[:i]
		}
	}
	return k.KeyID
}

// ClaimedKey is a key handed out by a claim, and whether it was a fallback
// key rather than a one-time one.
type ClaimedKey struct {
	OneTimeKey
	Fallback bool
}

// KeySignature is one signature of a key by another key, uploaded apart
// from the key it signs.
//
// The keys a client uploads come signed by the device or key that made
// them. Cross-signing adds signatures afterwards — a device by the user's
// self-signing key, the user's master key by a device, another user's
// master key by the user-signing key — and those arrive through
// /keys/signatures/upload, long after the key they sign. They are kept
// apart and merged in when the key is read, rather than written into it,
// so that a key's own JSON stays exactly what its owner uploaded.
type KeySignature struct {
	// TargetUserID owns the signed key, and TargetKeyID names it: a device
	// ID for a device key, the unpadded base64 public key for a
	// cross-signing key.
	TargetUserID string
	TargetKeyID  string

	// SignerUserID made the signature with SignerKeyID, which is
	// "ed25519:" and the signing key's ID.
	SignerUserID string
	SignerKeyID  string

	Signature string
}

// DeviceListUpdate is one move of a user's device list: the position this
// change was given, and the position of the change before it. A remote
// server is told both, and asks for the whole list when the previous one
// is not the last it heard of.
type DeviceListUpdate struct {
	UserID string

	// DeviceID is the device the change was about, or empty for a change
	// to the user's cross-signing keys.
	DeviceID string

	StreamPos int64

	// PrevPos is the user's position before this change, or 0 for their
	// first.
	PrevPos int64
}

// DeviceListChange is a user whose device list moved, at the position of
// the latest move.
type DeviceListChange struct {
	UserID    string
	StreamPos int64
}

// RemoteDevice is one device of a remote user, as their server described
// it.
type RemoteDevice struct {
	DeviceID    string
	DisplayName string

	// JSON is the device_keys object as the user's server sent it.
	JSON []byte
}

// RemoteKeys is the cached copy of a remote user's keys.
type RemoteKeys struct {
	UserID string

	// StreamID is the position in the user's server's device list stream
	// the copy is at: the stream_id of the snapshot it was fetched as, or
	// of the last update applied to it since.
	StreamID int64

	// Stale says an update was missed, or did not make sense, and the copy
	// is not to be served until it has been fetched again. A stale copy
	// may have no devices at all: an update about a user with nothing
	// cached leaves a stale marker, so that a snapshot fetched before the
	// update happened is not taken for a current one.
	Stale bool

	Devices []RemoteDevice

	// CrossSigning is the user's master and self-signing keys by purpose.
	CrossSigning map[string][]byte
}

// RemoteDeviceUpdate is one m.device_list_update, as the store applies it.
type RemoteDeviceUpdate struct {
	UserID      string
	DeviceID    string
	DisplayName string

	// JSON is the device's new keys, or nil for a device that was deleted
	// or that the update carried no keys for.
	JSON    []byte
	Deleted bool

	StreamID int64
	PrevIDs  []int64
}

// RemoteUpdateResult is what applying a RemoteDeviceUpdate did.
type RemoteUpdateResult int

const (
	// RemoteUpdateApplied is an update that followed on from the copy
	// held, and is now part of it.
	RemoteUpdateApplied RemoteUpdateResult = iota + 1

	// RemoteUpdateDuplicate is an update at or before the position the
	// copy is already at: a retransmission, changing nothing.
	RemoteUpdateDuplicate

	// RemoteUpdateStale is an update to a copy that is stale: one that did
	// not follow on from the copy held, which is now marked so, or any
	// update to a copy already marked. Either way the copy has to be
	// fetched again.
	RemoteUpdateStale

	// RemoteUpdateNotCached is an update about a user with nothing cached
	// at all. A stale marker is written at the update's position, so that
	// a snapshot fetched before it is not stored as current.
	RemoteUpdateNotCached
)

// BackupVersion is one version of a user's server-side key backup.
type BackupVersion struct {
	UserID  string
	Version string

	Algorithm string

	// AuthData is the algorithm's parameters as the client wrote them: the
	// backup's public key and the signatures over it.
	AuthData []byte

	// ETag changes every time a key in the backup does, so a client can
	// tell its copy of the backup is stale without downloading it.
	ETag string

	// Count is how many keys the backup holds.
	Count int
}

// BackupKey is one room key in a backup.
type BackupKey struct {
	RoomID    string
	SessionID string

	FirstMessageIndex int64
	ForwardedCount    int64
	IsVerified        bool

	// SessionData is the encrypted key, which the server cannot read.
	SessionData []byte
}

// Replaces reports whether a key should replace one already backed up for
// the same session, by the rule the spec gives: a verified key beats an
// unverified one, then the key that decrypts more of the session — the
// lower first message index — and then the one that has been forwarded
// fewer times.
func (k BackupKey) Replaces(existing BackupKey) bool {
	if k.IsVerified != existing.IsVerified {
		return k.IsVerified
	}
	if k.FirstMessageIndex != existing.FirstMessageIndex {
		return k.FirstMessageIndex < existing.FirstMessageIndex
	}
	return k.ForwardedCount < existing.ForwardedCount
}

// KeyStore persists end-to-end encryption keys: the storage half of the
// keyserver component.
//
// Users are full user IDs here rather than localparts, because most of what
// is stored is about users of other servers too: signatures of their keys,
// and the device list stream, which records their changes as well as ours
// so that a local client knows to fetch their keys again.
//
// Every write that changes what /keys/query would answer about a user moves
// the device list stream, inside the same transaction, and reports the
// position it moved to. That is what makes "whose keys changed since this
// sync token" a question with one exact answer.
type KeyStore interface {
	// SetDeviceKeys stores a device's identity keys. Storing what is
	// already there changes nothing and returns a zero update; anything
	// else moves the device list stream.
	SetDeviceKeys(ctx context.Context, keys DeviceKeys) (DeviceListUpdate, error)

	// DeviceKeys reads the identity keys of a user's devices. An empty
	// deviceIDs reads every device that has uploaded keys.
	DeviceKeys(ctx context.Context, userID string, deviceIDs []string) ([]DeviceKeys, error)

	// DeleteDeviceKeys drops everything stored for the given devices —
	// identity keys, one-time keys and fallback keys — and moves the
	// device list stream once for each device that had identity keys. An
	// empty deviceIDs drops every device of the user.
	DeleteDeviceKeys(ctx context.Context, userID string, deviceIDs []string) ([]DeviceListUpdate, error)

	// AddOneTimeKeys stores one-time keys. A key under an ID that already
	// holds the same key is ignored; one under an ID holding a different
	// key fails the whole batch with ErrKeyConflict.
	AddOneTimeKeys(ctx context.Context, userID, deviceID string, keys []OneTimeKey) error

	// OneTimeKeyCounts counts a device's unclaimed one-time keys by
	// algorithm.
	OneTimeKeyCounts(ctx context.Context, userID, deviceID string) (map[string]int, error)

	// SetFallbackKeys stores fallback keys, replacing any earlier one of
	// the same algorithm. A key the same as the one stored keeps its used
	// flag: a client re-uploading its fallback key on every start must not
	// make a used one look fresh.
	SetFallbackKeys(ctx context.Context, userID, deviceID string, keys []OneTimeKey) error

	// UnusedFallbackKeyTypes lists the algorithms a device has a fallback
	// key for that has not been claimed.
	UnusedFallbackKeyTypes(ctx context.Context, userID, deviceID string) ([]string, error)

	// ClaimKey hands out one key of an algorithm for a device, deleting it:
	// a one-time key if there is one, the fallback key otherwise, which is
	// marked used and kept. The boolean is false when the device has
	// neither.
	ClaimKey(ctx context.Context, userID, deviceID, algorithm string) (ClaimedKey, bool, error)

	// SetCrossSigningKeys stores a user's cross-signing keys by purpose and
	// moves the device list stream. A purpose mapped to nil is deleted,
	// which is what replacing the master key does to the keys it signed.
	SetCrossSigningKeys(ctx context.Context, userID string, keys map[string][]byte) (DeviceListUpdate, error)

	// CrossSigningKeys reads a user's cross-signing keys by purpose.
	CrossSigningKeys(ctx context.Context, userID string) (map[string][]byte, error)

	// AddSignatures stores signatures, replacing an earlier one by the same
	// key of the same key, and moves the device list stream of every user
	// whose key gained one.
	AddSignatures(ctx context.Context, sigs []KeySignature) ([]DeviceListUpdate, error)

	// Signatures reads the stored signatures of a user's keys.
	Signatures(ctx context.Context, targetUserID string) ([]KeySignature, error)

	// MarkDeviceListsChanged moves the device list stream for users whose
	// keys changed somewhere this store does not hold them: on another
	// server, which said so in an m.device_list_update.
	MarkDeviceListsChanged(ctx context.Context, userIDs []string) (int64, error)

	// DeviceListChanges returns which of the given users' device lists
	// moved after from and at or before to.
	DeviceListChanges(ctx context.Context, userIDs []string, from, to int64) ([]DeviceListChange, error)

	// DeviceListPosition is a user's latest position in the device list
	// stream, or 0 when their list has never moved.
	DeviceListPosition(ctx context.Context, userID string) (int64, error)

	// RemoteKeys reads the cached copy of a remote user's keys. The
	// boolean is false when nothing at all is held for them.
	RemoteKeys(ctx context.Context, userID string) (RemoteKeys, bool, error)

	// SetRemoteKeys replaces the cached copy of a remote user's keys with
	// a snapshot fetched from their server, and clears the stale mark. A
	// snapshot older than the position already held — fetched while an
	// update was arriving — is not stored, and the boolean is false.
	SetRemoteKeys(ctx context.Context, keys RemoteKeys) (bool, error)

	// UpdateRemoteDevice applies an m.device_list_update to a cached copy:
	// only when one of its PrevIDs is the position held, otherwise the
	// copy is marked stale. The position held never moves backwards.
	UpdateRemoteDevice(ctx context.Context, update RemoteDeviceUpdate) (RemoteUpdateResult, error)

	// SetRemoteCrossSigningKeys replaces a cached user's master and
	// self-signing keys; a purpose missing from keys is deleted. The
	// boolean is false, and nothing is written, when the user has no
	// fresh copy cached.
	SetRemoteCrossSigningKeys(ctx context.Context, userID string, keys map[string][]byte) (bool, error)

	// RemoteKeyUsers lists every user something is cached for.
	RemoteKeyUsers(ctx context.Context) ([]string, error)

	// DeleteRemoteKeys drops everything cached for the given users.
	DeleteRemoteKeys(ctx context.Context, userIDs []string) error

	// CreateBackup creates a new key backup version, which becomes the
	// user's current one, and returns its version.
	CreateBackup(ctx context.Context, userID, algorithm string, authData []byte) (string, error)

	// Backup reads one backup version, or the current one when version is
	// empty. A deleted or unknown version is ErrNoSuchBackup.
	Backup(ctx context.Context, userID, version string) (BackupVersion, error)

	// UpdateBackup replaces a backup version's auth data.
	UpdateBackup(ctx context.Context, userID, version string, authData []byte) error

	// DeleteBackup deletes a backup version and every key in it. Its
	// number is not handed out again.
	DeleteBackup(ctx context.Context, userID, version string) error

	// PutBackupKeys stores room keys in a backup version, each replacing
	// the key already there for its session only when BackupKey.Replaces
	// says it should, and returns the version as it then stands.
	PutBackupKeys(ctx context.Context, userID, version string, keys []BackupKey) (BackupVersion, error)

	// BackupKeys reads the keys of a backup version: all of them, one
	// room's (roomID set), or one session's (both set).
	BackupKeys(ctx context.Context, userID, version, roomID, sessionID string) ([]BackupKey, error)

	// DeleteBackupKeys deletes keys from a backup version with the same
	// scoping as BackupKeys, and returns the version as it then stands.
	DeleteBackupKeys(ctx context.Context, userID, version, roomID, sessionID string) (BackupVersion, error)
}