package paillier

import (
	"crypto/rand"
	"encoding/hex"
	"math/big"
	"testing"
	"time"

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

func bigFromBase16(s string) *big.Int {
	b, err := hex.DecodeString(s)
	if err != nil {
		panic(err)
	}
	return new(big.Int).SetBytes(b)
}

var p = bigFromBase16("ff03b1a74827c746db83d2eaff00067622f545b62584321256e62b01509f1096" +
	"2f9c5c8fd0b7f5184a9ce8e81f439df47dda14563dd55a221799d2aa57ed2713" +
	"271678a5a0b8b40a84ad13d5b6e6599e6467c670109cf1f45ccfed8f75ea3b81" +
	"4548ab294626fe4d14ff764dd8b091f11a0943a2dd2b983b0df02f4c4d00b413")

var q = bigFromBase16("dacaabc1dc57faa9fd6a4274c4d588765a1d3311c22e57d8101431b07eb3ddcb" +
	"05d77d9a742ac2322fe6a063bd1e05acb13b0fe91c70115c2b1eee1155e07252" +
	"7011a5f849de7072a1ce8e6b71db525fbcda7a89aaed46d27aca5eaeaf35a262" +
	"70a4a833c5cda681ffd49baa0f610bad100cdf47cc86e5034e2a0b2179e04ec7")

var n = new(big.Int).Mul(p, q)
var lambda = new(big.Int)

// lcm(lambda, new(big.Int).Sub(p, gmath.BigInt1), new(big.Int).Sub(q, gmath.BigInt1))

func TestPaillier(t *testing.T) {
	// pk, sk := GenerateKey(1024)

	sk, err := GenerateKeyFromPassword(2048, grand.GetRandom(32), grand.GetRandom(32), 1024)
	assert.Nil(t, err)
	pk := sk.Public()

	m1, _ := rand.Int(grand.Reader, pk.N)
	m2, _ := rand.Int(grand.Reader, pk.N)
	c1, _ := Encrypt(m1, pk, grand.Reader)
	c2, _ := Encrypt(m2, pk, grand.Reader)

	c := new(Cipher).HomomorphicAdd(c1, c2, pk)
	m, _ := Decrypt(c, sk)

	mm := new(big.Int).Add(m1, m2)
	mm.Mod(mm, pk.N)
	assert.Zero(t, m.Cmp(mm))
}

func TestKeyMarshal(t *testing.T) {

	sk, pk, err := GenerateKey(2048, grand.Reader)
	assert.Nil(t, err)

	assert.Equal(t, pk.Size(), uint(2048))

	skBytes, err := sk.Marshal()
	assert.Nil(t, err)
	sk2 := new(PrivateKey)
	assert.Nil(t, sk2.Unmarshal(skBytes))
	assert.Zero(t, sk.p.Cmp(sk2.p)|sk.q.Cmp(sk2.q)|sk.lambda.Cmp(sk2.lambda)|sk.lambdaInv.Cmp(sk2.lambdaInv))

	pkBytes, err := pk.Marshal()
	assert.Nil(t, err)
	pk2 := new(PublicKey)
	assert.Nil(t, pk2.Unmarshal(pkBytes))
	assert.Zero(t, pk.N.Cmp(pk2.N)|pk.N2.Cmp(pk2.N2))
}

func TestPaillierScalarMul(t *testing.T) {
	sk, err := GenerateKeyFromPassword(2048, grand.GetRandom(32), grand.GetRandom(32), 1024)
	assert.Nil(t, err)
	pk := sk.Public()

	limit := big.NewInt(1)
	limit.Lsh(limit, 256)
	ch := time.After(3 * time.Second)
	done := false
	for !done {
		select {
		case <-ch:
			done = true
		default:
			m, _ := rand.Int(grand.Reader, limit)
			k, _ := rand.Int(grand.Reader, limit)

			c, _ := Encrypt(m, pk, grand.Reader)
			c.HomomorphicScalarMul(c, k, sk.Public())

			m1, err := Decrypt(c, sk)
			assert.Nil(t, err)
			m2 := new(big.Int).Mul(m, k)
			assert.Zero(t, m1.Cmp(m2))
		}
	}
}

func BenchmarkParallelPaillierEnc(b *testing.B) {
	_, pk, _ := GenerateKey(2048, grand.Reader)

	m, _ := rand.Int(grand.Reader, pk.N)
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			Encrypt(m, pk, grand.Reader)
		}
	})
}

func BenchmarkPaillierEnc(b *testing.B) {
	_, pk, _ := GenerateKey(2048, grand.Reader)

	m, _ := rand.Int(grand.Reader, pk.N)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		Encrypt(m, pk, grand.Reader)
	}
}

