package rsa

import (
	"math/big"

	"xdx.jelly/xgcl/api/common"
	"xdx.jelly/xgcl/gmath"
	"xdx.jelly/xgcl/grand"
	"xdx.jelly/xgcl/identifier"
	"xdx.jelly/xgcl/rsa/internal"
	"xdx.jelly/xgcl/sm/sm4"
)

// GenerateKeyPairRSA corresponding to SDF_GenerateKeyPair_RSA
func (PKCS1v15) GenerateKeyPairRSA(keyBits uint32) (*PrivateKey, *PublicKey, error) {
	if keyBits > common.RSAref_MAX_BITS {
		return nil, nil, common.SDR_KEYERR
	}
	privateKey, err := internal.GenerateKey(grand.Reader, int(keyBits))
	// privateKey, err := internal.GenerateKey(rand.New(rand.NewSource(1234)), int(keyBits))

	if err != nil {
		return nil, nil, common.SDR_KEYERR
	}

	return privateKey, &privateKey.PublicKey, nil
}

// GenerateEnvelopedKeyPairRSA 生成密钥对保护结构并用publicKey加密输出。
// 对应SDF_GenerateRSAKeyPairWithIPK_RSA和SDF_GenerateRSAKeyPairWithEPK_RSA
// 注意：内部不检查keyBits.
func (p PKCS1v15) GenerateEnvelopedKeyPairRSA(keyBits uint32, publicKey *PublicKey) (*common.EnvelopedRSAKey, error) {
	sm4key := make([]byte, sm4.BlockSize)
	n, err := grand.GenerateRandom(sm4key)
	if n < sm4.BlockSize || err != nil {
		return nil, common.SDR_RANDERR
	}

	sk, pk, err := p.GenerateKeyPairRSA(keyBits)
	if err != nil {
		return nil, err
	}
	enveloped := common.EnvelopedRSAKey{
		SymmAlgID: identifier.SGDSM4ECB,
		Bits:      keyBits,
	}

	var data [common.RSAref_MAX_LEN]byte
	buf := data[:]
	buf0 := data[:common.RSAref_MAX_PLEN]

	gmath.FillBytes(sk.N, buf)
	copy(enveloped.EncryptedPriKey[0][:], buf[:common.RSAref_MAX_PLEN])
	copy(enveloped.EncryptedPriKey[1][:], buf[common.RSAref_MAX_PLEN:])

	gmath.FillBytes(new(big.Int).SetInt64(int64(sk.E)), buf)
	copy(enveloped.EncryptedPriKey[2][:], buf[:common.RSAref_MAX_PLEN])
	copy(enveloped.EncryptedPriKey[3][:], buf[common.RSAref_MAX_PLEN:])

	gmath.FillBytes(sk.D, buf)
	copy(enveloped.EncryptedPriKey[4][:], buf[:common.RSAref_MAX_PLEN])
	copy(enveloped.EncryptedPriKey[5][:], buf[common.RSAref_MAX_PLEN:])

	if len(sk.Primes) != 2 {
		return nil, common.SDR_KEYERR
	}
	gmath.FillBytes(sk.Primes[0], buf0)
	copy(enveloped.EncryptedPriKey[6][:], buf0)
	gmath.FillBytes(sk.Primes[1], buf0)
	copy(enveloped.EncryptedPriKey[7][:], buf0)

	gmath.FillBytes(sk.Precomputed.Dp, buf0)
	copy(enveloped.EncryptedPriKey[8][:], buf0)

	gmath.FillBytes(sk.Precomputed.Dq, buf0)
	copy(enveloped.EncryptedPriKey[9][:], buf0)

	gmath.FillBytes(sk.Precomputed.Qinv, buf0)
	copy(enveloped.EncryptedPriKey[10][:], buf0)

	for i := 0; i < len(enveloped.EncryptedPriKey); i++ {
		sm4.EncryptECB(enveloped.EncryptedPriKey[i][:], sm4key, enveloped.EncryptedPriKey[i][:])
	}

	err = pk.MarshalSDF(&enveloped.PubKey)
	if err != nil {
		return nil, common.SDR_UNKNOWERR
	}

	C, err := internal.EncryptPKCS1v15(grand.Reader, publicKey, sm4key)
	if err != nil {
		return nil, common.SDR_KEYERR
	}
	enveloped.L = uint32(len(C))
	copy(enveloped.C[:], C)
	return &enveloped, nil
}

