package federation

import (
	"context"
	"crypto/ed25519"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"sort"
	"sync"
	"time"

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

// Server keys: publishing ours, and fetching, caching and using everybody
// else's.
//
// Every signature another server makes — on an event, on a request — is
// checked against a key that server publishes at /_matrix/key/v2/server.
// The key ring fetches that document, checks that it is signed by the keys
// it lists (a key response vouches for itself; what makes it trustworthy
// is that it came from the server over TLS, or from a notary we chose to
// trust), and caches each key for as long as the server said it would stay
// valid — but never more than a week, which is the spec's cap on how long
// one fetch may be believed.

// ErrUnknownKey means a server's key could not be found: not cached, and
// neither the server nor any notary produced it.
var ErrUnknownKey = errors.New("signing key not found")

// maxKeyValidity caps how far ahead a fetched key's valid_until_ts is
// believed.
const maxKeyValidity = 7 * 24 * time.Hour

// fetchFailureBackoff stops a server whose key server is down from being
// asked again on every request that needs one of its keys.
const fetchFailureBackoff = time.Minute

// verifyKey and oldVerifyKey are the key objects of a key response.
type verifyKey struct {
	Key string `json:"key"`
}

type oldVerifyKey struct {
	Key       string `json:"key"`
	ExpiredTS int64  `json:"expired_ts"`
}

// serverKeyResponse is the document at /_matrix/key/v2/server.
type serverKeyResponse struct {
	ServerName    string                  `json:"server_name"`
	ValidUntilTS  int64                   `json:"valid_until_ts"`
	VerifyKeys    map[string]verifyKey    `json:"verify_keys"`
	OldVerifyKeys map[string]oldVerifyKey `json:"old_verify_keys"`
}

// KeyRingOptions configure a KeyRing.
type KeyRingOptions struct {
	ServerName string
	Key        *signing.KeyPair

	// Store caches fetched keys across restarts.
	Store storage.FederationStore

	// Client fetches keys. It refuses everything while federation is off,
	// which is exactly when no remote key can be needed.
	Client *Client

	// Notaries are the servers asked for a key the server itself cannot
	// be reached for.
	Notaries []string

	// Validity is how long our own key response says our key is valid.
	Validity time.Duration

	Logger *slog.Logger
	Now    func() time.Time
}

// KeyRing holds this server's key and everybody else's.
type KeyRing struct {
	serverName string
	key        *signing.KeyPair
	store      storage.FederationStore
	client     *Client
	notaries   []string
	validity   time.Duration
	log        *slog.Logger
	now        func() time.Time

	mu       sync.Mutex
	cache    map[string]map[string]storage.ServerKey // server -> key ID -> key
	failedAt map[string]time.Time
	inflight map[string]*fetchCall
}

type fetchCall struct {
	done chan struct{}
	err  error
}

// NewKeyRing builds a key ring.
func NewKeyRing(opts KeyRingOptions) *KeyRing {
	if opts.Logger == nil {
		opts.Logger = slog.Default()
	}
	if opts.Now == nil {
		opts.Now = time.Now
	}
	if opts.Validity <= 0 {
		opts.Validity = maxKeyValidity
	}
	return &KeyRing{
		serverName: opts.ServerName,
		key:        opts.Key,
		store:      opts.Store,
		client:     opts.Client,
		notaries:   opts.Notaries,
		validity:   opts.Validity,
		log:        opts.Logger,
		now:        opts.Now,
		cache:      map[string]map[string]storage.ServerKey{},
		failedAt:   map[string]time.Time{},
		inflight:   map[string]*fetchCall{},
	}
}

// OwnKeyResponse is this server's signed key document, as served at
// /_matrix/key/v2/server and handed out by the notary endpoints.
func (k *KeyRing) OwnKeyResponse() (json.RawMessage, error) {
	doc := serverKeyResponse{
		ServerName:    k.serverName,
		ValidUntilTS:  k.now().Add(k.validity).UnixMilli(),
		VerifyKeys:    map[string]verifyKey{k.key.KeyID(): {Key: k.key.PublicBase64()}},
		OldVerifyKeys: map[string]oldVerifyKey{},
	}
	raw, err := json.Marshal(doc)
	if err != nil {
		return nil, err
	}
	return signing.SignJSON(k.serverName, k.key, raw)
}

// PublicKey returns a server's key, fetching it when it is not cached or
// is not known to be valid at atTS.
//
// atTS is the time the key has to have been valid at, in milliseconds:
// the event's origin_server_ts when the room version enforces key
// validity, now for a request. Zero accepts any key the server has ever
// published, which is the rule for events in room versions before 5.
func (k *KeyRing) PublicKey(ctx context.Context, serverName, keyID string, atTS int64) (ed25519.PublicKey, error) {
	if serverName == k.serverName {
		if keyID == k.key.KeyID() {
			return k.key.Public, nil
		}
		return nil, fmt.Errorf("%w: %s is not a key of this server", ErrUnknownKey, keyID)
	}

	if key, ok, err := k.cached(ctx, serverName, keyID); err != nil {
		return nil, err
	} else if ok && usableAt(key, atTS) {
		return decodeKey(key)
	}

	if err := k.fetch(ctx, serverName); err != nil {
		// A key we already have may still be usable for an event whose
		// timestamp the old validity covers; only the refresh failed.
		if key, ok, _ := k.cached(ctx, serverName, keyID); ok && usableAt(key, atTS) {
			return decodeKey(key)
		}
		return nil, err
	}
	key, ok, err := k.cached(ctx, serverName, keyID)
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, fmt.Errorf("%w: %s has no key %s", ErrUnknownKey, serverName, keyID)
	}
	if !usableAt(key, atTS) {
		return nil, fmt.Errorf("%w: key %s of %s was not valid at %d", ErrUnknownKey, keyID, serverName, atTS)
	}
	return decodeKey(key)
}

