package sm2a

import (
	"bytes"
	"fmt"
	"testing"
	"time"

	"xdx.jelly/xgcl/grand"
)

func TestExchangeBoth(t *testing.T) {
	// 发起方密钥生成
	vv := false
	sponsorClientKeyCtx := NewClientEncKeyGenContext()
	out, _ := sponsorClientKeyCtx.ClientKeyGen_one(grand.Reader)
	sponsorServerKey, sponsorPublicKey, out, _ := ServerEncKeyGen(out, grand.Reader)
	sponsorClientKeyCtx.ClientKeyGen_two(out)
	sponsorClientKey := sponsorClientKeyCtx.ClientKey
	if vv {
		printLog("发起方客户端私钥分量=", sponsorClientKeyCtx.ClientKey.Bytes())
		printLog("发起方服务端私钥分量=", sponsorServerKey.Bytes())
		printLog("发起方公钥=", sponsorClientKeyCtx.Pubkey.Bytes())
	}
	// 响应方密钥生成
	responsorClientKeyCtx := NewClientEncKeyGenContext()
	out, _ = responsorClientKeyCtx.ClientKeyGen_one(grand.Reader)
	responsorServerKey, responsorPublicKey, out, _ := ServerEncKeyGen(out, grand.Reader)
	responsorClientKeyCtx.ClientKeyGen_two(out)
	responsorClientKey := responsorClientKeyCtx.ClientKey
	if vv {
		printLog("响应方客户端私钥分量=", responsorClientKeyCtx.ClientKey.Bytes())
		printLog("响应方服务端私钥分量=", responsorServerKey.Bytes())
		printLog("响应方公钥=", responsorClientKeyCtx.Pubkey.Bytes())
	}

	sponsor := []byte("alice")
	responsor := []byte("bob")
	keyLength := 16

	cs := NewClientSponsor(sponsor, sponsorClientKey, sponsorPublicKey)
	cr := NewClientResponsor(responsor, responsorClientKey, responsorPublicKey)
	tempKeyOfSponsor, err := cs.GenerateAgreementData(grand.GetRandom(32))
	if err != nil {
		t.Fatal("GenerateAgreementData failed")
	}
	if vv {
		printLog("Ra=", tempKeyOfSponsor.Bytes())
	}
	// 协同计算GenerateAgreementDataAndKey
	toServer, err := cr.GenerateAgreementDataAndKey_1of2(sponsorPublicKey, tempKeyOfSponsor)
	if err != nil {
		t.Fatal("cr.GenerateAgreementDataAndKey_1of2 failed")
	}
	if vv {
		printLog("Qa=", toServer)
	}

	toClient, err := ServerKeyExchange(toServer, responsorServerKey)
	if err != nil {
		t.Fatal("cr.GenerateAgreementDataAndKey_1of2 failed")
	}
	if vv {
		printLog("[h*db_s]Qa=", toClient)
	}

	keyOfResponsor, tempKeyOfResponsor, err := cr.GenerateAgreementDataAndKey_2of2(keyLength, sponsor, sponsorPublicKey, toClient, grand.GetRandom(32))
	if err != nil {
		t.Fatal("cr.GenerateAgreementDataAndKey_2of2 failed")
	}
	if vv {
		printLog("Rb=", tempKeyOfResponsor.Bytes())
	}
	toServer, err = cs.GenerateKey_1of2(responsorPublicKey, tempKeyOfResponsor)
	if err != nil {
		t.Fatal("cs.GenerateKey_1of2 failed")
	}
	if vv {
		printLog("Qb=", toServer)
	}
	toClient, err = ServerKeyExchange(toServer, sponsorServerKey)
	if err != nil {
		t.Fatal("cs.GenerateKey_1of2 failed")
	}
	if vv {
		printLog("[h*db_s]Qb=", toClient)
	}
	keyOfSponsor, err := cs.GenerateKey_2of2(keyLength, responsor, toClient, responsorPublicKey, sponsorClientKey)
	if err != nil {
		t.Fatal("cs.GenerateKey_2of2 failed")
	}
	printLog("ka=", keyOfSponsor)
	if vv {
		fmt.Printf("响应方密钥: %X\n", keyOfResponsor)
		fmt.Printf("发起方密钥: %X\n", keyOfSponsor)
	}
	if bytes.Compare(keyOfResponsor, keyOfSponsor) != 0 {
		t.Fatal("Key agreement failed")
	}
}

