package rsa

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

type KeyFormat int

const (
	RSA KeyFormat = iota
	COMMON
)

// 将指定的私钥编码为PEM格式，如果不指定具体格式将使用RSA PRIVATE KEY(PKCS1)格式。
func EncodePrivateKey(key *rsa.PrivateKey, format ...KeyFormat) ([]byte, error) {
	var block *pem.Block
	switch append(format, RSA)[0] {
	case COMMON:
		key, err := x509.MarshalPKCS8PrivateKey(key)
		if err != nil {
			return nil, fmt.Errorf("不支持的私钥格式，%w", err)
		}
		block = &pem.Block{
			Type:  "PRIVATE KEY",
			Bytes: key,
		}
	case RSA:
		fallthrough
	default:
		key := x509.MarshalPKCS1PrivateKey(key)
		block = &pem.Block{
			Type:  "RSA PRIVATE KEY",
			Bytes: key,
		}
	}
	return pem.EncodeToMemory(block), nil
}

// 将指定的公钥编码为PEM格式，如果不指定具体格式将使用RSA PUBLIC KEY(PKCS1)格式。
func EncodePublicKey(key *rsa.PublicKey, format ...KeyFormat) ([]byte, error) {
	var block *pem.Block
	switch append(format, RSA)[0] {
	case COMMON:
		key, err := x509.MarshalPKIXPublicKey(key)
		if err != nil {
			return nil, fmt.Errorf("不支持的公钥格式，%w", err)
		}
		block = &pem.Block{
			Type:  "PUBLIC KEY",
			Bytes: key,
		}
	case RSA:
		fallthrough
	default:
		key := x509.MarshalPKCS1PublicKey(key)
		block = &pem.Block{
			Type:  "RSA PUBLIC KEY",
			Bytes: key,
		}
	}
	return pem.EncodeToMemory(block), nil
}

// 将给定的PEM证书内容解析为RSA私钥。
// 仅能读取PEM证书中的第一个私钥编码段。
func DecodePrivateKey(cert []byte) (*rsa.PrivateKey, error) {
	var (
		block          *pem.Block
		pemRestContent = cert
	)
	for {
		block, pemRestContent = pem.Decode(pemRestContent)
		if block == nil {
			return nil, &KeyNotFoundError{}
		}
		if block.Type == "PRIVATE KEY" || block.Type == "RSA PRIVATE KEY" {
			break
		}
	}
	switch block.Type {
	case "PRIVATE KEY":
		key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("不支持的私钥格式，%w", err)
		}
		pKey, ok := key.(*rsa.PrivateKey)
		if !ok {
			return nil, fmt.Errorf("提供的私钥不是RSA专用私钥，%w", err)
		}
		return pKey, nil
	case "RSA PRIVATE KEY":
		key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("不支持的私钥格式，%w", err)
		}
		return key, nil
	default:
		return nil, fmt.Errorf("不支持的私钥格式或者私钥不存在，%s", block.Type)
	}
}

// 将给定的PEM证书内容解析为RSA公钥。
// 仅能读取PEM证书中的第一个公钥编码段。
func DecodePublicKey(cert []byte) (*rsa.PublicKey, error) {
	var (
		block          *pem.Block
		pemRestContent = cert
	)
	for {
		block, pemRestContent = pem.Decode(pemRestContent)
		if block == nil {
			return nil, &KeyNotFoundError{}
		}
		if block.Type == "PUBLIC KEY" || block.Type == "RSA PUBLIC KEY" {
			break
		}
	}
	switch block.Type {
	case "PUBLIC KEY":
		key, err := x509.ParsePKIXPublicKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("不支持的公钥格式，%w", err)
		}
		pKey, ok := key.(*rsa.PublicKey)
		if !ok {
			return nil, fmt.Errorf("提供的公钥不是RSA专用公钥，%w", err)
		}
		return pKey, nil
	case "RSA PUBLIC KEY":
		key, err := x509.ParsePKCS1PublicKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("不支持的公钥格式，%w", err)
		}
		return key, nil
	default:
		return nil, fmt.Errorf("不支持的公钥格式或者公钥不存在，%s", block.Type)
	}
}