// usableAt reports whether a key may verify something signed at atTS.
func usableAt(key storage.ServerKey, atTS int64) bool {
	if atTS == 0 {
		return true
	}
	if key.ExpiredTS != 0 {
		return atTS <= key.ExpiredTS
	}
	return atTS <= key.ValidUntilTS
}

func decodeKey(key storage.ServerKey) (ed25519.PublicKey, error) {
	raw, err := signing.DecodeBase64(key.PublicKey)
	if err != nil || len(raw) != ed25519.PublicKeySize {
		return nil, fmt.Errorf("%w: key %s of %s is malformed", ErrUnknownKey, key.KeyID, key.ServerName)
	}
	return ed25519.PublicKey(raw), nil
}

// cached reads a key from memory, then from the store.
func (k *KeyRing) cached(ctx context.Context, serverName, keyID string) (storage.ServerKey, bool, error) {
	k.mu.Lock()
	key, ok := k.cache[serverName][keyID]
	k.mu.Unlock()
	if ok {
		return key, true, nil
	}

	keys, err := k.store.ServerKeys(ctx, serverName, []string{keyID})
	if err != nil {
		return storage.ServerKey{}, false, err
	}
	if len(keys) == 0 {
		return storage.ServerKey{}, false, nil
	}
	k.remember(keys)
	return keys[0], true, nil
}

func (k *KeyRing) remember(keys []storage.ServerKey) {
	k.mu.Lock()
	defer k.mu.Unlock()
	for _, key := range keys {
		if k.cache[key.ServerName] == nil {
			k.cache[key.ServerName] = map[string]storage.ServerKey{}
		}
		k.cache[key.ServerName][key.KeyID] = key
	}
}

