package node

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

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

type peerRoute struct {
	IP            byte
	Up            bool // True if data can be sent on the route.
	Relay         bool // True if the peer is a relay.
	Direct        bool // True if this is a direct connection.
	PubSignKey    []byte
	ControlCipher *controlCipher
	DataCipher    *dataCipher
	RemoteAddr    netip.AddrPort // Remote address if directly connected.
}

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

type udpAddrPortWriter interface {
	WriteToUDPAddrPort([]byte, netip.AddrPort) (int, error)
}

type marshaller interface {
	Marshal([]byte) []byte
}

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

type connWriter struct {
	localIP byte
	conn    udpAddrPortWriter

	// For sending control packets.
	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 newConnWriter(conn udpAddrPortWriter, localIP byte) *connWriter {
	w := &connWriter{
		localIP: localIP,
		conn:    conn,
		cBuf1:   make([]byte, bufferSize),
		cBuf2:   make([]byte, bufferSize),
		dBuf1:   make([]byte, bufferSize),
		dBuf2:   make([]byte, bufferSize),
	}
	for i := range w.counters {
		w.counters[i] = uint64(time.Now().Unix()<<30 + 1)
	}
	return w
}

// Not safe for concurrent use. Should only be called by supervisor.
func (w *connWriter) SendControlPacket(pkt marshaller, route *peerRoute) {
	buf := pkt.Marshal(w.cBuf1)
	h := header{
		StreamID: controlStreamID,
		Counter:  atomic.AddUint64(&w.counters[route.IP], 1),
		SourceIP: w.localIP,
		DestIP:   route.IP,
	}
	buf = route.ControlCipher.Encrypt(h, buf, w.cBuf2)
	w.writeTo(buf, route.RemoteAddr)
}

func (w *connWriter) RelayControlPacket(pkt marshaller, route, relay *peerRoute) {
	buf := pkt.Marshal(w.cBuf1)
	h := header{
		StreamID: controlStreamID,
		Counter:  atomic.AddUint64(&w.counters[route.IP], 1),
		SourceIP: w.localIP,
		DestIP:   route.IP,
	}
	buf = route.ControlCipher.Encrypt(h, buf, w.cBuf2)
	w.relayPacket(buf, w.cBuf1, route, relay)
}

// Not safe for concurrent use. Should only be called by ifReader.
func (w *connWriter) SendDataPacket(pkt []byte, route, relay *peerRoute) {
	h := header{
		StreamID: dataStreamID,
		Counter:  atomic.AddUint64(&w.counters[route.IP], 1),
		SourceIP: w.localIP,
		DestIP:   route.IP,
	}

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

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

	w.relayPacket(enc, w.dBuf2, route, relay)
}

// TODO: RelayDataPacket

// Safe for concurrent use. Should only be called by connReader.
//
// This function will send pkt to the peer directly. This is used when a peer
// is acting as a relay and is forwarding already encrypted data for another
// peer.
func (w *connWriter) SendEncryptedDataPacket(pkt []byte, route *peerRoute) {
	w.writeTo(pkt, route.RemoteAddr)
}

func (w *connWriter) relayPacket(data, buf []byte, route, relay *peerRoute) {
	if relay == nil || !relay.Up {
		return
	}

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

	enc := relay.DataCipher.Encrypt(h, data, buf)
	w.writeTo(enc, relay.RemoteAddr)
}

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