package peer

import (
	"crypto/rand"
	"log"
	"net/netip"
	"sync/atomic"

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

type cryptoKeys struct {
	PubKey      []byte
	PrivKey     []byte
	PubSignKey  []byte
	PrivSignKey []byte
}

func generateKeys() cryptoKeys {
	pubKey, privKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		log.Fatalf("Failed to generate encryption keys: %v", err)
	}

	pubSignKey, privSignKey, err := sign.GenerateKey(rand.Reader)
	if err != nil {
		log.Fatalf("Failed to generate signing keys: %v", err)
	}

	return cryptoKeys{pubKey[:], privKey[:], pubSignKey[:], privSignKey[:]}
}

// ----------------------------------------------------------------------------

// Route must have a ControlCipher.
func encryptControlPacket(
	localIP byte,
	route *peerRoute,
	pkt marshaller,
	tmp []byte,
	out []byte,
) []byte {
	h := header{
		StreamID: controlStreamID,
		Counter:  atomic.AddUint64(route.Counter, 1),
		SourceIP: localIP,
		DestIP:   route.IP,
	}
	tmp = pkt.Marshal(tmp)
	return route.ControlCipher.Encrypt(h, tmp, out)
}

// Returns a controlMsg[PacketType]. Route must have ControlCipher.
//
// This function also drops packets with duplicate sequence numbers.
func decryptControlPacket(
	route *peerRoute,
	fromAddr netip.AddrPort,
	h header,
	encrypted []byte,
	tmp []byte,
) (any, error) {
	out, ok := route.ControlCipher.Decrypt(encrypted, tmp)
	if !ok {
		return nil, errDecryptionFailed
	}

	if route.DupCheck.IsDup(h.Counter) {
		return nil, errDuplicateSeqNum
	}

	msg, err := parseControlMsg(h.SourceIP, fromAddr, out)
	if err != nil {
		return nil, err
	}

	return msg, nil
}

// ----------------------------------------------------------------------------

func encryptDataPacket(
	localIP byte,
	destIP byte,
	route *peerRoute,
	data []byte,
	out []byte,
) []byte {
	h := header{
		StreamID: dataStreamID,
		Counter:  atomic.AddUint64(route.Counter, 1),
		SourceIP: localIP,
		DestIP:   destIP,
	}
	return route.DataCipher.Encrypt(h, data, out)
}

func decryptDataPacket(
	route *peerRoute,
	h header,
	encrypted []byte,
	out []byte,
) ([]byte, error) {
	dec, ok := route.DataCipher.Decrypt(encrypted, out)
	if !ok {
		return nil, errDecryptionFailed
	}

	if route.DupCheck.IsDup(h.Counter) {
		return nil, errDuplicateSeqNum
	}

	return dec, nil
}