// fetch refreshes a server's keys, once at a time per server.
func (k *KeyRing) fetch(ctx context.Context, serverName string) error {
	k.mu.Lock()
	if at, failed := k.failedAt[serverName]; failed && k.now().Sub(at) < fetchFailureBackoff {
		k.mu.Unlock()
		return fmt.Errorf("%w: fetching the keys of %s failed recently", ErrUnknownKey, serverName)
	}
	if call, ok := k.inflight[serverName]; ok {
		k.mu.Unlock()
		select {
		case <-call.done:
			return call.err
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	call := &fetchCall{done: make(chan struct{})}
	k.inflight[serverName] = call
	k.mu.Unlock()

	call.err = k.fetchKeys(ctx, serverName)

	k.mu.Lock()
	delete(k.inflight, serverName)
	if call.err != nil {
		k.failedAt[serverName] = k.now()
	} else {
		delete(k.failedAt, serverName)
	}
	k.mu.Unlock()
	close(call.done)
	return call.err
}

// fetchKeys asks the server, then each notary in turn.
func (k *KeyRing) fetchKeys(ctx context.Context, serverName string) error {
	raw, err := k.client.ServerKeys(ctx, serverName)
	if err == nil {
		if err = k.acceptResponse(ctx, serverName, raw, ""); err == nil {
			return nil
		}
	}
	k.log.Debug("could not fetch keys directly", "server_name", serverName, "error", err)

	for _, notary := range k.notaries {
		if notary == serverName {
			continue
		}
		responses, nerr := k.client.QueryKeys(ctx, notary, map[string]map[string]KeyCriteria{serverName: {}})
		if nerr != nil {
			k.log.Debug("notary could not supply keys", "notary", notary, "server_name", serverName, "error", nerr)
			continue
		}
		for _, resp := range responses {
			if nerr = k.acceptResponse(ctx, serverName, resp, notary); nerr == nil {
				return nil
			}
		}
	}
	return fmt.Errorf("%w: no keys for %s: %v", ErrUnknownKey, serverName, err)
}

// acceptResponse checks a key response and stores its keys.
//
// The response has to be about the server asked for, and every signature
// the server made on it has to come from one of the keys it lists and has
// to verify — a response that cites a key it does not publish is either
// broken or forged. A response relayed by a notary must also carry the
// notary's own signature, which is what the notary is trusted for.
func (k *KeyRing) acceptResponse(ctx context.Context, serverName string, raw []byte, notary string) error {
	var doc serverKeyResponse
	if err := json.Unmarshal(raw, &doc); err != nil {
		return fmt.Errorf("decode key response: %w", err)
	}
	if doc.ServerName != serverName {
		return fmt.Errorf("key response is for %q, not %q", doc.ServerName, serverName)
	}
	sigs, err := signing.ExtractSignatures(raw)
	if err != nil {
		return err
	}
	if len(sigs[serverName]) == 0 {
		return fmt.Errorf("key response from %s is not signed by it", serverName)
	}
	for keyID := range sigs[serverName] {
		vk, ok := doc.VerifyKeys[keyID]
		if !ok {
			return fmt.Errorf("key response from %s is signed with %s, which it does not list", serverName, keyID)
		}
		pub, err := signing.DecodeBase64(vk.Key)
		if err != nil || len(pub) != ed25519.PublicKeySize {
			return fmt.Errorf("key %s of %s is malformed", keyID, serverName)
		}
		if err := signing.VerifyJSON(serverName, keyID, pub, raw); err != nil {
			return fmt.Errorf("key response from %s: %w", serverName, err)
		}
	}
	if notary != "" {
		if err := k.verifyJSONAt(ctx, notary, 0, raw); err != nil {
			return fmt.Errorf("key response for %s relayed by %s: %w", serverName, notary, err)
		}
	}

	now := k.now()
	validUntil := min(doc.ValidUntilTS, now.Add(maxKeyValidity).UnixMilli())
	var keys []storage.ServerKey
	for keyID, vk := range doc.VerifyKeys {
		keys = append(keys, storage.ServerKey{
			ServerName: serverName, KeyID: keyID, PublicKey: vk.Key,
			ValidUntilTS: validUntil, FetchedAtMS: now.UnixMilli(),
		})
	}
	for keyID, ok := range doc.OldVerifyKeys {
		keys = append(keys, storage.ServerKey{
			ServerName: serverName, KeyID: keyID, PublicKey: ok.Key,
			ValidUntilTS: ok.ExpiredTS, ExpiredTS: ok.ExpiredTS, FetchedAtMS: now.UnixMilli(),
		})
	}
	sort.Slice(keys, func(i, j int) bool { return keys[i].KeyID < keys[j].KeyID })
	if err := k.store.StoreServerKeys(ctx, keys); err != nil {
		return err
	}
	k.remember(keys)
	return nil
}

// VerifyJSON checks that a signed object carries a valid signature from
// serverName, made with a key valid now.
func (k *KeyRing) VerifyJSON(ctx context.Context, serverName string, msg []byte) error {
	return k.verifyJSONAt(ctx, serverName, k.now().UnixMilli(), msg)
}

// verifyJSONAt is VerifyJSON with the time the key must be valid at.
//
// One good signature is enough. A server rotating its key signs with the
// new one, and an object it signed with both is still its object if we
// can only fetch one of them.
func (k *KeyRing) verifyJSONAt(ctx context.Context, serverName string, atTS int64, msg []byte) error {
	sigs, err := signing.ExtractSignatures(msg)
	if err != nil {
		return err
	}
	keyIDs := make([]string, 0, len(sigs[serverName]))
	for keyID := range sigs[serverName] {
		if algorithm, _, ok := signing.SplitKeyID(keyID); ok && algorithm == signing.KeyIDPrefix {
			keyIDs = append(keyIDs, keyID)
		}
	}
	if len(keyIDs) == 0 {
		return fmt.Errorf("%w: from %s", signing.ErrNoSignature, serverName)
	}
	sort.Strings(keyIDs)

	var lastErr error
	for _, keyID := range keyIDs {
		pub, err := k.PublicKey(ctx, serverName, keyID, atTS)
		if err != nil {
			lastErr = err
			continue
		}
		if err := signing.VerifyJSON(serverName, keyID, pub, msg); err != nil {
			lastErr = err
			continue
		}
		return nil
	}
	return lastErr
}

// VerifyEventSignature implements auth.Verifier, and is what every event
// from another server is checked with.
func (k *KeyRing) VerifyEventSignature(ctx context.Context, serverName string, e *event.Event) error {
	keyIDs := make([]string, 0, len(e.Signatures()[serverName]))
	for keyID := range e.Signatures()[serverName] {
		if algorithm, _, ok := signing.SplitKeyID(keyID); ok && algorithm == signing.KeyIDPrefix {
			keyIDs = append(keyIDs, keyID)
		}
	}
	if len(keyIDs) == 0 {
		return fmt.Errorf("%w: event %s is not signed by %s", signing.ErrNoSignature, e.EventID(), serverName)
	}
	sort.Strings(keyIDs)

	var atTS int64
	if e.Rules().EnforceSigningKeyValidity {
		atTS = e.OriginServerTS()
	}
	var lastErr error
	for _, keyID := range keyIDs {
		pub, err := k.PublicKey(ctx, serverName, keyID, atTS)
		if err != nil {
			lastErr = err
			continue
		}
		if err := e.VerifySignature(serverName, keyID, pub); err != nil {
			lastErr = err
			continue
		}
		return nil
	}
	return lastErr
}

// NotaryResponse answers a /key/v2/query for this server's own keys. The
// key ring does not relay other servers' responses — it keeps their keys,
// not the documents they came in — so a query about anybody else is
// answered with nothing, which a requester treats as "ask somebody else".
func (k *KeyRing) NotaryResponse(serverName string) ([]json.RawMessage, error) {
	if serverName != k.serverName {
		return nil, nil
	}
	own, err := k.OwnKeyResponse()
	if err != nil {
		return nil, err
	}
	// The response is already signed by this server; as the notary it is
	// the same signature, and canonical JSON keeps it byte-identical.
	canonical, err := canonicaljson.Encode(own)
	if err != nil {
		return nil, err
	}
	return []json.RawMessage{canonical}, nil
}