package busauth

import (
	"encoding/base64"

	server "github.com/nats-io/nats-server/v2/server"
	"github.com/nats-io/nkeys"
)

// nkeyAuthenticator is a NATS server.Authentication that authorizes a client by
// verifying the nkey signature over the server-presented nonce and then
// consulting the bus user allowlist. Authorization is checked on every new
// connection via the injected predicate (not a static Options.Nkeys map), so
// revoking a user denies its next connection without restarting the server.
type nkeyAuthenticator struct {
	// isAuthorized reports whether the lowercase-hex Ed25519 public key behind an
	// nkey belongs to an active bus user. Injected (membership.Store.IsAuthorized)
	// so this package stays free of the store dependency.
	isAuthorized func(signPubHex string) bool
}

// NewNkeyAuthenticator builds a NATS custom authenticator backed by isAuthorized.
// Pass it to embeddednats so the data plane only accepts registered identities.
func NewNkeyAuthenticator(isAuthorized func(signPubHex string) bool) server.Authentication {
	return &nkeyAuthenticator{isAuthorized: isAuthorized}
}

// Check verifies the client's nkey signature against the nonce the server
// presented, then maps the nkey to its allowlist key and checks authorization.
// Any malformed input or failed verification yields false (fail closed).
func (a *nkeyAuthenticator) Check(c server.ClientAuthentication) bool {
	signPubHex, ok := verifyNkey(c)
	if !ok {
		return false
	}
	return a.isAuthorized(signPubHex)
}

// verifyNkey performs the shared nkey verification: it checks the client's
// signature against the server-presented nonce and returns the lowercase-hex
// Ed25519 public key behind the nkey. ok is false on any malformed input or
// failed verification (fail closed). The signature decoding mirrors
// nats-server's own (raw-url base64, then std base64 fallback) so genuine
// clients using nats.Nkey are accepted unchanged.
func verifyNkey(c server.ClientAuthentication) (signPubHex string, ok bool) {
	opts := c.GetOpts()
	if opts.Nkey == "" {
		return "", false
	}
	sig, err := base64.RawURLEncoding.DecodeString(opts.Sig)
	if err != nil {
		sig, err = base64.StdEncoding.DecodeString(opts.Sig)
		if err != nil {
			return "", false
		}
	}
	pub, err := nkeys.FromPublicKey(opts.Nkey)
	if err != nil {
		return "", false
	}
	if err := pub.Verify(c.GetNonce(), sig); err != nil {
		return "", false
	}
	signPubHex, err = SignPubHexFromNkey(opts.Nkey)
	if err != nil {
		return "", false
	}
	return signPubHex, true
}

// PermissionsFunc maps a connecting identity (lowercase-hex Ed25519 signing key)
// to the NATS permissions it should be granted for this connection. Returning an
// error denies the connection (fail closed). It is how the data plane enforces
// per-subject access from room membership (issue 0003e, audit H4 residual).
type PermissionsFunc func(signPubHex string) (*server.Permissions, error)

// nkeyAuthenticatorACL is the nkey authenticator that ALSO scopes the connection
// to per-subject permissions derived from room membership. NATS evaluates
// permissions once, at connect time, so a peer that joins a room after
// connecting must reconnect (client.RefreshSession) to gain that room's subject
// — the dynamic-membership reconnection model the audit deferred to this issue.
type nkeyAuthenticatorACL struct {
	isAuthorized func(signPubHex string) bool
	perms        PermissionsFunc
	// internalPubHex is the lowercase-hex Ed25519 public key of membershipd's own
	// ephemeral internal service identity. A connection that proves that key is
	// granted full permissions WITHOUT consulting the allowlist, so the service can
	// bootstrap and manage JetStream (the replicated nonce bucket and, when
	// decentralized, the control-plane KV buckets) against its own embedded server
	// even while the data plane confines every client to its rooms. Empty disables
	// the internal-identity path entirely (behavior identical to a plain ACL
	// authenticator).
	internalPubHex string
}

// NewNkeyAuthenticatorACL builds an authenticator that authorizes by the bus
// allowlist AND registers per-subject permissions from perms. A registered but
// permission-less peer can no longer subscribe to or publish on arbitrary
// subjects: it is confined to the subjects of the rooms it belongs to (plus the
// client infrastructure subjects perms includes). This is the per-subject ACL
// the 0004 hardening left as a residual.
func NewNkeyAuthenticatorACL(isAuthorized func(signPubHex string) bool, perms PermissionsFunc) server.Authentication {
	return &nkeyAuthenticatorACL{isAuthorized: isAuthorized, perms: perms}
}

// NewNkeyAuthenticatorACLInternal is NewNkeyAuthenticatorACL that also recognizes
// membershipd's internal service identity (internalPubHex, the lowercase hex of
// its ephemeral Ed25519 public key): a connection proving that key is granted
// full permissions without an allowlist lookup, so the service can create and
// manage JetStream against its own embedded server under enforce (issue 0006a/c —
// the replicated nonce bucket and the control-plane KV). Every other identity
// goes through the allowlist + per-subject ACL unchanged. An empty internalPubHex
// is identical to NewNkeyAuthenticatorACL, so this is a superset and safe to use
// everywhere the plain constructor was used.
func NewNkeyAuthenticatorACLInternal(isAuthorized func(signPubHex string) bool, perms PermissionsFunc, internalPubHex string) server.Authentication {
	return &nkeyAuthenticatorACL{isAuthorized: isAuthorized, perms: perms, internalPubHex: internalPubHex}
}

// fullPermissions grants publish and subscribe on every subject (">"). It is the
// permission set for membershipd's own internal service connection, which must
// manage the JetStream control plane (nonce bucket + KV buckets) over NATS. It is
// NEVER granted to a bus user — only to the process's own ephemeral internal
// identity, recognized by exact public-key match in Check.
func fullPermissions() *server.Permissions {
	sp := &server.SubjectPermission{Allow: []string{">"}}
	return &server.Permissions{Publish: sp, Subscribe: sp}
}

// Check verifies the nkey, authorizes against the allowlist, then derives and
// registers the connection's subject permissions. A permissions-derivation
// error denies the connection (fail closed) rather than granting open access.
func (a *nkeyAuthenticatorACL) Check(c server.ClientAuthentication) bool {
	signPubHex, ok := verifyNkey(c)
	if !ok {
		return false
	}
	// membershipd's own internal service identity bypasses the allowlist and is
	// granted full permissions so the service can bootstrap JetStream under
	// enforce. The key is matched exactly against the cryptographically verified
	// connecting key, so no other identity can claim these permissions.
	if a.internalPubHex != "" && signPubHex == a.internalPubHex {
		c.RegisterUser(&server.User{Permissions: fullPermissions()})
		return true
	}
	if !a.isAuthorized(signPubHex) {
		return false
	}
	perms, err := a.perms(signPubHex)
	if err != nil {
		return false // fail closed: never grant open access on a derivation error
	}
	c.RegisterUser(&server.User{Permissions: perms})
	return true
}