package relay_test

import (
	"bytes"
	"testing"
	"time"

	"go-blockchain/identity"
	"go-blockchain/relay"
)

func mustGenerateKeyPair(t *testing.T) *relay.KeyPair {
	t.Helper()
	kp, err := relay.GenerateKeyPair()
	if err != nil {
		t.Fatalf("GenerateKeyPair: %v", err)
	}
	return kp
}

func mustGenerateIdentity(t *testing.T) *identity.Identity {
	t.Helper()
	id, err := identity.Generate()
	if err != nil {
		t.Fatalf("identity.Generate: %v", err)
	}
	return id
}

// TestSealOpenRoundTrip seals a message and opens it with the correct recipient key.
func TestSealOpenRoundTrip(t *testing.T) {
	sender := mustGenerateKeyPair(t)
	recipient := mustGenerateKeyPair(t)
	plaintext := []byte("hello relay world")

	env, err := relay.Seal(sender, nil, recipient.Pub, plaintext, 0, time.Now().Unix())
	if err != nil {
		t.Fatalf("Seal: %v", err)
	}

	got, err := relay.Open(recipient, env)
	if err != nil {
		t.Fatalf("Open: %v", err)
	}
	if !bytes.Equal(got, plaintext) {
		t.Errorf("plaintext mismatch: got %q, want %q", got, plaintext)
	}
}

// TestOpenWrongKey attempts to open an envelope with a different keypair.
func TestOpenWrongKey(t *testing.T) {
	sender := mustGenerateKeyPair(t)
	recipient := mustGenerateKeyPair(t)
	wrong := mustGenerateKeyPair(t)

	env, err := relay.Seal(sender, nil, recipient.Pub, []byte("secret"), 0, time.Now().Unix())
	if err != nil {
		t.Fatalf("Seal: %v", err)
	}

	_, err = relay.Open(wrong, env)
	if err == nil {
		t.Fatal("Open with wrong key should return an error")
	}
}

// TestIsAddressedTo checks that IsAddressedTo returns true for the correct
// recipient and false for a different keypair.
func TestIsAddressedTo(t *testing.T) {
	sender := mustGenerateKeyPair(t)
	recipient := mustGenerateKeyPair(t)
	other := mustGenerateKeyPair(t)

	env, err := relay.Seal(sender, nil, recipient.Pub, []byte("msg"), 0, time.Now().Unix())
	if err != nil {
		t.Fatalf("Seal: %v", err)
	}

	if !env.IsAddressedTo(recipient) {
		t.Error("IsAddressedTo should return true for the correct recipient")
	}
	if env.IsAddressedTo(other) {
		t.Error("IsAddressedTo should return false for a different keypair")
	}
}

// TestSealWithFee seals a message with a positive fee and a real sender identity.
// FeeSig must be non-nil and SenderEd25519PubKey must be populated.
func TestSealWithFee(t *testing.T) {
	sender := mustGenerateKeyPair(t)
	senderID := mustGenerateIdentity(t)
	recipient := mustGenerateKeyPair(t)

	env, err := relay.Seal(sender, senderID, recipient.Pub, []byte("paid msg"), 1000, time.Now().Unix())
	if err != nil {
		t.Fatalf("Seal: %v", err)
	}

	if len(env.FeeSig) == 0 {
		t.Error("FeeSig should be non-nil when feeUT > 0 and senderID is provided")
	}
	if env.SenderEd25519PubKey == "" {
		t.Error("SenderEd25519PubKey should be set when senderID is provided")
	}
	if env.SenderEd25519PubKey != senderID.PubKeyHex() {
		t.Errorf("SenderEd25519PubKey mismatch: got %s, want %s",
			env.SenderEd25519PubKey, senderID.PubKeyHex())
	}
	if env.FeeUT != 1000 {
		t.Errorf("FeeUT: got %d, want 1000", env.FeeUT)
	}
}

// TestSealNilSenderID seals with nil senderID and feeUT=0.
// FeeSig should be nil and SenderEd25519PubKey should be empty.
func TestSealNilSenderID(t *testing.T) {
	sender := mustGenerateKeyPair(t)
	recipient := mustGenerateKeyPair(t)

	env, err := relay.Seal(sender, nil, recipient.Pub, []byte("free msg"), 0, time.Now().Unix())
	if err != nil {
		t.Fatalf("Seal: %v", err)
	}

	if env.FeeSig != nil {
		t.Error("FeeSig should be nil for zero-fee envelope with nil senderID")
	}
	if env.SenderEd25519PubKey != "" {
		t.Errorf("SenderEd25519PubKey should be empty, got %s", env.SenderEd25519PubKey)
	}
}

// TestEnvelopeIDUnique verifies that two seals of the same message produce
// different IDs because random nonces are used each time.
func TestEnvelopeIDUnique(t *testing.T) {
	sender := mustGenerateKeyPair(t)
	recipient := mustGenerateKeyPair(t)
	msg := []byte("same message")
	sentAt := time.Now().Unix()

	env1, err := relay.Seal(sender, nil, recipient.Pub, msg, 0, sentAt)
	if err != nil {
		t.Fatalf("Seal 1: %v", err)
	}
	env2, err := relay.Seal(sender, nil, recipient.Pub, msg, 0, sentAt)
	if err != nil {
		t.Fatalf("Seal 2: %v", err)
	}

	if env1.ID == env2.ID {
		t.Error("two seals of the same message should produce different IDs")
	}
}

// TestOpenTamperedCiphertext flips a byte in the ciphertext and expects Open to fail.
func TestOpenTamperedCiphertext(t *testing.T) {
	sender := mustGenerateKeyPair(t)
	recipient := mustGenerateKeyPair(t)

	env, err := relay.Seal(sender, nil, recipient.Pub, []byte("tamper me"), 0, time.Now().Unix())
	if err != nil {
		t.Fatalf("Seal: %v", err)
	}

	// Flip the first byte of the ciphertext.
	env.Ciphertext[0] ^= 0xFF

	_, err = relay.Open(recipient, env)
	if err == nil {
		t.Fatal("Open with tampered ciphertext should return an error")
	}
}

// TestOpenTamperedNonce flips a byte in the nonce and expects Open to fail.
func TestOpenTamperedNonce(t *testing.T) {
	sender := mustGenerateKeyPair(t)
	recipient := mustGenerateKeyPair(t)

	env, err := relay.Seal(sender, nil, recipient.Pub, []byte("tamper nonce"), 0, time.Now().Unix())
	if err != nil {
		t.Fatalf("Seal: %v", err)
	}

	// Flip the first byte of the nonce.
	env.Nonce[0] ^= 0xFF

	_, err = relay.Open(recipient, env)
	if err == nil {
		t.Fatal("Open with tampered nonce should return an error")
	}
}