package cryptography

import (
	"bytes"
	"testing"
)

func TestBoxEncryptDecrypt(t *testing.T) {
	bk, err := GenerateBoxKey()
	if err != nil {
		t.Fatalf("GenerateBoxKey: %v", err)
	}

	plain := []byte("necropolis secure message")
	cipher, err := EncryptMessage(plain, bk.PublicKey)
	if err != nil {
		t.Fatalf("EncryptMessage: %v", err)
	}

	got, err := DecryptMessage(cipher, bk)
	if err != nil {
		t.Fatalf("DecryptMessage: %v", err)
	}
	if !bytes.Equal(got, plain) {
		t.Fatalf("roundtrip mismatch: got %q, want %q", got, plain)
	}

	if len(cipher) > 0 {
		cipher[0] ^= 0x01
	}
	_, err = DecryptMessage(cipher, bk)
	if err == nil {
		t.Fatal("expected error on tampered ciphertext, got nil")
	}
}

func TestSignVerify(t *testing.T) {
	op, err := GenerateOperatorKey()
	if err != nil {
		t.Fatalf("GenerateOperatorKey: %v", err)
	}

	data := []byte("necropolis command")
	sig, err := op.PrivateKey.Sign(data)
	if err != nil {
		t.Fatalf("Sign: %v", err)
	}

	ok, err := Verify(op.PublicKey, data, sig)
	if err != nil {
		t.Fatalf("Verify: %v", err)
	}
	if !ok {
		t.Fatal("signature verification failed")
	}

	data[0] ^= 0x01
	ok, err = Verify(op.PublicKey, data, sig)
	if err != nil {
		t.Fatalf("Verify tampered: %v", err)
	}
	if ok {
		t.Fatal("expected verification failure on tampered data")
	}
}

func TestBoxKeyPersistenceRoundtrip(t *testing.T) {
	orig, err := GenerateBoxKey()
	if err != nil {
		t.Fatalf("GenerateBoxKey: %v", err)
	}

	privBytes := orig.Marshal()
	if len(privBytes) != 32 {
		t.Fatalf("Marshal returned %d bytes, want 32", len(privBytes))
	}

	pubBytes := orig.BoxPubKeyBytes()
	if len(pubBytes) != 32 {
		t.Fatalf("BoxPubKeyBytes returned %d bytes, want 32", len(pubBytes))
	}

	reloaded := LoadBoxKey(privBytes)
	if reloaded == nil {
		t.Fatal("LoadBoxKey returned nil")
	}

	if !bytes.Equal(reloaded.PublicKey[:], pubBytes) {
		t.Fatal("reloaded public key does not match original public key")
	}

	plain := []byte("necropolis beacon payload")
	cipher, err := EncryptMessage(plain, reloaded.PublicKey)
	if err != nil {
		t.Fatalf("EncryptMessage with reloaded pub: %v", err)
	}

	got, err := DecryptMessage(cipher, reloaded)
	if err != nil {
		t.Fatalf("DecryptMessage with reloaded keys: %v", err)
	}
	if !bytes.Equal(got, plain) {
		t.Fatalf("reloaded roundtrip mismatch: got %q, want %q", got, plain)
	}

	got2, err := DecryptMessage(cipher, orig)
	if err != nil {
		t.Fatalf("DecryptMessage with original keys: %v", err)
	}
	if !bytes.Equal(got2, plain) {
		t.Fatalf("original roundtrip mismatch: got %q, want %q", got2, plain)
	}
}

func TestKeyPersistence(t *testing.T) {
	orig, err := GenerateOperatorKey()
	if err != nil {
		t.Fatalf("GenerateOperatorKey: %v", err)
	}

	marshaled, err := MarshalPrivateKey(orig.PrivateKey)
	if err != nil {
		t.Fatalf("MarshalPrivateKey: %v", err)
	}

	loaded, err := LoadPrivateKey(marshaled)
	if err != nil {
		t.Fatalf("LoadPrivateKey: %v", err)
	}

	if !orig.PublicKey.Equals(loaded.GetPublic()) {
		t.Fatal("public key mismatch after marshal roundtrip")
	}
}