package sm2

import (
	"bytes"
	"crypto/elliptic"
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"io"
	"log"
	"math/big"
	"testing"

	"github.com/stretchr/testify/assert"
	"xdx.jelly/xgcl/grand"
	"xdx.jelly/xgcl/internal"
)

func hexDecode(s string) []byte {
	b, err := hex.DecodeString(s)
	if err != nil {
		panic(err)
	}
	return b
}

func TestAudit(t *testing.T) {
	if !(Auditor{}).Correctness() {
		t.Fatal("Correctness check failed")
	}
}

// TestCurve test speed of sm2 curve on ScalarMult and ScalarBaseMult.
// A litter slower than P256 curve.
func TestCurve(t *testing.T) {
	gx, gy := Curve().Params().Gx, Curve().Params().Gy
	k := make([]byte, 32)
	grand.GenerateRandom(k)

	count, duation := internal.SingleThreadTester(func() {
		sm2Curve.ScalarMult(gx, gy, k)
	})
	fmt.Printf("SM2Curve ScalarMult: %d times/Sec\n", int(internal.Rate(count, duation)))

	count, duation = internal.SingleThreadTester(func() {
		sm2Curve.ScalarBaseMult(k)
	})
	fmt.Printf("SM2Curve ScalarBaseMult: %d times/Sec\n", int(internal.Rate(count, duation)))
}

// TestP256Curve test speed of P256 curve on ScalarMult and ScalarBaseMult.
func TestP256Curve(t *testing.T) {
	c := elliptic.P256()
	gx, gy := c.Params().Gx, c.Params().Gy
	k := make([]byte, 32)
	grand.GenerateRandom(k)

	count, duation := internal.SingleThreadTester(func() {
		c.ScalarMult(gx, gy, k)
	})
	fmt.Printf("P256 ScalarMult: %d times/Sec\n", int(internal.Rate(count, duation)))

	count, duation = internal.SingleThreadTester(func() {
		c.ScalarBaseMult(k)
	})
	fmt.Printf("P256 ScalarBaseMult: %d times/Sec\n", int(internal.Rate(count, duation)))
}

// func TestKeyGen(t *testing.T) {
// 	sk, _ := GenPrivateKey(grand.GetRandom(32))
// 	pk := GenPublicKey(sk)

// 	sm4Key := make([]byte, 16)
// 	buf := make([]byte, 0)
// 	fp, err := os.OpenFile("/Users/fengwd/Codes/go/src/xdx.jelly/xgcl/sm/sm2/key", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0666)
// 	if err != nil {
// 		panic("openfile failed")
// 	}
// 	for i := 0; i < 1000; i++ {
// 		buf = buf[:0]
// 		if i%1000 == 0 {
// 			fmt.Printf("%d\n", i)
// 		}
// 		usk, _ := GenPrivateKey(grand.GetRandom(32))
// 		upk := GenPublicKey(sk)
// 		grand.GenerateRandom(sm4Key)
// 		tmp, _ := usk.MarshalBinary()
// 		buf = append(buf, tmp...)
// 		tmp, _ = upk.MarshalBinary()
// 		buf = append(buf, tmp...)
// 		buf = padding.P7.Pad(buf, 16)
// 		buf, _ = sm4.EncryptECB(buf, sm4Key, buf)
// 		cipher, _ := Encrypt(pk, sm4Key, nil)
// 		tmp, _ = cipher.MarshalBinary()
// 		buf = append(buf, tmp...)
// 		fp.Write(buf)

// 	}

// 	fp.Close()
// }

func TestKeyValid(t *testing.T) {
	{
		k, _ := hex.DecodeString("000001000000000000000000000000000000000000000000000000000000000000000000972b80815e2c29e837fd006dc05534373950b221a090ad8dbeaa6e787bbbc85b00000000000000000000000000000000000000000000000000000000000000009ec521bc8fe61d2cfbc3bb8d4748ee6eb56e33d5955dba236f21085264095fa8")
		// pk := NewPublicKey()
		pk := &PublicKey{}
		pk.UnmarshalBinary(k)
		b, _ := pk.MarshalBinary()
		assert.Equal(t, k, b)
	}
	{
		k := NewPrivateKey().Random(grand.Reader)
		pk := (&PublicKey{}).Generate(k)
		b, _ := pk.MarshalBinary()
		pk1 := &PublicKey{}
		pk1.UnmarshalBinary(b)
		assert.True(t, pk1.Equal(pk))
		assert.True(t, pk.IsValid())
		assert.True(t, pk1.IsValid())
	}
}
func TestPreComputeWithIdAndPubkey(t *testing.T) {
	var r [32]byte
	r[31] = 1
	sk, _ := GenPrivateKey(r[:])
	pk := GenPublicKey(sk)
	e := PreComputeWithIdAndPubkey(nil, pk)
	assert.Equal(t, e, hexDecode("5b32bfe35482899b195d72c09d33ccdb465b2ded883240ff91f120a68bc91de8"))
}