func BenchmarkPaillierDec(b *testing.B) {
	sk, pk, _ := GenerateKey(2048, grand.Reader)

	m, _ := rand.Int(grand.Reader, pk.N)
	c, _ := Encrypt(m, pk, grand.Reader)
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		Decrypt(c, sk)
	}
}

func BenchmarkPaillierHomomorphicScalarMul(b *testing.B) {
	sk, pk, _ := GenerateKey(2048, grand.Reader)

	limit := big.NewInt(1)
	limit.Lsh(limit, 256)

	m, _ := rand.Int(grand.Reader, limit)
	c, _ := Encrypt(m, pk, grand.Reader)

	k, _ := rand.Int(grand.Reader, limit)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		c.HomomorphicScalarMul(c, k, sk.Public())
	}
}
func TestPaillierHomomorphicEncSpeed(t *testing.T) {
	_, pk, _ := GenerateKey(2048, grand.Reader)

	m, _ := rand.Int(grand.Reader, pk.N)

	times := 1000
	start := time.Now()
	for i := 0; i < times; i++ {
		Encrypt(m, pk, grand.Reader)
	}
	end := time.Now()
	elapsed := end.Sub(start)
	t.Logf("HomomorphicScalarMul: %.2f per second\n", float64(times)/elapsed.Seconds())
}

func TestPaillierHomomorphicScalarMulSpeed(t *testing.T) {
	_, pk, _ := GenerateKey(2048, grand.Reader)

	limit := big.NewInt(1)
	limit.Lsh(limit, 256)

	m, _ := rand.Int(grand.Reader, limit)
	c, _ := Encrypt(m, pk, grand.Reader)

	k, _ := rand.Int(grand.Reader, limit)

	times := 1000
	start := time.Now()
	for i := 0; i < times; i++ {
		c.HomomorphicScalarMul(c, k, pk)
	}
	end := time.Now()
	elapsed := end.Sub(start)
	t.Logf("HomomorphicScalarMul: %.2f per second\n", float64(times)/elapsed.Seconds())
}

func BenchmarkBlind1(b *testing.B) {
	_, pk, _ := GenerateKey(2048, grand.Reader)

	limit := big.NewInt(1)
	limit.Lsh(limit, 256)

	k, _ := rand.Int(grand.Reader, limit)
	nc, _ := pk.Encrypt(sm2.OrderN(), grand.Reader)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		(&Cipher{}).HomomorphicScalarMul(nc, k, pk)
	}
}

func BenchmarkBlind2(b *testing.B) {
	_, pk, _ := GenerateKey(2048, grand.Reader)

	limit := big.NewInt(1)
	limit.Lsh(limit, 256)

	k, _ := rand.Int(grand.Reader, limit)
	kn := new(big.Int).Mul(k, sm2.OrderN())
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		kn.Mul(k, sm2.OrderN())
		pk.Encrypt(kn, grand.Reader)
	}
}

// func BenchmarkGMPIntExp(b *testing.B) {
// 	sk, _, _ := GenerateKey(2048, grand.Reader)
// 	n := gmp.NewInt(1).SetBytes(sk.N.Bytes())
// 	n2 := gmp.NewInt(1).Mul(n, n)

// 	r := gmp.NewInt(1).SetBytes(grand.GetRandom(128))
// 	b.ResetTimer()
// 	for i := 0; i < b.N; i++ {
// 		r.Exp(r, n, n2)
// 	}
// }

func BenchmarkIntExp(b *testing.B) {
	sk, _, _ := GenerateKey(2048, grand.Reader)
	n := sk.N
	n2 := sk.N2

	r := new(big.Int).SetBytes(grand.GetRandom(128))
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		r.Exp(r, n, n2)
	}
}

func TestGenerateKeyFromPassword(t *testing.T) {
	b, _ := hex.DecodeString("12345678abcdef")
	sk, _ := GenerateKeyFromPassword(2048, []byte("1234567812345678"), b, 1024)
	assert.Equal(t, sk.N.Text(16), "86257d4f684edb672420a6e0d6a6e66b00a316080c0b025d96e9d22c57d6fc7166bcddc07dadc2f0dca7440c6f43a90fab494ea810304af2abc18d54bcdeb8803c8c245270d70b897f3bf2e43c3a5e9d777fb256f84cf5b4705cba19283952dd10aa2b15793a401195ee2a418b999bbf72e35bd94f48267597ab7cfb699b3bbab78a6201a9ee18a22d25e841a66c54d5ebf45c10491ea43cca5e9e6f487c7aedc24d15cda20e6f80a755477b6e2b1d1a9e669ba0773bb31e2fb76e550fdef1be4d02ec820aba5ebfbfd0fbb37bd1c748af1f4a0c17f5cfb41e4cb44a05587bcff2937415eb2e13398da514224def6c03f61a07319512362912caa08bded076dd")
	assert.Equal(t, sk.Size(), uint(2048))
}