package jwtauth

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

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

// Verifier validates JWT strings. Use this in services that receive but do not
// issue tokens (e.g. microservices given only the RSA public key).
type Verifier interface {
	Verify(tokenString string) (*jwt.Token, error)
}

// Signer signs and verifies JWTs.
// NewHMACSigner and NewRSASigner return implementations backed by HS256 and RS256.
type Signer interface {
	Verifier
	Sign(claims jwt.Claims) (string, error)
}

// --- HMAC (HS256) ---

type hmacSigner struct{ secret []byte }

// NewHMACSigner returns a Signer backed by HMAC-SHA256.
// secret should be at least 32 bytes; shorter values are accepted but weakened.
func NewHMACSigner(secret []byte) Signer {
	return &hmacSigner{secret: secret}
}

func (s *hmacSigner) Sign(claims jwt.Claims) (string, error) {
	return jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString(s.secret)
}

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

// --- RSA (RS256) ---

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

// NewRSASigner returns a Signer backed by RSA-SHA256.
// 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
	})
}

// --- RSA public-key-only verifier ---

type rsaPublicVerifier struct{ public *rsa.PublicKey }

// NewRSAPublicKeyVerifier returns a Verifier backed by an RSA public key.
// Use this in services that verify tokens but never issue them.
func NewRSAPublicKeyVerifier(publicKey *rsa.PublicKey) Verifier {
	return &rsaPublicVerifier{public: publicKey}
}

// NewRSAPublicKeyVerifierFromPEM parses a PKIX or PKCS#1 PEM-encoded RSA public key.
func NewRSAPublicKeyVerifierFromPEM(pemKey []byte) (Verifier, error) {
	block, _ := pem.Decode(pemKey)
	if block == nil {
		return nil, fmt.Errorf("no PEM block found")
	}
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		rsaPub, err2 := x509.ParsePKCS1PublicKey(block.Bytes)
		if err2 != nil {
			return nil, fmt.Errorf("parse RSA public key: %w", err)
		}
		return NewRSAPublicKeyVerifier(rsaPub), nil
	}
	rsaPub, ok := pub.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("PEM key is not an RSA public key")
	}
	return NewRSAPublicKeyVerifier(rsaPub), nil
}

func (v *rsaPublicVerifier) 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 v.public, nil
	})
}