func TestSignStd(t *testing.T) {
	m := []byte("message digest")
	buf, _ := hex.DecodeString("3945208F7B2144B13F36E38AC6D39F95889393692860B51A42FB81EF4DF7C5B8")
	d, _ := GenPrivateKey(buf)
	pk := GenPublicKey(d)
	e := PreComputeWithIdAndPubkeyAndMessage(nil, m, pk)
	k, err := hex.DecodeString("59276E27D506861A16680F3AD9C02DCCEF3CC1FA3CDBE4CE6D54B80DEAC1BC21")
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%x\n", e)
	fmt.Println(pk.X.Text(16))
	fmt.Println(pk.Y.Text(16))

	sig, err := Sign(e, k, d)
	if err != nil {
		t.Log()
		t.Fail()
	}

	data, err := sig.MarshalBinary()
	if err != nil {
		t.Log()
		t.Fail()
	}
	if hex.EncodeToString(data) != "0000000000000000000000000000000000000000000000000000000000000000f5a03b0648d2c4630eeac513e1bb81a15944da3827d5b74143ac7eaceee720b30000000000000000000000000000000000000000000000000000000000000000b1b6aa29df212fd8763182bc0d421ca1bb9038fd1f7f42d4840b69c485bbc1aa" {
		t.Log()
		t.Fail()
	}
	assert.True(t, Verify(e, pk, sig))
}

func TestSign(t *testing.T) {
	e := hexDecode("3854C463FA3F73783621B1CE4EF83F7C78048AAC79B221FCDD290866CC131174")
	k := hexDecode("F026AD9A7EB94401A800C8D8C3277E69972C7F3778ACE4D537012023EDFB69FF")
	d, err := GenPrivateKey(hexDecode("C242939DDAB6FCC07B6676C07D2DC117EC68A09142C25C008630B9756786162D"))
	assert.Nil(t, err)

	sig, err := Sign(e, k, d)
	assert.Nil(t, err)

	data, err := sig.MarshalBinary()
	assert.Nil(t, err)
	assert.Equal(t, data, hexDecode("00000000000000000000000000000000000000000000000000000000000000006e5db49dbd0992b97040080a96003c721cdb9cf64c88d74321fc2f630adf377400000000000000000000000000000000000000000000000000000000000000002f6dff453dfc8d7a506d3f52301bee529e62fddd38948f0d5d2cbcbc55900cfa"))

	sig1 := NewSignature()
	assert.Nil(t, sig1.UnmarshalBinary(data))
	assert.Equal(t, sig1.R.Cmp(sig.R), 0)
	assert.Equal(t, sig1.S.Cmp(sig.S), 0)

	pk := GenPublicKey(d)
	fmt.Println(pk)
	data, err = pk.MarshalBinary()
	assert.Nil(t, err)

	pk1 := NewPublicKey()
	assert.Nil(t, pk1.UnmarshalBinary(data))
	assert.Equal(t, fmt.Sprintf("%#v", pk), fmt.Sprintf("%#v", pk1))

	assert.True(t, Verify(e, pk, sig))
}

// (r,s)是e,pk的签名，则(r, -(s + 2rd/(1+d))也是e,pk的签名
func TestSignDual(t *testing.T) {
	e := hexDecode("3854C463FA3F73783621B1CE4EF83F7C78048AAC79B221FCDD290866CC131174")
	k := hexDecode("F026AD9A7EB94401A800C8D8C3277E69972C7F3778ACE4D537012023EDFB69FF")
	sk, _ := GenPrivateKey(hexDecode("C242939DDAB6FCC07B6676C07D2DC117EC68A09142C25C008630B9756786162D"))
	pk := GenPublicKey(sk)

	var sig *Signature
	sig, err := Sign(e, k, sk)
	assert.Nil(t, err)

	assert.True(t, Verify(e, pk, sig))

	// s' = -(s + 2rd(1+d)^-1)
	tmp := big.NewInt(1)
	tmp.Add(tmp, sk.D)
	tmp.ModInverse(tmp, OrderN())
	tmp.Mul(tmp, sk.D)
	tmp.Mul(tmp, sig.R)
	tmp.Lsh(tmp, 1)
	tmp.Add(tmp, sig.S)
	tmp.Mod(tmp, OrderN())
	tmp.Sub(OrderN(), tmp)
	sig.S.Set(tmp)

	assert.True(t, Verify(e, pk, sig))
}

