package sm2a

import (
	"fmt"
	"math/rand"
	"sync"
	"testing"
	"time"

	"xdx.jelly/xgcl/gmath"
	"xdx.jelly/xgcl/grand"
	"xdx.jelly/xgcl/sm/sm2"
)

func TestGenSignKeyRaw(t *testing.T) {
	serverKeyCtx := NewServerSignKeyGenContext()
	clientKeyCtx := NewClientSignKeyGenContext(grand.Reader)

	buf, err := serverKeyCtx.ServerGenKey_one(grand.Reader)
	if err != nil {
		fmt.Println(err)
		return
	}
	serverctx, err := serverKeyCtx.MarshalBinary()
	if err != nil {
		t.Log(err)
		t.Fail()
		return
	}
	buf, err = clientKeyCtx.ClientKeyGen_one(buf)
	if err != nil {
		fmt.Println(err)
		return
	}
	_ = serverKeyCtx.UnmarshalBinary(serverctx)
	buf, err = serverKeyCtx.ServerGenKey_two(buf, grand.Reader)
	if err != nil {
		fmt.Println(err)
		return
	}
	serverctx, _ = serverKeyCtx.MarshalBinary()

	buf, err = clientKeyCtx.ClientKeyGen_two(buf)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println("Client public Key = ", clientKeyCtx.PubKey)

	_ = serverKeyCtx.UnmarshalBinary(serverctx)
	err = serverKeyCtx.ServerGenKey_three(buf)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println("Server public Key = ", serverKeyCtx.PubKey)

	dc, _ := clientKeyCtx.GetClientKey()
	ds, _ := serverKeyCtx.GetServerKey()
	pkc, _ := clientKeyCtx.GetPublicKey()
	pks, _ := serverKeyCtx.GetPublicKey()
	if !pkc.Equals(pks) {
		t.Fatal("Public keys of server and client are not equal")
	}

	// check (1+d)^{-1} = dc+ds
	d := sm2.NewPrivateKey()
	d.D.Add(dc.D, ds.D)
	d.D.ModInverse(d.D, sm2.OrderN())
	d.D.Sub(d.D, gmath.BigInt1)
	px, py := sm2.Curve256.ScalarBaseMult(d.Bytes())
	fmt.Println("[d]*G=(", px.Text(16)+", "+py.Text(16)+")")
	if px.Cmp(pkc.X) != 0 || py.Cmp(pkc.Y) != 0 {
		t.Fatal("Private keys of server and client are not march")
	}

	d.D.Mul(clientKeyCtx.ClientSubKey.D, serverKeyCtx.ServerSubKey.D)
	d.D.Sub(d.D, gmath.BigInt1)
	d.D.Mod(d.D, sm2.OrderN())
	px, py = sm2.Curve256.ScalarBaseMult(d.Bytes())
	fmt.Println("[d]*G=(", px.Text(16)+", "+py.Text(16)+")")
	if px.Cmp(pkc.X) != 0 || py.Cmp(pkc.Y) != 0 {
		t.Fatal("Private keys of server and client are not march")
	}
}

func TestGenSignKeyTps(t *testing.T) {
	rand1 := rand.New(rand.NewSource(time.Now().UnixNano()))
	wg := sync.WaitGroup{}

	totalSuccess := 0
	totalFailed := 0

	var (
		fail  float32
		total float32
	)

	beg := time.Now().UnixNano()

	for j := 0; j < 10; j++ {
		wg.Add(1)

		wtt := time.Duration(rand1.Int() % 100)
		time.Sleep(wtt * time.Millisecond)

		go func(wg *sync.WaitGroup, t *testing.T) {
			defer wg.Done()

			for mi := 0; mi < 20; mi++ {
				serverKeySeg := NewServerSignKeyGenContext()
				clientKeySeg := NewClientSignKeyGenContext(grand.Reader)

				buf, err := serverKeySeg.ServerGenKey_one(grand.Reader)
				if err != nil {
					fmt.Println(err)
					return
				}

				serverctx, _ := serverKeySeg.MarshalBinary()
				fmt.Printf("%d\n", len(serverctx))
				//rand.Reader will got error? why

				buf, err = clientKeySeg.ClientKeyGen_one(buf)
				if err != nil {
					fmt.Println(err)
					return
				}

				_ = serverKeySeg.UnmarshalBinary(serverctx)
				buf, err = serverKeySeg.ServerGenKey_two(buf, grand.Reader)
				if err != nil {
					fmt.Println(err)
					return
				}
				serverctx, _ = serverKeySeg.MarshalBinary()
				fmt.Printf("%X\n", serverctx)

				buf, err = clientKeySeg.ClientKeyGen_two(buf)
				if err != nil {
					fmt.Println(err)
					return
				}
				fmt.Println("Client public Key = ", clientKeySeg.PubKey)

				_ = serverKeySeg.UnmarshalBinary(serverctx)
				err = serverKeySeg.ServerGenKey_three(buf)
				if err != nil {
					fmt.Println(err)
					return
				}
				fmt.Println("Server public Key = ", serverKeySeg.PubKey)

				dc, _ := clientKeySeg.GetClientKey()
				ds, _ := serverKeySeg.GetServerKey()
				pkc, _ := clientKeySeg.GetPublicKey()
				pks, _ := serverKeySeg.GetPublicKey()
				if !pkc.Equals(pks) {
					t.Fatal("Public keys of server and client are not equal")
				}

				// check (1+d)^{-1} = dc+ds
				d := sm2.NewPrivateKey()
				d.D.Add(dc.D, ds.D)
				d.D.ModInverse(d.D, sm2.OrderN())
				d.D.Sub(d.D, gmath.BigInt1)
				px, py := sm2.Curve256.ScalarBaseMult(d.Bytes())
				fmt.Println("p=(", px.Text(16)+", "+py.Text(16)+")")

				if px.Cmp(pkc.X) != 0 || py.Cmp(pkc.Y) != 0 {
					t.Fatal("Private keys of server and client are not march")
				}

				totalSuccess += 1
			}
		}(&wg, t)

	}

	wg.Wait()
	end := time.Now().UnixNano()

	fail = float32(totalFailed)
	total = float32(totalSuccess + totalFailed)
	rate := fail / total

	elapseDenom := float32(end - beg)
	elapse := elapseDenom / 1000000000
	one := elapse / total

	fmt.Printf("错误率: %f %% \n", rate*100)
	fmt.Printf("耗时时间(单位:秒): %f seconds \n", elapse)
	fmt.Printf("单次请求平均耗时(单位:秒) %f seconds \n", one)
	fmt.Println("TPS: ", total/elapse)
	fmt.Println("")
	fmt.Println("")
	fmt.Println("")

}

