package federation

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"

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

// The X-Matrix authorization scheme: how one server proves to another that
// a request came from it.
//
// The sender signs a small JSON object naming the request — method, URI,
// origin, destination and, when there is one, the JSON body — and sends
// the signature in the Authorization header:
//
//	Authorization: X-Matrix origin="a.example",destination="b.example",key="ed25519:k",sig="..."
//
// The receiver rebuilds the same object from the request it actually got
// and checks the signature against the origin's published key. Nothing in
// the header is trusted until that check passes: the origin it names is a
// claim, and it is exactly the claim being checked.

// ErrBadAuthorization is the root of every X-Matrix parsing and
// verification failure. Each one is answered with M_UNAUTHORIZED.
var ErrBadAuthorization = errors.New("invalid X-Matrix authorization")

// XMatrix is one parsed X-Matrix header.
type XMatrix struct {
	Origin string

	// Destination is the server the request was addressed to. Older
	// servers leave it out, and a receiver then has to assume it was
	// itself.
	Destination string

	// Key is the full key ID, algorithm included.
	Key string

	// Sig is the unpadded base64 signature.
	Sig string
}

// maxAuthorizationLength bounds the header before anything is parsed. A
// real one is a couple of hundred bytes; this leaves room for long server
// names without letting a client make the parser chew on a megabyte.
const maxAuthorizationLength = 4096

// ParseXMatrix parses an Authorization header value.
//
// The grammar is the spec's: the scheme, then comma-separated key=value
// parameters whose values may be quoted. Quoted values may contain
// backslash escapes, and unknown parameters are ignored, so that a header
// carrying something a newer server added still parses.
func ParseXMatrix(header string) (XMatrix, error) {
	if len(header) > maxAuthorizationLength {
		return XMatrix{}, fmt.Errorf("%w: header is %d bytes", ErrBadAuthorization, len(header))
	}
	scheme, rest, ok := strings.Cut(strings.TrimSpace(header), " ")
	if !ok || !strings.EqualFold(scheme, "X-Matrix") {
		return XMatrix{}, fmt.Errorf("%w: not an X-Matrix header", ErrBadAuthorization)
	}

	var x XMatrix
	params, err := parseAuthParams(rest)
	if err != nil {
		return XMatrix{}, err
	}
	for name, value := range params {
		switch strings.ToLower(name) {
		case "origin":
			x.Origin = value
		case "destination":
			x.Destination = value
		case "key":
			x.Key = value
		case "sig":
			x.Sig = value
		}
	}

	switch {
	case x.Origin == "":
		return XMatrix{}, fmt.Errorf("%w: no origin", ErrBadAuthorization)
	case x.Key == "":
		return XMatrix{}, fmt.Errorf("%w: no key", ErrBadAuthorization)
	case x.Sig == "":
		return XMatrix{}, fmt.Errorf("%w: no sig", ErrBadAuthorization)
	}
	if _, _, ok := signing.SplitKeyID(x.Key); !ok {
		return XMatrix{}, fmt.Errorf("%w: key %q is not algorithm:id", ErrBadAuthorization, x.Key)
	}
	return x, nil
}

