package services

import (
	"crypto/rand"
	"crypto/rsa"
	"encoding/base64"
	"fmt"
	"math/big"
)

type KeyManager struct {
	privateKey *rsa.PrivateKey
	publicKey  *rsa.PublicKey
	kid        string
}

type JSONWebKey struct {
	Kty string `json:"kty"`
	Kid string `json:"kid"`
	Use string `json:"use"`
	N   string `json:"n"`
	E   string `json:"e"`
	Alg string `json:"alg"`
}

type JSONWebKeySet struct {
	Keys []JSONWebKey `json:"keys"`
}

func NewKeyManager() (*KeyManager, error) {
	// 生成 RSA 密钥对
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, fmt.Errorf("failed to generate RSA key: %v", err)
	}

	// 生成密钥 ID
	kid := generateKeyID()

	return &KeyManager{
		privateKey: privateKey,
		publicKey:  &privateKey.PublicKey,
		kid:        kid,
	}, nil
}

func (km *KeyManager) GetJWKS() (*JSONWebKeySet, error) {
	// 将公钥转换为 JWK 格式
	jwk := JSONWebKey{
		Kty: "RSA",
		Kid: km.kid,
		Use: "sig",
		Alg: "RS256",
		N:   base64.RawURLEncoding.EncodeToString(km.publicKey.N.Bytes()),
		E:   base64.RawURLEncoding.EncodeToString(big.NewInt(int64(km.publicKey.E)).Bytes()),
	}

	return &JSONWebKeySet{
		Keys: []JSONWebKey{jwk},
	}, nil
}

func (km *KeyManager) GetPrivateKey() *rsa.PrivateKey {
	return km.privateKey
}

func (km *KeyManager) GetKID() string {
	return km.kid
}

func generateKeyID() string {
	b := make([]byte, 16)
	rand.Read(b)
	return base64.RawURLEncoding.EncodeToString(b)
}