func TestEncKeyGen(t *testing.T) {
	clientKeyCtx := NewClientEncKeyGenContext()

	buf, _ := clientKeyCtx.ClientKeyGen_one(grand.Reader)
	printLog("客户端第一步输出: ", buf)

	ds, pk, buf, _ := ServerEncKeyGen(buf, grand.Reader)
	printLog("服务端第一步输出: ", buf)

	clientKeyCtx.ClientKeyGen_two(buf)
	dc := clientKeyCtx.ClientKey

	printLog("SM2加密密钥生成, 服务端生成公钥", pk.Bytes())
	printLog("SM2加密密钥生成, 客户端生成公钥", clientKeyCtx.Pubkey.Bytes())
	printLog("SM2加密密钥生成, 服务端私钥分量", ds.Bytes())
	printLog("SM2加密密钥生成, 客户端私钥分量", dc.Bytes())

	// 验证 [ds+dc]·G = pk
	d := sm2.NewPrivateKey()
	d.D.Add(dc.D, ds.D)
	d.D.Mod(d.D, sm2.OrderN())
	px, py := sm2.Curve256.ScalarBaseMult(d.Bytes())
	if px.Cmp(pk.X) != 0 || py.Cmp(pk.Y) != 0 {
		t.Fatal("Private keys of server and client are not march")
	}
}

func TestGenSignKey(t *testing.T) {
	serverKeyCtx := NewServerSignKeyGenContext()
	clientKeyCtx := NewClientSignKeyGenContext(grand.Reader)

	buf, err := serverKeyCtx.ServerGenKey_one(grand.Reader)
	if err != nil {
		fmt.Println(err)
		return
	}

	printLog("服务端第一步输出: ", buf)

	buf, err = clientKeyCtx.ClientKeyGen_one(buf)
	if err != nil {
		fmt.Println(err)
		return
	}

	printLog("客户端第一步输出: ", buf)

	buf, err = serverKeyCtx.ServerGenKey_two(buf, grand.Reader)
	if err != nil {
		fmt.Println(err)
		return
	}
	printLog("服务端第二步输出: ", buf)

	buf, err = clientKeyCtx.ClientKeyGen_two(buf)
	if err != nil {
		fmt.Println(err)
		return
	}
	printLog("客户端第二步输出: ", buf)

	err = serverKeyCtx.ServerGenKey_three(buf)
	if err != nil {
		fmt.Println(err)
		return
	}

	dc, _ := clientKeyCtx.GetClientKey()
	ds, _ := serverKeyCtx.GetServerKey()
	pkc, _ := clientKeyCtx.GetPublicKey()
	pks, _ := serverKeyCtx.GetPublicKey()
	if !pkc.Equals(pks) {
		t.Fatal("错误：客户端和服务端生成的公钥不相等")
	}

	fmt.Printf("客户端生成密钥分量 = %X\n", dc.Bytes())
	fmt.Printf("服务端生成密钥分量 = %X\n", ds.Bytes())

	fmt.Printf("客户端生成公钥 = %X\n", clientKeyCtx.PubKey.Bytes())
	fmt.Printf("服务端生成公钥 = %X\n", serverKeyCtx.PubKey.Bytes())

	// 验证密钥配对 (1+d)^{-1} = dc+ds
	d := sm2.NewPrivateKey()
	d.D.Add(dc.D, ds.D)
	d.D.ModInverse(d.D, sm2.OrderN())
	d.D.Sub(d.D, gmath.BigInt1)
	px, py := sm2.Curve256.ScalarBaseMult(d.Bytes())
	if px.Cmp(pkc.X) != 0 || py.Cmp(pkc.Y) != 0 {
		t.Fatal("错误：客户端和服务端密钥不配对")
	}

}