package keyserver

import (
	"crypto/ed25519"
	"encoding/json"
	"fmt"
	"strings"

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

// The shapes of the keys a client uploads, as far as the server reads
// them. Everything else in a key is kept as it was sent: a key is a signed
// object, and the server re-encoding a field it did not understand would
// break the signature.

// deviceKeysObject is a device's identity keys.
type deviceKeysObject struct {
	UserID     string            `json:"user_id"`
	DeviceID   string            `json:"device_id"`
	Algorithms []string          `json:"algorithms"`
	Keys       map[string]string `json:"keys"`
}

// crossSigningKeyObject is a cross-signing key.
type crossSigningKeyObject struct {
	UserID string            `json:"user_id"`
	Usage  []string          `json:"usage"`
	Keys   map[string]string `json:"keys"`
}

// canonical re-encodes a key in canonical JSON, which is what the store
// compares when it decides whether an upload is the same key again.
func canonical(raw json.RawMessage) ([]byte, error) {
	out, err := canonicaljson.Encode(raw)
	if err != nil {
		return nil, invalidParamf("a key is not valid JSON: %v", err)
	}
	return out, nil
}

// object decodes a JSON object into its fields.
func object(raw []byte) (map[string]json.RawMessage, error) {
	var obj map[string]json.RawMessage
	if err := json.Unmarshal(raw, &obj); err != nil || obj == nil {
		return nil, invalidParamf("a key must be a JSON object")
	}
	return obj, nil
}

// withoutField drops one field from a JSON object, in canonical form.
func withoutField(raw []byte, field string) ([]byte, error) {
	obj, err := object(raw)
	if err != nil {
		return nil, err
	}
	delete(obj, field)
	return canonicaljson.Marshal(obj)
}

// signaturesOf reads a key's signatures block.
func signaturesOf(raw []byte) (signing.Signatures, error) {
	sigs, err := signing.ExtractSignatures(raw)
	if err != nil {
		return nil, invalidParamf("a key's signatures are malformed: %v", err)
	}
	return sigs, nil
}

// withSignatures merges stored signatures into a key, and optionally sets
// its unsigned.device_display_name. Signatures the key already carries are
// kept; a stored one under the same signer and key ID wins, since it is
// the more recent.
func withSignatures(raw []byte, sigs []storage.KeySignature, displayName string) (json.RawMessage, error) {
	if len(sigs) == 0 && displayName == "" {
		return raw, nil
	}
	obj, err := object(raw)
	if err != nil {
		return nil, err
	}
	if len(sigs) > 0 {
		merged, err := signaturesOf(raw)
		if err != nil {
			return nil, err
		}
		if merged == nil {
			merged = signing.Signatures{}
		}
		for _, sig := range sigs {
			if merged[sig.SignerUserID] == nil {
				merged[sig.SignerUserID] = map[string]string{}
			}
			merged[sig.SignerUserID][sig.SignerKeyID] = sig.Signature
		}
		if obj["signatures"], err = json.Marshal(merged); err != nil {
			return nil, err
		}
	}
	if displayName != "" {
		unsigned := map[string]json.RawMessage{}
		if existing, ok := obj["unsigned"]; ok {
			_ = json.Unmarshal(existing, &unsigned)
		}
		if unsigned["device_display_name"], err = json.Marshal(displayName); err != nil {
			return nil, err
		}
		if obj["unsigned"], err = json.Marshal(unsigned); err != nil {
			return nil, err
		}
	}
	return canonicaljson.Marshal(obj)
}

// ed25519Key returns the one ed25519 key of a cross-signing key, by its
// unpadded base64 public key, which is also its key ID.
func (o crossSigningKeyObject) ed25519Key() (string, bool) {
	if len(o.Keys) != 1 {
		return "", false
	}
	for keyID, value := range o.Keys {
		algorithm, name, ok := signing.SplitKeyID(keyID)
		if !ok || algorithm != "ed25519" || name != value {
			return "", false
		}
		return value, true
	}
	return "", false
}

// publicKey decodes an unpadded base64 ed25519 public key.
func publicKey(encoded string) (ed25519.PublicKey, error) {
	raw, err := signing.DecodeBase64(encoded)
	if err != nil || len(raw) != ed25519.PublicKeySize {
		return nil, fmt.Errorf("%q is not an ed25519 public key", encoded)
	}
	return ed25519.PublicKey(raw), nil
}

// verify checks that a key carries a valid signature by userID's key
// "ed25519:"+keyName, whose public key is pub.
func verify(raw []byte, userID, keyName, pub string) error {
	key, err := publicKey(pub)
	if err != nil {
		return invalidSignaturef("%v", err)
	}
	if err := signing.VerifyJSON(userID, "ed25519:"+keyName, key, raw); err != nil {
		return invalidSignaturef("%v", err)
	}
	return nil
}

// validKeyID reports whether a one-time key ID has the "algorithm:id"
// shape.
func validKeyID(keyID string) bool {
	algorithm, name, ok := strings.Cut(keyID, ":")
	return ok && algorithm != "" && name != ""
}