package romulus_go

import (
	"bufio"
	"encoding/hex"
	"fmt"
	"os"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
)

type testData struct {
	Name  string
	Key   []byte
	Nonce []byte
	// Plaintext
	PT []byte
	// Addional data
	AD []byte
	// Cipher text
	CT []byte
}

var tests = loadTests()

func loadTests() []testData {
	var tData []testData
	dat, _ := os.ReadFile("LWC_AEAD_KAT_128_128.txt")
	scanner := bufio.NewScanner(strings.NewReader((string)(dat)))

	var currentTest testData
	for scanner.Scan() {
		noSpace := strings.ReplaceAll(scanner.Text(), " ", "")
		split := strings.Split(noSpace, "=")
		if len(split) > 0 {
			switch split[0] {
			case "Count":
				currentTest.Name = fmt.Sprintf("Test Number %v", split[1])
			case "Key":
				currentTest.Key, _ = hex.DecodeString(split[1])
			case "Nonce":
				currentTest.Nonce, _ = hex.DecodeString(split[1])
			case "PT":
				currentTest.PT, _ = hex.DecodeString(split[1])
			case "AD":
				currentTest.AD, _ = hex.DecodeString(split[1])
			case "CT":
				currentTest.CT, _ = hex.DecodeString(split[1])
				tData = append(tData, currentTest)
				currentTest = testData{}
			}
		}
	}

	return tData
}

func TestEnc(t *testing.T) {

	for _, r := range tests {
		t.Run(r.Name, func(t *testing.T) {
			c := make([]byte, len(r.PT)+16)
			var clen uint64
			romulus_m_encrypt(c[:], &clen, r.PT, (uint64)(len(r.PT)), r.AD, (uint64)(len(r.AD)), nil, r.Nonce, r.Key)
			assert.Equal(t, (uint64)(len(r.CT)), clen)
			assert.Equal(t, c, r.CT)
		})
	}
}

func TestDec(t *testing.T) {

	for _, r := range tests {
		t.Run(r.Name, func(t *testing.T) {
			d := make([]byte, len(r.CT))
			var dlen uint64
			romulus_m_decrypt(d, &dlen, nil, r.CT, (uint64)(len(r.CT)), r.AD, (uint64)(len(r.AD)), r.Nonce, r.Key)
			assert.Equal(t, (uint64)(len(r.PT)), dlen)
			assert.Equal(t, d[:dlen], r.PT)
		})
	}
}