package signing

import (
	"crypto/ed25519"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"maps"
	"strings"

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

// Signatures is the signatures block of a signed JSON object: server name,
// then key ID ("ed25519:abc"), then the unpadded base64 signature.
//
// An object can carry several, and adding one must not invalidate the
// others — which is why signatures, along with unsigned, is left out of the
// bytes a signature covers.
type Signatures map[string]map[string]string

// Get returns one signature, and whether it was there.
func (s Signatures) Get(serverName, keyID string) (string, bool) {
	sig, ok := s[serverName][keyID]
	return sig, ok
}

// Errors from verification. ErrNoSignature and ErrBadSignature are
// separate because only the second one is evidence of anything: a missing
// signature usually means we have not fetched the right key yet, while a
// wrong one means the event has been tampered with.
var (
	// ErrNoSignature means the object carries no signature from the server
	// and key we were asked about.
	ErrNoSignature = errors.New("no signature from that server and key")

	// ErrBadSignature means a signature was present and did not verify.
	ErrBadSignature = errors.New("signature does not verify")
)

// SigningBytes returns the bytes a signature over msg covers: the object
// with its signatures and unsigned properties removed, in canonical JSON.
//
// Both ends of a signature have to agree on these bytes exactly, so this
// is the single place the rule is written down.
func SigningBytes(msg []byte) ([]byte, error) {
	obj, err := decodeObject(msg)
	if err != nil {
		return nil, err
	}
	delete(obj, "signatures")
	delete(obj, "unsigned")
	return canonicaljson.Marshal(obj)
}

// SignJSON returns msg with a signature from serverName added.
//
// Signatures already on the object are kept; one from the same server and
// key ID is replaced. The result is canonical JSON, so signing is
// idempotent in the sense that matters: signing the output again produces
// the same bytes.
func SignJSON(serverName string, key *KeyPair, msg []byte) ([]byte, error) {
	if serverName == "" {
		return nil, errors.New("sign JSON: server name must not be empty")
	}
	obj, err := decodeObject(msg)
	if err != nil {
		return nil, err
	}

	existing, err := signaturesOf(obj)
	if err != nil {
		return nil, err
	}

	// The signature covers the object without signatures or unsigned, so
	// take them out before canonicalising and put them back afterwards.
	unsigned, hasUnsigned := obj["unsigned"]
	delete(obj, "signatures")
	delete(obj, "unsigned")

	signed, err := canonicaljson.Marshal(obj)
	if err != nil {
		return nil, fmt.Errorf("sign JSON: %w", err)
	}
	sig := ed25519.Sign(key.Private, signed)

	if existing == nil {
		existing = Signatures{}
	}
	if existing[serverName] == nil {
		existing[serverName] = map[string]string{}
	}
	existing[serverName][key.KeyID()] = base64.RawStdEncoding.EncodeToString(sig)

	obj["signatures"], err = json.Marshal(existing)
	if err != nil {
		return nil, fmt.Errorf("sign JSON: encode signatures: %w", err)
	}
	if hasUnsigned {
		obj["unsigned"] = unsigned
	}
	return canonicaljson.Marshal(obj)
}

// VerifyJSON checks msg for a valid signature from serverName made with
// keyID, using the public key given.
//
// keyID is the full identifier including the algorithm, as it appears in
// the signatures block: "ed25519:abc123".
func VerifyJSON(serverName, keyID string, pub ed25519.PublicKey, msg []byte) error {
	obj, err := decodeObject(msg)
	if err != nil {
		return err
	}
	sigs, err := signaturesOf(obj)
	if err != nil {
		return err
	}
	encoded, ok := sigs.Get(serverName, keyID)
	if !ok {
		return fmt.Errorf("%w: %s %s", ErrNoSignature, serverName, keyID)
	}
	sig, err := decodeBase64(encoded)
	if err != nil {
		return fmt.Errorf("%w: %s %s: decode: %w", ErrBadSignature, serverName, keyID, err)
	}
	if len(sig) != ed25519.SignatureSize {
		return fmt.Errorf("%w: %s %s: signature is %d bytes, want %d",
			ErrBadSignature, serverName, keyID, len(sig), ed25519.SignatureSize)
	}

	delete(obj, "signatures")
	delete(obj, "unsigned")
	signed, err := canonicaljson.Marshal(obj)
	if err != nil {
		return err
	}
	if !ed25519.Verify(pub, signed, sig) {
		return fmt.Errorf("%w: %s %s", ErrBadSignature, serverName, keyID)
	}
	return nil
}

// ExtractSignatures returns the signatures block of a signed object. It is
// how a caller finds out which keys it needs before it has any of them.
func ExtractSignatures(msg []byte) (Signatures, error) {
	obj, err := decodeObject(msg)
	if err != nil {
		return nil, err
	}
	return signaturesOf(obj)
}

// SplitKeyID splits "ed25519:abc123" into its algorithm and identifier.
func SplitKeyID(keyID string) (algorithm, id string, ok bool) {
	algorithm, id, ok = strings.Cut(keyID, ":")
	if !ok || algorithm == "" || id == "" {
		return "", "", false
	}
	return algorithm, id, true
}

// decodeObject parses a JSON object into its properties, leaving each value
// as the bytes it arrived as.
func decodeObject(msg []byte) (map[string]json.RawMessage, error) {
	var obj map[string]json.RawMessage
	if err := json.Unmarshal(msg, &obj); err != nil {
		return nil, fmt.Errorf("parse signed JSON: %w", err)
	}
	if obj == nil {
		return nil, errors.New("parse signed JSON: value is null, not an object")
	}
	return obj, nil
}

// signaturesOf decodes the signatures property, which may be absent.
func signaturesOf(obj map[string]json.RawMessage) (Signatures, error) {
	raw, ok := obj["signatures"]
	if !ok {
		return nil, nil
	}
	var sigs Signatures
	if err := json.Unmarshal(raw, &sigs); err != nil {
		return nil, fmt.Errorf("parse signatures: %w", err)
	}
	// Clone so that a caller mutating the result cannot reach back into
	// the object we parsed.
	out := make(Signatures, len(sigs))
	for server, keys := range sigs {
		out[server] = maps.Clone(keys)
	}
	return out, nil
}