package outsource

import (
	"math/big"
	"testing"

	"github.com/stretchr/testify/assert"
	"xdx.jelly/xgcl/gmath"
	"xdx.jelly/xgcl/grand"
	"xdx.jelly/xgcl/he/paillier"
	"xdx.jelly/xgcl/sm/sm2"
	"xdx.jelly/xgcl/sm/sm3"
	"xdx.jelly/xgcl/tpc/sm2/sm2m"
)

func BenchmarkOS(b *testing.B) {
	// O、客户端生成paillier密钥
	salt := grand.GetRandom(16)
	paiPrivKey, err := paillier.GenerateKeyFromPassword(2048, []byte("password123"), salt, 1024)
	assert.Nil(b, err)
	paiPubKey := paiPrivKey.Public()

	// 一、密钥生成

	// 1）客户端第一步, a1, a2 发外包服务器
	clientKeyGen := &ClientKeyGenerator{}
	a1, a2, err := clientKeyGen.Step1(paiPubKey, grand.Reader)
	assert.Nil(b, err)

	// 2) 外包服务器第一步
	// 保存encryptedClientKey1，encryptedClientKey2,
	// 把 P 发客户端
	// encryptedClientKey1 -签名用
	// encryptedClientKey2 -解密加密密钥保护结构用
	osKenGen := &OSKeyGenerator{}
	encryptedClientKey1, _, P, err := osKenGen.Step1(a1, a2, paiPubKey, grand.Reader)
	assert.Nil(b, err)

	// 3) 客户端第二步
	// clientTempKey发协同签名服务端
	clientTempKey, err := clientKeyGen.Step2(P)

	// 4）协同签名服务器
	// serverTempKey 发客户端，保存serverKey-服务端密钥分量
	serverKey, serverTempKey, publicKey, _ := sm2m.ServerGenSignKey(clientTempKey, grand.GetRandom(32))

	// 5）客户端第三步
	// 把(serverTempKey,S)发送给外包服务器
	S, err := clientKeyGen.Step3(serverTempKey)
	assert.Nil(b, err)

	// 6）外包服务器第二步
	// 把T， osPublicKey发给客户端。保存osPublicKey为用户签名公钥
	T, osPublicKey, err := osKenGen.Step2(serverTempKey, S)
	assert.Nil(b, err)

	// 7）客户端第四步
	// clientPublicKey, 客户端生成的公钥，如果没有返错，则应与ocPublicKey一致。

	clientPublicKey, err := clientKeyGen.Step4(T)
	assert.Nil(b, err)
	assert.True(b, clientPublicKey.Equals(osPublicKey))

	e := grand.GetRandom(32)
	// 签名
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		// 1）客户端发起请求

		// 2）外包方计算
		b.StartTimer()
		outsourcintCtx := new(OSSignContext)
		PPrime, err := outsourcintCtx.Step1(grand.Reader) // P'
		assert.Nil(b, err)

		// outsourcintCtx把PPrime发给客户端，保存outsourcintCtx.Marshal()

		// 3) 客户端组合数据data = (e,p)=e||px||py并发送给协同服务端
		b.StopTimer()
		data := make([]byte, sm3.Size+2*sm2.ByteSize())
		pos := copy(data, e)
		pos += copy(data[pos:], gmath.BigIntToNByte(PPrime.X, sm2.ByteSize()))
		copy(data[pos:], gmath.BigIntToNByte(PPrime.Y, sm2.ByteSize()))

		// 4) 协同服务端计算,发回data
		b.StartTimer()
		data, err = sm2m.ServerSign(serverKey, data, grand.Reader)
		assert.Nil(b, err)

		// 5) 客户端解析data = r || s1 || s2, 把s1, s2发给外包服务器。
		b.StopTimer()
		r := new(big.Int)
		r.SetBytes(data[:sm2.ByteSize()])
		s1 := new(big.Int).SetBytes(data[sm2.ByteSize() : 2*sm2.ByteSize()])
		s2 := new(big.Int).SetBytes(data[2*sm2.ByteSize():])

		// 6) 外包服务器解析data并计算c,把c发送给客户端
		b.StartTimer()
		c, err := outsourcintCtx.Step2(s1, s2, encryptedClientKey1, paiPubKey)
		assert.Nil(b, err)

		// 7) 客户端计算签名值
		b.StopTimer()
		s, _ := paillier.Decrypt(c, paiPrivKey)
		s.Sub(s, r)
		s.Mod(s, sm2.OrderN())
		sig := &sm2.Signature{
			R: r,
			S: s,
		}
		// 8) 客户端验证签名
		assert.True(b, sm2.Verify(e, publicKey, sig))

	}
}