package cryptoutil

import (
	"bytes"
	"encoding/hex"
	"io"
	"testing"
)

func mustHex(t testing.TB, s string) []byte {
	t.Helper()
	bytes, err := hex.DecodeString(s)
	if err != nil {
		t.Fatalf("Bad hex %q: %s", s, err)
	}
	return bytes
}

func TestPassword(t *testing.T) {
	const password = "eclair scroll gliding sled shining removed ascend android cheer confetti"
	hash, err := Hash(password)
	if err != nil {
		t.Fatalf("Hash(%q) failed: %s", password, err)
	}
	_, err = hash.CheckPassword(password)
	if err != nil {
		t.Errorf("%x.CheckPassword(%q) failed: %s", hash, password, err)
	}
}

func TestSignature(t *testing.T) {
	key := HMACKey(mustHex(t, "669e06ec457778b9a8133edb0a87ea82c6b141ffbbc63c038da96258175eb35c"))
	msg := []byte("test message")
	signedMsg := key.Sign(msg, "info")
	got, ok := key.Verify(signedMsg, "info")
	if !ok {
		t.Fatalf("Verify(%x) rejected the message", signedMsg)
	}
	if !bytes.Equal(got, msg) {
		t.Errorf("Verify(%x) = %x, want %x", signedMsg, got, msg)
	}
}

func TestEncrypt(t *testing.T) {
	key := EncryptionKey(mustHex(t, "b6de26860e0a39aa134732e58055c06ba028675453f73a6912bc932e58d5743d0e423217f06487d3e96a59980301fcc97dfcc4c6b19765f2947de3c33ab7ef9e"))
	for _, tc := range []struct {
		desc string
		msg  []byte
	}{{
		desc: "empty",
		msg:  nil,
	}, {
		desc: "short",
		msg:  []byte("test message"),
	}, {
		desc: "long",
		msg: func() []byte {
			msg := make([]byte, 1024*1024)
			for i := range msg {
				msg[i] = byte(i)
			}
			return msg
		}(),
	}} {
		t.Run(tc.desc, func(t *testing.T) {
			encryptedMsg := new(bytes.Buffer)
			w := key.NewWriter(encryptedMsg, []byte("additional data"))
			if _, err := w.Write(tc.msg); err != nil {
				t.Fatalf("EncryptingWriter.Write(%q) failed: %s", tc.msg, err)
			}
			if err := w.Close(); err != nil {
				t.Fatalf("EncryptingWriter.Close() failed: %s", err)
			}
			got, err := io.ReadAll(key.NewReader(bytes.NewReader(encryptedMsg.Bytes()), []byte("additional data")))
			if err != nil {
				t.Fatalf("DecryptingReader.Read(%x) failed: %s", encryptedMsg, err)
			}
			if !bytes.Equal(got, tc.msg) {
				t.Errorf("DecryptingReader.Read(%x) = %x, want %x", encryptedMsg, got, tc.msg)
			}
		})
	}
}

func TestDeriveKey(t *testing.T) {
	const password = "automaker moisture botch tubular kelp rinse rule giggly donated stock"
	hash, err := Hash(password)
	if err != nil {
		t.Fatalf("Hash(%q) failed: %s", password, err)
	}
	rawKey, err := hash.CheckPassword(password)
	if err != nil {
		t.Fatalf("%x.CheckPassword(%q) failed: %s", hash, password, err)
	}
	key, err := rawKey.EncryptionKey()
	if err != nil {
		t.Fatalf("EncryptionKey() failed: %s", err)
	}
	msg := []byte("test message")
	encryptedMsg := new(bytes.Buffer)
	w := key.NewWriter(encryptedMsg, []byte("additional data"))
	if _, err := w.Write(msg); err != nil {
		t.Fatalf("EncryptingWriter.Write(%q) failed: %s", msg, err)
	}
	if err := w.Close(); err != nil {
		t.Fatalf("EncryptingWriter.Close() failed: %s", err)
	}
	got, err := io.ReadAll(key.NewReader(encryptedMsg, []byte("additional data")))
	if err != nil {
		t.Fatalf("DecryptingReader.Read(%x) failed: %s", encryptedMsg, err)
	}
	if !bytes.Equal(got, msg) {
		t.Errorf("DecryptingReader.Read(%x) = %x, want %x", encryptedMsg, got, msg)
	}
}

func BenchmarkEncrypt(b *testing.B) {
	key := EncryptionKey(mustHex(b, "b6de26860e0a39aa134732e58055c06ba028675453f73a6912bc932e58d5743d0e423217f06487d3e96a59980301fcc97dfcc4c6b19765f2947de3c33ab7ef9e"))
	msg := make([]byte, 1024*1024)
	for i := range msg {
		msg[i] = byte(i)
	}
	encryptedMsg := bytes.NewBuffer(make([]byte, 0, 35+len(msg)+16*(len(msg)+segmentSize-1)/segmentSize /* ?? */))
	additionalData := []byte("additional data")
	for b.Loop() {
		w := key.NewWriter(encryptedMsg, additionalData)
		if _, err := w.Write(msg); err != nil {
			b.Fatal(err)
		}
		if err := w.Close(); err != nil {
			b.Fatal(err)
		}
		encryptedMsg.Reset()
	}
}

func BenchmarkDecrypt(b *testing.B) {
	key := EncryptionKey(mustHex(b, "b6de26860e0a39aa134732e58055c06ba028675453f73a6912bc932e58d5743d0e423217f06487d3e96a59980301fcc97dfcc4c6b19765f2947de3c33ab7ef9e"))
	msg := make([]byte, 1024*1024)
	for i := range msg {
		msg[i] = byte(i)
	}
	encryptedMsg := new(bytes.Buffer)
	additionalData := []byte("additional data")
	w := key.NewWriter(encryptedMsg, additionalData)
	if _, err := w.Write(msg); err != nil {
		b.Fatal(err)
	}
	if err := w.Close(); err != nil {
		b.Fatal(err)
	}
	decryptedMsg := make([]byte, len(msg))
	for b.Loop() {
		r := key.NewReader(bytes.NewReader(encryptedMsg.Bytes()), additionalData)
		if _, err := io.ReadFull(r, decryptedMsg); err != nil {
			b.Fatal(err)
		}
	}
}