func TestExchangeBenchmark(t *testing.T) {
	// 发起方密钥生成
	sponsorClientKeyCtx := NewClientEncKeyGenContext()
	out, _ := sponsorClientKeyCtx.ClientKeyGen_one(grand.Reader)
	sponsorServerKey, sponsorPublicKey, out, _ := ServerEncKeyGen(out, grand.Reader)
	sponsorClientKeyCtx.ClientKeyGen_two(out)
	sponsorClientKey := sponsorClientKeyCtx.ClientKey

	// 响应方密钥生成
	responsorClientKeyCtx := NewClientEncKeyGenContext()
	out, _ = responsorClientKeyCtx.ClientKeyGen_one(grand.Reader)
	responsorServerKey, responsorPublicKey, out, _ := ServerEncKeyGen(out, grand.Reader)
	responsorClientKeyCtx.ClientKeyGen_two(out)
	responsorClientKey := responsorClientKeyCtx.ClientKey

	sponsor := []byte("alice")
	responsor := []byte("bob")
	keyLength := 16

	cnt := 1000
	errors := 0
	start := time.Now()
	for i := 0; i < cnt; i++ {

		cs := NewClientSponsor(sponsor, sponsorClientKey, sponsorPublicKey)
		cr := NewClientResponsor(responsor, responsorClientKey, responsorPublicKey)
		tempKeyOfSponsor, err := cs.GenerateAgreementData(grand.GetRandom(32))
		if err != nil {
			t.Fatal("GenerateAgreementData failed")
		}

		// 协同计算GenerateAgreementDataAndKey
		toServer, err := cr.GenerateAgreementDataAndKey_1of2(sponsorPublicKey, tempKeyOfSponsor)
		if err != nil {
			t.Fatal("cr.GenerateAgreementDataAndKey_1of2 failed")
		}
		toClient, err := ServerKeyExchange(toServer, responsorServerKey)
		if err != nil {
			t.Fatal("cr.GenerateAgreementDataAndKey_1of2 failed")
		}
		keyOfResponsor, tempKeyOfResponsor, err := cr.GenerateAgreementDataAndKey_2of2(keyLength, sponsor, sponsorPublicKey, toClient, grand.GetRandom(32))
		if err != nil {
			t.Fatal("cr.GenerateAgreementDataAndKey_2of2 failed")
		}

		toServer, err = cs.GenerateKey_1of2(responsorPublicKey, tempKeyOfResponsor)
		if err != nil {
			t.Fatal("cs.GenerateKey_1of2 failed")
		}
		toClient, err = ServerKeyExchange(toServer, sponsorServerKey)
		if err != nil {
			t.Fatal("cs.GenerateKey_1of2 failed")
		}
		keyOfSponsor, err := cs.GenerateKey_2of2(keyLength, responsor, toClient, responsorPublicKey, sponsorClientKey)
		if err != nil {
			t.Fatal("cs.GenerateKey_2of2 failed")
		}

		if bytes.Compare(keyOfResponsor, keyOfSponsor) != 0 {
			errors += 1
		}

	}
	end := time.Now()
	elapsed := end.Sub(start)
	fmt.Printf("总运行次数：%d\n", cnt)
	fmt.Printf("总失败次数：%d\n", errors)
	fmt.Printf("错误率：%f\n", float64(errors)/float64(cnt))
	fmt.Printf("耗时时间（单位：秒）：%f\n", elapsed.Seconds()*20)
	fmt.Printf("单次请求平均耗时（单位：秒）：%f\n", elapsed.Seconds()/float64(cnt)*20)
	fmt.Printf("TPS：%f\n", float64(cnt)/elapsed.Seconds()/20)

}