package apptoken

import (
	"context"
	"crypto"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"errors"
	"fmt"
	"time"

	jose "github.com/go-jose/go-jose/v4"
	"github.com/hashicorp/vault/sdk/logical"
)

const (
	algEdDSA = "EdDSA"
	algRS256 = "RS256"

	rsaBits = 2048
)

// signingKey is one generation of the engine's signing key. The private
// material is stored (seal-wrapped) inside the barrier; the kid is the RFC 7638
// JWK thumbprint of the public key so it is stable and collision-resistant.
type signingKey struct {
	KeyID     string    `json:"key_id"`
	Algorithm string    `json:"algorithm"`
	Created   time.Time `json:"created"`
	// PrivateKey is PKCS#8 DER for both Ed25519 and RSA keys.
	PrivateKey []byte `json:"private_key"`
}

// keyset is the ordered set of signing keys: index 0 is the current signer,
// the remainder are retired keys kept in the JWKS so tokens signed before a
// rotation still validate.
type keyset struct {
	Keys []signingKey `json:"keys"`
}

func (k *signingKey) signer() (crypto.Signer, error) {
	priv, err := x509.ParsePKCS8PrivateKey(k.PrivateKey)
	if err != nil {
		return nil, fmt.Errorf("parsing stored private key: %w", err)
	}
	s, ok := priv.(crypto.Signer)
	if !ok {
		return nil, errors.New("stored key is not a crypto.Signer")
	}
	return s, nil
}

func (k *signingKey) joseAlgorithm() jose.SignatureAlgorithm {
	return jose.SignatureAlgorithm(k.Algorithm)
}

// publicJWK returns the public half of the key as a JWK, tagged with its kid,
// algorithm and signing use — the exact form published in the JWKS.
func (k *signingKey) publicJWK() (jose.JSONWebKey, error) {
	signer, err := k.signer()
	if err != nil {
		return jose.JSONWebKey{}, err
	}
	return jose.JSONWebKey{
		Key:       signer.Public(),
		KeyID:     k.KeyID,
		Algorithm: k.Algorithm,
		Use:       "sig",
	}, nil
}

// generateSigningKey creates a fresh key for the given algorithm and derives
// its kid from the JWK thumbprint of the public key.
func generateSigningKey(algorithm string, now time.Time) (signingKey, error) {
	var pub crypto.PublicKey
	var priv crypto.PrivateKey

	switch algorithm {
	case "", algEdDSA:
		algorithm = algEdDSA
		pk, sk, err := ed25519.GenerateKey(rand.Reader)
		if err != nil {
			return signingKey{}, fmt.Errorf("generating ed25519 key: %w", err)
		}
		pub, priv = pk, sk
	case algRS256:
		sk, err := rsa.GenerateKey(rand.Reader, rsaBits)
		if err != nil {
			return signingKey{}, fmt.Errorf("generating rsa key: %w", err)
		}
		pub, priv = &sk.PublicKey, sk
	default:
		return signingKey{}, fmt.Errorf("unsupported algorithm %q (supported: EdDSA, RS256)", algorithm)
	}

	der, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return signingKey{}, fmt.Errorf("marshaling private key: %w", err)
	}

	kid, err := thumbprint(pub, algorithm)
	if err != nil {
		return signingKey{}, err
	}

	return signingKey{
		KeyID:      kid,
		Algorithm:  algorithm,
		Created:    now,
		PrivateKey: der,
	}, nil
}

// thumbprint returns the base64url-encoded RFC 7638 JWK thumbprint of a public
// key, used as its stable kid.
func thumbprint(pub crypto.PublicKey, algorithm string) (string, error) {
	jwk := jose.JSONWebKey{Key: pub, Algorithm: algorithm, Use: "sig"}
	tp, err := jwk.Thumbprint(crypto.SHA256)
	if err != nil {
		return "", fmt.Errorf("computing key thumbprint: %w", err)
	}
	return base64URL(tp), nil
}

// current returns the active signing key (index 0).
func (ks *keyset) current() *signingKey {
	if ks == nil || len(ks.Keys) == 0 {
		return nil
	}
	return &ks.Keys[0]
}

// rotate makes a freshly generated key the current signer and trims retired
// keys so that at most retained previous keys remain published.
func (ks *keyset) rotate(algorithm string, retained int, now time.Time) error {
	nk, err := generateSigningKey(algorithm, now)
	if err != nil {
		return err
	}
	ks.Keys = append([]signingKey{nk}, ks.Keys...)
	if max := retained + 1; len(ks.Keys) > max {
		ks.Keys = ks.Keys[:max]
	}
	return nil
}

// jwks assembles the public JSON Web Key Set from every retained key.
func (ks *keyset) jwks() (jose.JSONWebKeySet, error) {
	out := jose.JSONWebKeySet{}
	if ks == nil {
		return out, nil
	}
	for i := range ks.Keys {
		jwk, err := ks.Keys[i].publicJWK()
		if err != nil {
			return out, err
		}
		out.Keys = append(out.Keys, jwk)
	}
	return out, nil
}

func (b *appTokenBackend) getKeyset(ctx context.Context, s logical.Storage) (*keyset, error) {
	entry, err := s.Get(ctx, keysetStoragePath)
	if err != nil {
		return nil, err
	}
	if entry == nil {
		return nil, nil
	}
	ks := &keyset{}
	if err := entry.DecodeJSON(ks); err != nil {
		return nil, err
	}
	return ks, nil
}

func (b *appTokenBackend) putKeyset(ctx context.Context, s logical.Storage, ks *keyset) error {
	entry, err := logical.StorageEntryJSON(keysetStoragePath, ks)
	if err != nil {
		return err
	}
	return s.Put(ctx, entry)
}

// getOrCreateKeyset returns the signing keyset, lazily generating an initial
// key with the configured algorithm on first use. Callers that mutate the
// keyset must hold b.keyLock.
func (b *appTokenBackend) getOrCreateKeyset(ctx context.Context, s logical.Storage, cfg *config) (*keyset, error) {
	ks, err := b.getKeyset(ctx, s)
	if err != nil {
		return nil, err
	}
	if ks != nil && len(ks.Keys) > 0 {
		return ks, nil
	}
	ks = &keyset{}
	if err := ks.rotate(cfg.Algorithm, cfg.RetainedKeys, time.Now()); err != nil {
		return nil, err
	}
	if err := b.putKeyset(ctx, s, ks); err != nil {
		return nil, err
	}
	return ks, nil
}