package util

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

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

func ParseBase64PrivateKey(b64 string) (*rsa.PrivateKey, error) {
	decoded, err := base64.StdEncoding.DecodeString(b64)
	if err != nil {
		return nil, fmt.Errorf("failed to decode base64 key: %v", err)
	}

	key, err := x509.ParsePKCS8PrivateKey(decoded)
	return key.(*rsa.PrivateKey), err
}

func ParseBase64PublicKey(b64 string) (*rsa.PublicKey, error) {
	decoded, err := base64.StdEncoding.DecodeString(b64)
	if err != nil {
		return nil, fmt.Errorf("failed to decode base64 key: %v", err)
	}

	pubInterface, err := x509.ParsePKIXPublicKey(decoded)
	if err != nil {
		return nil, fmt.Errorf("failed to parse public key: %v", err)
	}

	pubKey, ok := pubInterface.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("not an RSA public key")
	}

	return pubKey, nil
}

func SignJwtToken(claims jwt.Claims, key string) (string, error) {
	privateKey, err := ParseBase64PrivateKey(key)
	if err != nil {
		return "", err
	}

	token := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)

	token.Header["kid"] = "my-rsa-key-1"

	s, err := token.SignedString(privateKey)
	if err != nil {
		return "", err
	}

	return s, nil
}

func VerifyToken(token string, key string, claims jwt.Claims) (*jwt.Token, error) {
	publicKey, err := ParseBase64PublicKey(key)
	if err != nil {
		return nil, err
	}

	parsed, err := jwt.ParseWithClaims(token, claims, func(t *jwt.Token) (any, error) {
		if _, ok := t.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"])
		}
		return publicKey, nil
	})

	if err != nil {
		return nil, fmt.Errorf("invalid token: %w", err)
	}

	if !parsed.Valid {
		return nil, fmt.Errorf("token is not valid")
	}

	return parsed, nil
}