// parseAuthParams splits `a="x",b=y` into its parameters.
func parseAuthParams(s string) (map[string]string, error) {
	params := map[string]string{}
	i := 0
	for {
		// Skip separators and whitespace between parameters.
		for i < len(s) && (s[i] == ',' || s[i] == ' ' || s[i] == '\t') {
			i++
		}
		if i >= len(s) {
			return params, nil
		}

		eq := strings.IndexByte(s[i:], '=')
		if eq <= 0 {
			return nil, fmt.Errorf("%w: parameter without a value", ErrBadAuthorization)
		}
		name := strings.TrimSpace(s[i : i+eq])
		if name == "" || strings.ContainsAny(name, " \t\",") {
			return nil, fmt.Errorf("%w: malformed parameter name %q", ErrBadAuthorization, name)
		}
		i += eq + 1

		var value string
		if i < len(s) && s[i] == '"' {
			var b strings.Builder
			i++
			closed := false
			for i < len(s) {
				c := s[i]
				if c == '\\' && i+1 < len(s) {
					b.WriteByte(s[i+1])
					i += 2
					continue
				}
				if c == '"' {
					closed = true
					i++
					break
				}
				b.WriteByte(c)
				i++
			}
			if !closed {
				return nil, fmt.Errorf("%w: unterminated quoted value", ErrBadAuthorization)
			}
			value = b.String()
		} else {
			end := strings.IndexByte(s[i:], ',')
			if end < 0 {
				end = len(s) - i
			}
			value = strings.TrimSpace(s[i : i+end])
			i += end
		}

		if _, dup := params[name]; dup {
			// Two origins in one header is not something to resolve by
			// picking one: whichever we picked, the other server meant
			// something else.
			return nil, fmt.Errorf("%w: parameter %q given twice", ErrBadAuthorization, name)
		}
		params[name] = value
	}
}

// String renders the header value. Every parameter is quoted, which the
// spec allows and which every implementation accepts.
func (x XMatrix) String() string {
	quote := func(s string) string {
		return `"` + strings.NewReplacer(`\`, `\\`, `"`, `\"`).Replace(s) + `"`
	}
	return "X-Matrix origin=" + quote(x.Origin) +
		",destination=" + quote(x.Destination) +
		",key=" + quote(x.Key) +
		",sig=" + quote(x.Sig)
}

// requestJSON is the object an X-Matrix signature covers.
type requestJSON struct {
	Method      string          `json:"method"`
	URI         string          `json:"uri"`
	Origin      string          `json:"origin"`
	Destination string          `json:"destination"`
	Content     json.RawMessage `json:"content,omitempty"`
}

// signableRequest builds the bytes a request signature covers.
//
// uri is the request target exactly as it was sent — path and query, still
// escaped. Both ends must agree on it byte for byte, which is why the
// receiver reads it from the raw request line rather than reassembling it
// from a parsed URL.
func signableRequest(method, uri, origin, destination string, content []byte) ([]byte, error) {
	obj := requestJSON{Method: method, URI: uri, Origin: origin, Destination: destination}
	if len(content) > 0 {
		obj.Content = content
	}
	raw, err := json.Marshal(obj)
	if err != nil {
		return nil, fmt.Errorf("encode request for signing: %w", err)
	}
	return canonicaljson.Encode(raw)
}

// SignRequest produces the X-Matrix header for an outbound request.
func SignRequest(key *signing.KeyPair, origin, destination, method, uri string, content []byte) (XMatrix, error) {
	msg, err := signableRequest(method, uri, origin, destination, content)
	if err != nil {
		return XMatrix{}, err
	}
	signed, err := signing.SignJSON(origin, key, msg)
	if err != nil {
		return XMatrix{}, err
	}
	sigs, err := signing.ExtractSignatures(signed)
	if err != nil {
		return XMatrix{}, err
	}
	sig, _ := sigs.Get(origin, key.KeyID())
	return XMatrix{Origin: origin, Destination: destination, Key: key.KeyID(), Sig: sig}, nil
}

// requestSignedJSON rebuilds the signed object an inbound header claims to
// cover, with the header's signature attached, ready for VerifyJSON.
func requestSignedJSON(x XMatrix, destination, method, uri string, content []byte) ([]byte, error) {
	msg, err := signableRequest(method, uri, x.Origin, destination, content)
	if err != nil {
		return nil, err
	}
	var obj map[string]json.RawMessage
	if err := json.Unmarshal(msg, &obj); err != nil {
		return nil, err
	}
	sigs, err := json.Marshal(signing.Signatures{x.Origin: {x.Key: x.Sig}})
	if err != nil {
		return nil, err
	}
	obj["signatures"] = sigs
	return json.Marshal(obj)
}