package fpe

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

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

func TestNum(t *testing.T) {
	N := 1000
	A := make([]Numeral, N)
	B := make([]Numeral, N)

	for r := 2; r < 100; r++ {
		radix := NewRadix(r)
		for i := 1; i <= N; i++ {
			for j := 0; j < i; j++ {
				A[j] = Numeral(rand.Int31() % int32(r))
			}
			numSwitch = false
			a0 := radix.Num(A[:i])
			numSwitch = true
			a1 := radix.Num(A[:i])
			assert.Equal(t, a0.Cmp(a1), 0)

			radix.Str(B[:i], a0)
			assert.Equal(t, A[:i], B[:i])
		}
	}
}

func TestStr(t *testing.T) {
	f := Alpha
	// 24948832942366750129003083595837476696096111951836798179448429117431251233398075403977925597490
	A := []Numeral{
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
	}

	for i := 20; i <= len(A); i++ {
		a := f.Num(A[:i])
		B := make([]Numeral, i)
		f.Str(B, a)
		for j := range B {
			if B[j] != A[j] {
				t.Fatal()
			}
		}
	}
}

// BenchmarkAlphabet-10    	 1158810	      1020 ns/op	      88 B/op	       3 allocs/op
// BenchmarkAlphabet-10    	 3776529	       310.4 ns/op	      96 B/op	       4 allocs/op
func BenchmarkNum(b *testing.B) {
	f := Alpha
	A := []Numeral{
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		f.Num(A)
	}
}

// BenchmarkStr-10    	  556981	      2130 ns/op	     104 B/op	       2 allocs/op
// BenchmarkStr-10    	 3134851	       382.2 ns/op	     104 B/op	       2 allocs/op
func BenchmarkStr(b *testing.B) {
	f := Alpha
	A := []Numeral{
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
		1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
	}

	a := f.Num(A)
	B := make([]Numeral, len(A))
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		f.Str(B, a)
	}
}

func TestFF1(t *testing.T) {
	// fmt.Printf("%x\n", ^uint64(15))
	// 密钥
	key := grand.GetRandom(16)
	// 使用SM4加密算法
	b, _ := sm4.NewCipher(key)
	// tweak
	T := grand.GetRandom(16)

	// interface FPE
	f := NewFF1(b, NewAlphabet("abcdefghijklmnopqrstuvwxyz "))

	plainString0 := "hello world"
	plain0, err := f.Encode(plainString0)
	assert.Nil(t, err)

	cipher, err := f.Encrypt(T, plain0)
	assert.Nil(t, err)

	cipherString, err := f.Decode(cipher)
	assert.Nil(t, err)

	fmt.Println("Cipher string:", cipherString)
	plain1, err := f.Decrypt(T, cipher)
	assert.Nil(t, err)

	assert.Equal(t, plain0, plain1)
	plainString1, err := f.Decode(plain1)
	assert.Nil(t, err)

	fmt.Println("Decrypt string:", plainString1)
	assert.Equal(t, plainString0, plainString1)
}

// BenchmarkFF1-10    	  129781	      8683 ns/op	    4256 B/op	     154 allocs/op
// BenchmarkFF1-10    	  325104	      3694 ns/op	    2808 B/op	      99 allocs/op
func BenchmarkFF1(b *testing.B) {
	key := grand.GetRandom(16)
	block, _ := sm4.NewCipher(key)
	T := grand.GetRandom(16)
	f := NewFF1(block, Printable)

	plainString0 := "hello world"
	plain0, _ := f.Encode(plainString0)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, _ = f.Encrypt(T, plain0)
	}
}

// // 密文长度    加密速率(MBps)
// //   32         2.55
// //   64         3.99
// //  128         3.58
// //  256         3.82
// //  512         3.65
// // 1024         3.15
// // 2048         2.38
// // 4096         1.58
// // 8192         0.94
func TestSpeed(t *testing.T) {
	key := grand.GetRandom(16)
	block, _ := sm4.NewCipher(key)
	T := grand.GetRandom(16)
	f := NewFF1(block, Printable)

	fmt.Println("密文长度    加密速率(MBps)")
	for plainLen, cnt := 32, 10000; plainLen < 10000; plainLen, cnt = plainLen*2, cnt/2 {
		b := make([]byte, plainLen)
		for i := range b {
			b[i] = byte(rand.Uint32()%26 + 'a')
		}
		plain := string(b)

		start := time.Now()
		for i := 0; i < cnt; i++ {
			n, _ := f.Encode(plain)
			_, _ = f.Encrypt(T, n)

		}
		end := time.Now()
		elapsed := end.Sub(start)
		fmt.Printf("%4d         %.2f\n", plainLen, float64(len(plain)*cnt)/(1024*1024*elapsed.Seconds()))
	}
}