package jwk

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"math/big"
	"time"
)

// rsaKeyBits 為簽章金鑰位元數；RFC 7518 §3.5 規定 RS256 至少 2048 bits。
const rsaKeyBits = 2048

// SigningKey 為簽發 JWT 的 RSA 金鑰，對應 signing_keys 資料表。私鑰以
// PKCS#8 PEM 存於資料庫；Kid 為 RFC 7638 thumbprint，同時作為 JWKS 的
// kid 與 JWT header 的 kid，讓 RP 得以對應兩者。RetiredAt 為 nil 表示
// 使用中；輪替時舊金鑰先保留於 JWKS 一段時間（供已簽發的 token 驗證），
// 之後才退休停發，實現無縫金鑰輪替。
type SigningKey struct {
	ID            uint   `gorm:"primaryKey"`
	Kid           string `gorm:"uniqueIndex;size:43;not null"` // RFC 7638 thumbprint（32 bytes 的 base64url，43 字元）
	Algorithm     string `gorm:"size:8;not null;default:RS256"`
	PrivateKeyPEM string `gorm:"type:text;not null"` // PKCS#8 PEM 私鑰
	RetiredAt     *time.Time
	CreatedAt     time.Time
	UpdatedAt     time.Time
}

// Active 回傳金鑰是否使用中（未退休）。
func (k *SigningKey) Active() bool {
	return k.RetiredAt == nil
}

// NewSigningKey 產生新的 RSA-2048 簽章金鑰，並以公鑰的 RFC 7638
// SHA-256 thumbprint 作為 Kid。
func NewSigningKey() (*SigningKey, error) {
	key, err := rsa.GenerateKey(rand.Reader, rsaKeyBits)
	if err != nil {
		return nil, fmt.Errorf("generate rsa key: %w", err)
	}
	der, err := x509.MarshalPKCS8PrivateKey(key)
	if err != nil {
		return nil, fmt.Errorf("marshal private key: %w", err)
	}
	kid, err := thumbprint(&key.PublicKey)
	if err != nil {
		return nil, err
	}
	return &SigningKey{
		Kid:           kid,
		Algorithm:     AlgRS256,
		PrivateKeyPEM: string(pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: der})),
	}, nil
}

// PrivateKey 解析並回傳 RSA 私鑰，供簽發 JWT 使用。
func (k *SigningKey) PrivateKey() (*rsa.PrivateKey, error) {
	block, _ := pem.Decode([]byte(k.PrivateKeyPEM))
	if block == nil {
		return nil, errors.New("invalid PEM block")
	}
	parsed, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parse private key: %w", err)
	}
	key, ok := parsed.(*rsa.PrivateKey)
	if !ok {
		return nil, fmt.Errorf("not an RSA private key: %T", parsed)
	}
	return key, nil
}

// PublicJWK 回傳金鑰的公開 JWK（僅 kty、use、kid、alg、n、e），
// 供 JWKS 端點發佈。
func (k *SigningKey) PublicJWK() (*JWK, error) {
	key, err := k.PrivateKey()
	if err != nil {
		return nil, err
	}
	return &JWK{
		Kty: KeyTypeRSA,
		Use: KeyUseSig,
		Kid: k.Kid,
		Alg: k.Algorithm,
		N:   base64.RawURLEncoding.EncodeToString(key.PublicKey.N.Bytes()),
		E:   base64.RawURLEncoding.EncodeToString(big.NewInt(int64(key.PublicKey.E)).Bytes()),
	}, nil
}

// thumbprint 依 RFC 7638 計算 RSA 公鑰的 SHA-256 JWK thumbprint：對必要
// 參數依字典序（e、kty、n）組成的正規化 JSON 做 SHA-256，再以無填充
// base64url 編碼；此值即 JWK／JWT 的 kid。
func thumbprint(pub *rsa.PublicKey) (string, error) {
	// 欄位依字典序宣告，json.Marshal 的輸出即 RFC 7638 要求的正規化形式；
	// 成員僅含 base64url 字元，不會有 JSON 跳脫的差異。
	canonical := struct {
		E   string `json:"e"`
		Kty string `json:"kty"`
		N   string `json:"n"`
	}{
		E:   base64.RawURLEncoding.EncodeToString(big.NewInt(int64(pub.E)).Bytes()),
		Kty: KeyTypeRSA,
		N:   base64.RawURLEncoding.EncodeToString(pub.N.Bytes()),
	}
	b, err := json.Marshal(canonical)
	if err != nil {
		return "", fmt.Errorf("marshal canonical jwk: %w", err)
	}
	sum := sha256.Sum256(b)
	return base64.RawURLEncoding.EncodeToString(sum[:]), nil
}