package peer

import (
	"bytes"
	"crypto/rand"
	mrand "math/rand/v2"
	"reflect"
	"testing"
)

func TestDataCipher(t *testing.T) {
	maxSizePlaintext := make([]byte, bufferSize-dataHeaderSize-dataCipherOverhead)
	rand.Read(maxSizePlaintext)

	testCases := [][]byte{
		make([]byte, 0),
		{1},
		{255},
		{1, 2, 3, 4, 5},
		[]byte("Hello world"),
		maxSizePlaintext,
	}

	for _, plaintext := range testCases {
		h1 := Header{
			StreamID: dataStreamID,
			Counter:  235153,
			SourceIP: 4,
			DestIP:   88,
		}

		encrypted := make([]byte, bufferSize)

		dc1 := newDataCipher()
		encrypted = dc1.Encrypt(h1, plaintext, encrypted)
		h2 := Header{}
		h2.Parse(encrypted)

		dc2 := newDataCipherFromKey(dc1.Key())

		decrypted, ok := dc2.Decrypt(encrypted, make([]byte, bufferSize-dataHeaderSize))
		if !ok {
			t.Fatal(ok)
		}

		if !bytes.Equal(plaintext, decrypted) {
			t.Fatal("not equal")
		}

		if !reflect.DeepEqual(h1, h2) {
			t.Fatalf("%v != %v", h1, h2)
		}
	}
}

func TestDataCipher_ModifyCiphertext(t *testing.T) {
	maxSizePlaintext := make([]byte, bufferSize-dataHeaderSize-dataCipherOverhead)
	rand.Read(maxSizePlaintext)

	testCases := [][]byte{
		make([]byte, 0),
		{1},
		{255},
		{1, 2, 3, 4, 5},
		[]byte("Hello world"),
		maxSizePlaintext,
	}

	for _, plaintext := range testCases {
		h1 := Header{
			Counter:  235153,
			SourceIP: 4,
			DestIP:   88,
		}

		encrypted := make([]byte, bufferSize)

		dc1 := newDataCipher()
		encrypted = dc1.Encrypt(h1, plaintext, encrypted)
		encrypted[mrand.IntN(len(encrypted))]++

		dc2 := newDataCipherFromKey(dc1.Key())

		_, ok := dc2.Decrypt(encrypted, make([]byte, bufferSize-dataHeaderSize))
		if ok {
			t.Fatal(ok)
		}
	}
}

func TestDataCipher_ShortCiphertext(t *testing.T) {
	dc1 := newDataCipher()
	shortText := make([]byte, dataHeaderSize+dataCipherOverhead-1)
	rand.Read(shortText)
	_, ok := dc1.Decrypt(shortText, make([]byte, bufferSize))
	if ok {
		t.Fatal(ok)
	}
}

func BenchmarkDataCipher_Encrypt(b *testing.B) {
	h1 := Header{
		Counter:  235153,
		SourceIP: 4,
		DestIP:   88,
	}

	plaintext := make([]byte, bufferSize-dataHeaderSize-dataCipherOverhead)
	rand.Read(plaintext)

	encrypted := make([]byte, bufferSize)

	dc1 := newDataCipher()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		encrypted = dc1.Encrypt(h1, plaintext, encrypted)
	}
}

func BenchmarkDataCipher_Decrypt(b *testing.B) {
	h1 := Header{
		Counter:  235153,
		SourceIP: 4,
		DestIP:   88,
	}

	plaintext := make([]byte, bufferSize-dataHeaderSize-dataCipherOverhead)
	rand.Read(plaintext)

	encrypted := make([]byte, bufferSize)

	dc1 := newDataCipher()
	encrypted = dc1.Encrypt(h1, plaintext, encrypted)

	decrypted := make([]byte, bufferSize)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		decrypted, _ = dc1.Decrypt(encrypted, decrypted)
	}
}
