package keyserver

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

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

// Server-side key backup: a user's room keys, encrypted to a backup key
// only the user holds, so that a new device can read history the user's
// other devices received.
//
// The server stores what it is given and decides one thing: when two
// copies of the same session's key arrive, which to keep. Everything else
// about a backup — what the algorithm is, what the auth data says, whether
// the backup key is trusted — is the clients' business.

// BackupInfo is a backup version as GET /room_keys/version reports it.
type BackupInfo struct {
	Algorithm string          `json:"algorithm"`
	AuthData  json.RawMessage `json:"auth_data"`
	Count     int             `json:"count"`
	ETag      string          `json:"etag"`
	Version   string          `json:"version"`
}

// KeyBackupData is one backed-up session key on the wire.
type KeyBackupData struct {
	FirstMessageIndex int64           `json:"first_message_index"`
	ForwardedCount    int64           `json:"forwarded_count"`
	IsVerified        bool            `json:"is_verified"`
	SessionData       json.RawMessage `json:"session_data"`
}

// RoomKeyBackup is one room's keys: session ID to key.
type RoomKeyBackup struct {
	Sessions map[string]KeyBackupData `json:"sessions"`
}

// BackupKeysResult is what a write or delete of keys answers with.
type BackupKeysResult struct {
	Count int    `json:"count"`
	ETag  string `json:"etag"`
}

func infoOf(v storage.BackupVersion) BackupInfo {
	return BackupInfo{
		Algorithm: v.Algorithm,
		AuthData:  v.AuthData,
		Count:     v.Count,
		ETag:      v.ETag,
		Version:   v.Version,
	}
}

// authDataOf checks auth_data is an object and canonicalises it.
func authDataOf(raw json.RawMessage) ([]byte, error) {
	if len(raw) == 0 {
		return nil, missingParamf("auth_data is required")
	}
	if _, err := object(raw); err != nil {
		return nil, invalidParamf("auth_data must be an object")
	}
	out, err := canonicaljson.Encode(raw)
	if err != nil {
		return nil, invalidParamf("auth_data is not valid JSON: %v", err)
	}
	return out, nil
}

// CreateBackup creates a new backup version, which becomes the current
// one.
func (k *API) CreateBackup(ctx context.Context, userID, algorithm string, authData json.RawMessage) (string, error) {
	if algorithm == "" {
		return "", missingParamf("algorithm is required")
	}
	encoded, err := authDataOf(authData)
	if err != nil {
		return "", err
	}
	return k.store.CreateBackup(ctx, userID, algorithm, encoded)
}

// Backup reads a backup version, or the current one when version is empty.
func (k *API) Backup(ctx context.Context, userID, version string) (BackupInfo, error) {
	v, err := k.store.Backup(ctx, userID, version)
	if err != nil {
		return BackupInfo{}, err
	}
	return infoOf(v), nil
}

// UpdateBackup replaces a version's auth data. The algorithm cannot
// change — every key in the version is encrypted with it — and a version
// in the body has to be the one in the path.
func (k *API) UpdateBackup(ctx context.Context, userID, version, algorithm, bodyVersion string, authData json.RawMessage) error {
	if bodyVersion != "" && bodyVersion != version {
		return invalidParamf("the version in the body, %s, is not the one in the path, %s", bodyVersion, version)
	}
	existing, err := k.store.Backup(ctx, userID, version)
	if err != nil {
		return err
	}
	if algorithm != existing.Algorithm {
		return invalidParamf("a backup version's algorithm cannot change from %s", existing.Algorithm)
	}
	encoded, err := authDataOf(authData)
	if err != nil {
		return err
	}
	return k.store.UpdateBackup(ctx, userID, version, encoded)
}

// DeleteBackup deletes a backup version and its keys.
func (k *API) DeleteBackup(ctx context.Context, userID, version string) error {
	return k.store.DeleteBackup(ctx, userID, version)
}

// PutBackupKeys stores keys in a version, which has to be the current one:
// a client writing into an old version after another device replaced it
// would be backing up keys nobody can restore from.
func (k *API) PutBackupKeys(ctx context.Context, userID, version string, rooms map[string]RoomKeyBackup) (BackupKeysResult, error) {
	if version == "" {
		return BackupKeysResult{}, missingParamf("the version parameter is required")
	}
	current, err := k.store.Backup(ctx, userID, "")
	switch {
	case errors.Is(err, storage.ErrNoSuchBackup):
		return BackupKeysResult{}, ErrNoSuchBackup
	case err != nil:
		return BackupKeysResult{}, err
	case current.Version != version:
		if _, err := k.store.Backup(ctx, userID, version); errors.Is(err, storage.ErrNoSuchBackup) {
			return BackupKeysResult{}, ErrNoSuchBackup
		}
		return BackupKeysResult{}, &WrongBackupVersionError{Current: current.Version}
	}

	var keys []storage.BackupKey
	for roomID, room := range rooms {
		for sessionID, data := range room.Sessions {
			if len(data.SessionData) == 0 {
				return BackupKeysResult{}, missingParamf("session_data is required for %s in %s", sessionID, roomID)
			}
			sessionData, err := canonicaljson.Encode(data.SessionData)
			if err != nil {
				return BackupKeysResult{}, invalidParamf("session_data of %s is not valid JSON: %v", sessionID, err)
			}
			keys = append(keys, storage.BackupKey{
				RoomID: roomID, SessionID: sessionID,
				FirstMessageIndex: data.FirstMessageIndex,
				ForwardedCount:    data.ForwardedCount,
				IsVerified:        data.IsVerified,
				SessionData:       sessionData,
			})
		}
	}
	v, err := k.store.PutBackupKeys(ctx, userID, version, keys)
	if err != nil {
		return BackupKeysResult{}, err
	}
	return BackupKeysResult{Count: v.Count, ETag: v.ETag}, nil
}

// BackupKeys reads keys of a version, grouped by room.
func (k *API) BackupKeys(ctx context.Context, userID, version, roomID, sessionID string) (map[string]RoomKeyBackup, error) {
	if version == "" {
		return nil, missingParamf("the version parameter is required")
	}
	keys, err := k.store.BackupKeys(ctx, userID, version, roomID, sessionID)
	if err != nil {
		return nil, err
	}
	out := map[string]RoomKeyBackup{}
	for _, key := range keys {
		room, ok := out[key.RoomID]
		if !ok {
			room = RoomKeyBackup{Sessions: map[string]KeyBackupData{}}
			out[key.RoomID] = room
		}
		room.Sessions[key.SessionID] = KeyBackupData{
			FirstMessageIndex: key.FirstMessageIndex,
			ForwardedCount:    key.ForwardedCount,
			IsVerified:        key.IsVerified,
			SessionData:       key.SessionData,
		}
	}
	return out, nil
}

// DeleteBackupKeys deletes keys from a version.
func (k *API) DeleteBackupKeys(ctx context.Context, userID, version, roomID, sessionID string) (BackupKeysResult, error) {
	if version == "" {
		return BackupKeysResult{}, missingParamf("the version parameter is required")
	}
	v, err := k.store.DeleteBackupKeys(ctx, userID, version, roomID, sessionID)
	if err != nil {
		return BackupKeysResult{}, err
	}
	return BackupKeysResult{Count: v.Count, ETag: v.ETag}, nil
}