// Package relay implements NaCl-box encrypted envelope routing over gossipsub.
// Messages are sealed for a specific recipient's X25519 public key; relay nodes
// propagate them without being able to read the contents.
//
// Economic model: the sender pre-authorises a delivery fee by signing
// FeeAuthBytes(envelopeID, feeUT) with their Ed25519 identity key. When the
// relay delivers the envelope, it submits a RELAY_PROOF transaction on-chain
// that pulls feeUT from the sender's balance and credits the relay.
package relay

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"

	"golang.org/x/crypto/nacl/box"

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

// KeyPair holds an X25519 keypair used exclusively for relay envelope encryption.
// This is separate from the node's Ed25519 identity keypair.
type KeyPair struct {
	Pub  [32]byte
	Priv [32]byte
}

// GenerateKeyPair creates a fresh X25519 keypair.
func GenerateKeyPair() (*KeyPair, error) {
	pub, priv, err := box.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("generate relay keypair: %w", err)
	}
	return &KeyPair{Pub: *pub, Priv: *priv}, nil
}

// PubHex returns the hex-encoded X25519 public key.
func (kp *KeyPair) PubHex() string {
	return hex.EncodeToString(kp.Pub[:])
}

// Envelope is a sealed message routed via relay nodes.
// Only the holder of the matching X25519 private key can decrypt it.
type Envelope struct {
	// ID is the hex-encoded first 16 bytes of SHA-256(nonce || ciphertext).
	ID           string `json:"id"`
	RecipientPub string `json:"recipient_pub"` // hex X25519 public key
	SenderPub    string `json:"sender_pub"`    // hex X25519 public key (for decryption)

	// Fee authorization: sender pre-signs permission for relay to pull FeeUT.
	SenderEd25519PubKey string `json:"sender_ed25519_pub"` // sender's blockchain identity key (hex)
	FeeUT               uint64 `json:"fee_ut"`             // µT the relay may claim on delivery
	FeeSig              []byte `json:"fee_sig"`            // Ed25519 sig over FeeAuthBytes(ID, FeeUT)

	Nonce      []byte `json:"nonce"`      // 24 bytes
	Ciphertext []byte `json:"ciphertext"` // NaCl box ciphertext
	SentAt     int64  `json:"sent_at"`    // unix timestamp (informational)
}

// Seal encrypts msg for recipientPub and attaches a fee authorization.
// senderID is the sender's Ed25519 identity used to sign the fee authorisation.
// feeUT is the delivery fee offered to the relay — 0 means free delivery.
func Seal(
	sender *KeyPair,
	senderID *identity.Identity,
	recipientPub [32]byte,
	msg []byte,
	feeUT uint64,
	sentAt int64,
) (*Envelope, error) {
	var nonce [24]byte
	if _, err := rand.Read(nonce[:]); err != nil {
		return nil, fmt.Errorf("generate nonce: %w", err)
	}
	ct := box.Seal(nil, msg, &nonce, &recipientPub, &sender.Priv)
	envID := envelopeID(nonce[:], ct)

	var feeSig []byte
	if feeUT > 0 && senderID != nil {
		authBytes := blockchain.FeeAuthBytes(envID, feeUT)
		feeSig = senderID.Sign(authBytes)
	}

	return &Envelope{
		ID:                  envID,
		RecipientPub:        hex.EncodeToString(recipientPub[:]),
		SenderPub:           sender.PubHex(),
		SenderEd25519PubKey: func() string {
			if senderID != nil {
				return senderID.PubKeyHex()
			}
			return ""
		}(),
		FeeUT:      feeUT,
		FeeSig:     feeSig,
		Nonce:      nonce[:],
		Ciphertext: ct,
		SentAt:     sentAt,
	}, nil
}

// Open decrypts an envelope using the recipient's private key.
func Open(recipient *KeyPair, env *Envelope) ([]byte, error) {
	senderBytes, err := hex.DecodeString(env.SenderPub)
	if err != nil || len(senderBytes) != 32 {
		return nil, fmt.Errorf("invalid sender pub key")
	}
	if len(env.Nonce) != 24 {
		return nil, fmt.Errorf("invalid nonce: expected 24 bytes, got %d", len(env.Nonce))
	}
	var senderPub [32]byte
	var nonce [24]byte
	copy(senderPub[:], senderBytes)
	copy(nonce[:], env.Nonce)

	msg, ok := box.Open(nil, env.Ciphertext, &nonce, &senderPub, &recipient.Priv)
	if !ok {
		return nil, fmt.Errorf("decryption failed: not addressed to this key or data is corrupt")
	}
	return msg, nil
}

// Hash returns the SHA-256 of (nonce || ciphertext), used in RELAY_PROOF payloads.
func Hash(env *Envelope) []byte {
	h := sha256.Sum256(append(env.Nonce, env.Ciphertext...))
	return h[:]
}

// IsAddressedTo reports whether the envelope is addressed to the given keypair.
func (e *Envelope) IsAddressedTo(kp *KeyPair) bool {
	return e.RecipientPub == kp.PubHex()
}

func envelopeID(nonce, ct []byte) string {
	h := sha256.Sum256(append(nonce, ct...))
	return hex.EncodeToString(h[:16])
}