// 已知公钥P, 则可任意取r,s,并构造e = r-X(sG+(r+s)P)
// verify(r,s,P,e) = true
func TestSignForge(t *testing.T) {

	d, err := hex.DecodeString("C242939DDAB6FCC07B6676C07D2DC117EC68A09142C25C008630B9756786162D")
	if err != nil {
		log.Fatal(err)
	}
	sk, _ := GenPrivateKey(d)
	pk := GenPublicKey(sk)

	r, _ := rand.Int(grand.Reader, OrderN())
	s, _ := rand.Int(grand.Reader, OrderN())

	u := big.NewInt(0)
	u.Add(r, s)
	u.Mod(u, OrderN())
	x1, y1 := Curve().ScalarMult(pk.X, pk.Y, u.Bytes())
	x2, y2 := Curve().ScalarBaseMult(s.Bytes())
	x1, _ = Curve().Add(x1, y1, x2, y2)
	e := new(big.Int)
	e = e.Sub(r, x1)
	e.Mod(e, OrderN())

	sig := new(Signature)
	sig.S = s
	sig.R = r

	assert.True(t, Verify(e.Bytes(), pk, sig))

}

func TestEncrypt(t *testing.T) {
	sk, _ := GenPrivateKey([]byte{
		0x81, 0x98, 0x7C, 0xC9, 0x0C, 0xF5, 0x05, 0x7C,
		0x2D, 0xCA, 0xA7, 0x5D, 0x1F, 0xDD, 0xCA, 0x84,
		0xB2, 0x48, 0x62, 0xF0, 0xCA, 0xD7, 0x3C, 0x7F,
		0x67, 0x34, 0x9A, 0xE6, 0x99, 0xB9, 0x29, 0x83},
	)

	rnd := []byte{
		0x26, 0xD6, 0x16, 0x3F, 0xA1, 0x86, 0x03, 0xEE,
		0x2F, 0x3D, 0xE8, 0x93, 0x65, 0x44, 0xD5, 0xDF,
		0x12, 0x55, 0xA2, 0xDB, 0xEB, 0xA6, 0x3A, 0xFC,
		0x0D, 0x83, 0x3E, 0xC4, 0x49, 0xD2, 0xCB, 0x45,
	}

	msg := []byte{
		0xC3, 0x53, 0xC6, 0x8E, 0xF0, 0x5C, 0x4B, 0x34,
		0x2B, 0x37, 0x7D, 0xA0, 0x55, 0xD9, 0x09, 0xFB,
		0x1F, 0xAA, 0x42, 0x55, 0x66, 0x2F, 0x3B, 0xAB,
		0x8D, 0xDB, 0x35, 0x35, 0xE4, 0x0B, 0xC9, 0x3B,
	}

	wantCipher := "bf6cfcb8e6295dc22777376f8385c5d6aadd5e430d11e004246d6bebf99ec5249cb9ab2f9af688c77a1bdf9f3b0816a4eab7f5da22e5dacdc1c8f6e45499874e1fc32e35744161aa0ffa6c70fc811d3b66d4cacda3c0996b54768c603c6b24e0c85cdde8ad71a258b89ddb42da900bcf4f18ab52d7841134cac581d3cf7f58f7"
	pk := GenPublicKey(sk)

	// pk.X.Add(pk.X, gmath.BigInt1)
	// fmt.Println(hex.EncodeToString(pk.x[:]) + hex.EncodeToString(pk.y[:]))
	cipher, err := Encrypt(pk, msg, rnd)
	if err != nil {
		t.Log(err)
		t.Fail()
		return
	}

	if hex.EncodeToString(cipher.Bytes()) != wantCipher {
		t.Log("Encrypt Failed")
		t.Log("cipher    :", hex.EncodeToString(cipher.Bytes()))
		t.Log("WantCipher:", wantCipher)
		t.Fail()
		return
	}

	// marshal test
	// t.Log(cipher)
	data, _ := cipher.MarshalBinary()
	// t.Log(hex.EncodeToString(data))
	cipher1 := NewCipher()
	if err := cipher1.UnmarshalBinary(data); err != nil {
		t.Log("unmarshal failed")
	}
	if fmt.Sprintf("%#v", cipher) != fmt.Sprintf("%#v", cipher1) {
		t.Log("Marshal Test failed")
		t.Fail()
	}

	//Decrypt
	decryptedMsg, err := Decrypt(sk, cipher)
	if err != nil || bytes.Compare(decryptedMsg, msg) != 0 {
		t.Log(err)
		t.Log("decryptedMsg    :", hex.EncodeToString(decryptedMsg))
		t.Log("msg    :", hex.EncodeToString(msg))
		t.Fail()
	}

}

// func TestErrors(t *testing.T) {
// 	errors.PrintErrorInfos(os.Stdout)
// }

