package node

import (
	"log"
	"net"
	"net/netip"
	"sync"
	"sync/atomic"
	"time"
)

type controlPacketSender interface {
	SendControlPacket(pkt marshaller, route peerRoute)
}

type dataPacketSender interface {
	SendDataPacket(pkt []byte, route peerRoute)
	SendEncryptedDataPacket(pkt []byte, addr netip.AddrPort)
}

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

type packetSender struct {
	conn *net.UDPConn

	// For sending control packets.
	cLock sync.Mutex
	cBuf1 []byte
	cBuf2 []byte

	// For sending data packets.
	dBuf1 []byte
	dBuf2 []byte

	counters [256]uint64

	// Lock around for sending on UDP Conn.
	wLock sync.Mutex
}

func newPacketSender(conn *net.UDPConn) *packetSender {
	ps := &packetSender{
		conn:  conn,
		cBuf1: make([]byte, bufferSize),
		cBuf2: make([]byte, bufferSize),
		dBuf1: make([]byte, bufferSize),
		dBuf2: make([]byte, bufferSize),
	}
	for i := range ps.counters {
		ps.counters[i] = uint64(time.Now().Unix()<<30 + 1)
	}
	return ps
}

// Safe for concurrent use.
func (sender *packetSender) SendControlPacket(pkt marshaller, route peerRoute) {
	sender.cLock.Lock()
	defer sender.cLock.Unlock()

	buf := pkt.Marshal(sender.cBuf1)
	h := header{
		StreamID: controlStreamID,
		Counter:  atomic.AddUint64(&sender.counters[route.IP], 1),
		SourceIP: localIP,
		DestIP:   route.IP,
	}
	buf = route.ControlCipher.Encrypt(h, buf, sender.cBuf2)

	if route.Direct {
		sender.writeTo(buf, route.RemoteAddr)
		return
	}

	sender.relayPacket(route.IP, buf, sender.cBuf1)
}

// Not safe for concurrent use.
func (sender *packetSender) SendDataPacket(pkt []byte, route peerRoute) {
	h := header{
		StreamID: dataStreamID,
		Counter:  atomic.AddUint64(&sender.counters[route.IP], 1),
		SourceIP: localIP,
		DestIP:   route.IP,
	}

	enc := route.DataCipher.Encrypt(h, pkt, sender.dBuf1)

	if route.Direct {
		sender.writeTo(enc, route.RemoteAddr)
		return
	}

	sender.relayPacket(route.IP, enc, sender.dBuf2)
}

func (sender *packetSender) SendEncryptedDataPacket(pkt []byte, addr netip.AddrPort) {
	sender.writeTo(pkt, addr)
}

func (sender *packetSender) relayPacket(destIP byte, data, buf []byte) {
	ip := relayIP.Load()
	if ip == nil {
		return
	}
	relayRoute := routingTable[*ip].Load()
	if relayRoute == nil || !relayRoute.Up || !relayRoute.Relay {
		return
	}

	h := header{
		StreamID: dataStreamID,
		Counter:  atomic.AddUint64(&sender.counters[relayRoute.IP], 1),
		SourceIP: localIP,
		DestIP:   destIP,
	}

	enc := relayRoute.DataCipher.Encrypt(h, data, buf)
	sender.writeTo(enc, relayRoute.RemoteAddr)
}

func (sender *packetSender) writeTo(packet []byte, addr netip.AddrPort) {
	sender.wLock.Lock()
	if _, err := sender.conn.WriteToUDPAddrPort(packet, addr); err != nil {
		log.Printf("Failed to write to UDP port: %v", err)
	}
	sender.wLock.Unlock()
}