// ImportEnvelopedKeyPairRSA 导入RSA密钥保护结构。
func (PKCS1v15) ImportEnvelopedKeyPairRSA(envelopedRSAKey *common.EnvelopedRSAKey, privateKey *PrivateKey) (*PrivateKey, *PublicKey, error) {
	sk := internal.NewPrivateKey()

	sm4key, err := internal.DecryptPKCS1v15(grand.Reader, privateKey, envelopedRSAKey.C[:envelopedRSAKey.L])
	if err != nil || len(sm4key) != sm4.BlockSize {
		return nil, nil, common.SDR_INARGERR
	}
	for i := 0; i < len(envelopedRSAKey.EncryptedPriKey); i++ {
		sm4.DecryptECB(envelopedRSAKey.EncryptedPriKey[i][:], sm4key, envelopedRSAKey.EncryptedPriKey[i][:])
	}

	var data [common.RSAref_MAX_LEN]byte
	buf := data[:]

	copy(buf[:common.RSAref_MAX_PLEN], envelopedRSAKey.EncryptedPriKey[0][:])
	copy(buf[common.RSAref_MAX_PLEN:], envelopedRSAKey.EncryptedPriKey[1][:])
	sk.N.SetBytes(buf)

	copy(buf[:common.RSAref_MAX_PLEN], envelopedRSAKey.EncryptedPriKey[2][:])
	copy(buf[common.RSAref_MAX_PLEN:], envelopedRSAKey.EncryptedPriKey[3][:])
	sk.E = int(new(big.Int).SetBytes(buf).Int64())

	copy(buf[:common.RSAref_MAX_PLEN], envelopedRSAKey.EncryptedPriKey[4][:])
	copy(buf[common.RSAref_MAX_PLEN:], envelopedRSAKey.EncryptedPriKey[5][:])
	sk.D.SetBytes(buf)

	sk.Primes[0].SetBytes(envelopedRSAKey.EncryptedPriKey[6][:])
	sk.Primes[1].SetBytes(envelopedRSAKey.EncryptedPriKey[7][:])
	sk.Precompute()

	return sk, &sk.PublicKey, nil
}

// GenerateKeyRSA 生成会话密钥并用publicKey加密输出
func (PKCS1v15) GenerateKeyRSA(keyBits uint32, publicKey *PublicKey) ([]byte, []byte, error) {
	key := make([]byte, keyBits)
	c, err := internal.EncryptPKCS1v15(grand.Reader, publicKey, key)
	if err != nil {
		return nil, nil, common.SDR_KEYERR
	}
	return key, c, nil
}

// GenerateKeyRSA 导入会话密钥密文encryptedKey，并用privateKey解密。返回会话密钥。
func (PKCS1v15) ImportKeyRSA(keyBits uint32, privateKey *PrivateKey, encryptedKey []byte) ([]byte, error) {
	key, err := internal.DecryptPKCS1v15(grand.Reader, privateKey, encryptedKey)
	if err != nil || len(key) != int(keyBits) {
		return nil, common.SDR_INARGERR
	}

	return key, nil
}

// ExchangeDigitEnvelopeRSA 基于RSA的数字信封转换，privateKey解密会话密钥密文，再用publicKey加密会话密文。
// inKey: 输入会话密钥密文
// 返回: 转换后的会话密文
func (PKCS1v15) ExchangeDigitEnvelopeRSA(privateKey *PrivateKey, publicKey *PublicKey, inKey []byte) ([]byte, error) {
	key, err := internal.DecryptPKCS1v15(grand.Reader, privateKey, inKey)
	if err != nil {
		return nil, common.SDR_INARGERR
	}

	c, err := internal.EncryptPKCS1v15(grand.Reader, publicKey, key)
	if err != nil {
		return nil, common.SDR_KEYERR
	}

	return c, nil
}

// PublicKeyOperationRSA RSA公钥运算。数据格式由应用层封装。inData 应该小于 privateKey.Size()
func (PKCS1v15) PublicKeyOperationRSA(publicKey *PublicKey, inData []byte) ([]byte, error) {

	return internal.PublicKeyOperationRSA(publicKey, inData)
}

// PrivateKeyOperationRSA RSA私钥运算。数据格式由应用层封装。inData 应该小于 privateKey.Bits/8
func (PKCS1v15) PrivateKeyOperationRSA(privateKey *PrivateKey, inData []byte) ([]byte, error) {
	return internal.PrivateKeyOperationRSA(privateKey, inData)
}