func TestZeroHashSign(t *testing.T) {
	zeroHash := make([]byte, 32)

	for _, curve := range []elliptic.Curve{Curve()} {
		privKey, err := GenerateKey(curve, rand.Reader)
		if err != nil {
			panic(err)
		}

		// Sign a hash consisting of all zeros.
		// r, s, err := Sign(rand.Reader, privKey, zeroHash)
		k := make([]byte, 32)
		grand.GenerateRandom(k)
		sig, err := Sign(zeroHash, k, privKey)
		if err != nil {
			panic(err)
		}

		// Confirm that it can be verified.
		if !Verify(zeroHash, privKey.Public().(*PublicKey), sig) {
			t.Errorf("zero hash signature verify failed for %T", curve)
		}
	}
}

func TestEncEmpty(t *testing.T) {
	sk, _ := GenPrivateKey([]byte{
		0x81, 0x98, 0x7C, 0xC9, 0x0C, 0xF5, 0x05, 0x7C,
		0x2D, 0xCA, 0xA7, 0x5D, 0x1F, 0xDD, 0xCA, 0x84,
		0xB2, 0x48, 0x62, 0xF0, 0xCA, 0xD7, 0x3C, 0x7F,
		0x67, 0x34, 0x9A, 0xE6, 0x99, 0xB9, 0x29, 0x83},
	)

	rnd := []byte{
		0x26, 0xD6, 0x16, 0x3F, 0xA1, 0x86, 0x03, 0xEE,
		0x2F, 0x3D, 0xE8, 0x93, 0x65, 0x44, 0xD5, 0xDF,
		0x12, 0x55, 0xA2, 0xDB, 0xEB, 0xA6, 0x3A, 0xFC,
		0x0D, 0x83, 0x3E, 0xC4, 0x49, 0xD2, 0xCB, 0x45,
	}

	msg := []byte{}

	pk := &sk.PublicKey
	cipher, err := Encrypt(pk, msg, rnd)
	if err != nil {
		t.Fatal(err)
	}
	decryptedMsg, err := Decrypt(sk, cipher)
	if err != nil || bytes.Compare(decryptedMsg, msg) != 0 {
		if err != nil {
			t.Fatal(err)
		}
	}
}

func TestEncLong(t *testing.T) {
	sk, err := GenerateKey(Curve(), grand.Reader)
	pk := &sk.PublicKey

	if err != nil {
		t.Fatal(err)
	}

	data := make([]byte, 409600)
	if _, err := io.ReadFull(grand.Reader, data); err != nil {
		t.Fatal(err)
	}

	cipher, err := Encrypt(pk, data, grand.GetRandom(32))
	if err != nil {
		t.Fatal(err)
	}
	decryptedData, err := Decrypt(sk, cipher)
	if err != nil || bytes.Compare(decryptedData, data) != 0 {
		if err != nil {
			t.Fatal(err)
		}
	}
}

func TestSM2Point(t *testing.T) {
	x, _ := new(big.Int).SetString("54cb2e421e7119ddbba44ff0db131e9907ed90a1b0b9a23b9d3eb890519d5cae", 16)
	y, _ := new(big.Int).SetString("389646ab568a3bd7a566e1e49b98a4d41041b2fa2cf0af74fceef65ef1e986c6", 16)

	pk := &PublicKey{
		Curve: Curve(),
		X:     x,
		Y:     y,
	}
	assert.True(t, pk.IsValid())
}

// func TestSM2Verify(t *testing.T) {
// 	r, _ := new(big.Int).SetString("d3233fe6b680f1fc26bb4932b894af1e9e73c700554c39f407ac91c4d4d74347", 16)
// 	s, _ := new(big.Int).SetString("08bc7476b0fd496287d162dfab5a68a340515b3f8647ed03f75c4d1b50aeb957", 16)
// 	x, _ := new(big.Int).SetString("6a3f992330c0857eda4b914ebef614bf401266fb6168e3f90475ed1fcd3c2625", 16)
// 	y, _ := new(big.Int).SetString("36341b6383d996205996aa164480bad952aa07a0db942858d103bf8f2f4b460a", 16)
// 	pk := &PublicKey{
// 		Curve: Curve(),
// 		X:     x,
// 		Y:     y,
// 	}
// 	sig := &Signature{
// 		R: r,
// 		S: s,
// 	}
// 	assert.True(t, pk.IsValid())
// 	e, _ := hex.DecodeString("c3e283e5ad232e04b29d2422a88fb251cd08cad4f0e8b77cdcdd7859eb6fabd7")
// 	// e = PreComputeWithIdAndPubkeyAndMessage(GetDefaultID(), e, pk)

// 	assert.True(t, Verify(e, pk, sig))
// 	e, _ = hex.DecodeString("8d3eef66a23e3118a7b32a386cd1db41bc4e7a803098ee92c6b5cda43829ad96")
// 	// e = PreComputeWithIdAndPubkeyAndMessage(GetDefaultID(), e, pk)
// 	assert.True(t, Verify(e, pk, sig))
// }