package authjwt

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"

	"github.com/golang-jwt/jwt/v5"
)

var _ Signer = (*rsaSigner)(nil)

type rsaSigner struct {
	private *rsa.PrivateKey
	public  *rsa.PublicKey
}

// NewRSASigner returns a Signer backed by RSA-SHA256 (RS256).
// The public key is derived from the private key — no separate argument needed.
func NewRSASigner(privateKey *rsa.PrivateKey) Signer {
	return &rsaSigner{private: privateKey, public: &privateKey.PublicKey}
}

// NewRSASignerFromPEM parses a PKCS#8 or PKCS#1 PEM-encoded RSA private key.
func NewRSASignerFromPEM(pemKey []byte) (Signer, error) {
	block, _ := pem.Decode(pemKey)
	if block == nil {
		return nil, fmt.Errorf("no PEM block found")
	}
	key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		rsaKey, err2 := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err2 != nil {
			return nil, fmt.Errorf("parse RSA private key: %w", err)
		}
		return NewRSASigner(rsaKey), nil
	}
	rsaKey, ok := key.(*rsa.PrivateKey)
	if !ok {
		return nil, fmt.Errorf("PEM key is not an RSA private key")
	}
	return NewRSASigner(rsaKey), nil
}

func (s *rsaSigner) Sign(claims jwt.Claims) (string, error) {
	return jwt.NewWithClaims(jwt.SigningMethodRS256, claims).SignedString(s.private)
}

func (s *rsaSigner) Verify(tokenString string) (*jwt.Token, error) {
	return jwt.Parse(tokenString, func(t *jwt.Token) (any, error) {
		if _, ok := t.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, fmt.Errorf("unexpected signing method %q", t.Header["alg"])
		}
		return s.public, nil
	}, jwt.WithJSONNumber())
}