// Package zitadel talks to the ZITADEL management API with the controller's own
// machine key.
//
// The controller needs a credential to rotate credentials. It signs a short
// assertion with its private key, exchanges it for an access token, and reuses
// that token until shortly before it expires.
package zitadel

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"time"
)

// assertionLifetime is deliberately short. The assertion proves one request.
const assertionLifetime = 60 * time.Second

// newAssertionClaims builds the claim set for a machine-key assertion.
//
// ZITADEL expects `iss` to name the caller, `sub` to name the user it acts as,
// and `aud` to equal the issuer of the instance.
func newAssertionClaims(clientID, userID, issuer string, now time.Time) map[string]any {
	if userID == "" {
		userID = clientID
	}
	return map[string]any{
		"iss": clientID,
		"sub": userID,
		"aud": issuer,
		"iat": now.Unix(),
		"exp": now.Add(assertionLifetime).Unix(),
	}
}

func base64URL(b []byte) string {
	return base64.RawURLEncoding.EncodeToString(b)
}

// SignAssertion produces a signed RS256 JWT for the jwt-bearer grant.
//
// The header carries `kid`, so ZITADEL selects the matching registered key. That
// is what makes overlapping keys possible: two keys can be valid at once and the
// assertion says which one it used.
func SignAssertion(privateKeyPEM []byte, keyID, clientID, userID, issuer string, now time.Time) (string, error) {
	block, _ := pem.Decode(privateKeyPEM)
	if block == nil {
		return "", fmt.Errorf("private key is not PEM encoded")
	}

	key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		// ZITADEL exports PKCS#1 keys as well.
		if rsaKey, rsaErr := x509.ParsePKCS1PrivateKey(block.Bytes); rsaErr == nil {
			key = rsaKey
		} else {
			return "", fmt.Errorf("private key does not parse: %w", err)
		}
	}
	rsaKey, ok := key.(*rsa.PrivateKey)
	if !ok {
		return "", fmt.Errorf("private key is not RSA")
	}

	header := map[string]any{"alg": "RS256", "typ": "JWT", "kid": keyID}
	headerJSON, err := json.Marshal(header)
	if err != nil {
		return "", err
	}
	claimsJSON, err := json.Marshal(newAssertionClaims(clientID, userID, issuer, now))
	if err != nil {
		return "", err
	}

	signingInput := base64URL(headerJSON) + "." + base64URL(claimsJSON)
	digest := sha256.Sum256([]byte(signingInput))
	signature, err := rsa.SignPKCS1v15(rand.Reader, rsaKey, crypto.SHA256, digest[:])
	if err != nil {
		return "", fmt.Errorf("cannot sign the assertion: %w", err)
	}

	return signingInput + "." + base64URL(signature), nil
}