package peer

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

type ConnWriter struct {
	wLock sync.Mutex // Lock around for sending on UDP Conn.

	// Output.
	writeToUDPAddrPort func([]byte, netip.AddrPort) (int, error)

	// Shared state.
	rt *atomic.Pointer[RoutingTable]

	// For sending control packets.
	cBuf1 []byte
	cBuf2 []byte

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

func NewConnWriter(
	writeToUDPAddrPort func([]byte, netip.AddrPort) (int, error),
	rt *atomic.Pointer[RoutingTable],
) *ConnWriter {
	return &ConnWriter{
		writeToUDPAddrPort: writeToUDPAddrPort,
		rt:                 rt,
		cBuf1:              newBuf(),
		cBuf2:              newBuf(),
		dBuf1:              newBuf(),
		dBuf2:              newBuf(),
	}
}

// Called by ConnReader to forward already encrypted bytes to another peer.
func (w *ConnWriter) Forward(ip byte, pkt []byte) {
	peer := w.rt.Load().Peers[ip]
	if !(peer.Up && peer.Direct) {
		w.logf("Failed to forward to %d.", ip)
		return
	}
	w.writeTo(pkt, peer.DirectAddr)
}

// Called by IFReader to send data. Encryption will be applied, and packet will
// be relayed if appropriate.
func (w *ConnWriter) WriteData(ip byte, pkt []byte) {
	rt := w.rt.Load()
	peer := rt.Peers[ip]
	if !peer.Up {
		w.logf("Failed to send data to %d.", ip)
		return
	}

	enc := peer.EncryptDataPacket(ip, pkt, w.dBuf1)

	if peer.Direct {
		w.writeTo(enc, peer.DirectAddr)
		return
	}

	relay, ok := rt.GetRelay()
	if !ok {
		w.logf("Failed to send data to %d. No relay.", ip)
		return
	}

	enc = relay.EncryptDataPacket(ip, enc, w.dBuf2)
	w.writeTo(enc, relay.DirectAddr)
}

// Called by Supervisor to send control packets.
func (w *ConnWriter) WriteControl(peer RemotePeer, pkt Marshaller) {
	enc := peer.EncryptControlPacket(pkt, w.cBuf2, w.cBuf1)

	if peer.Direct {
		w.writeTo(enc, peer.DirectAddr)
		return
	}

	rt := w.rt.Load()
	relay, ok := rt.GetRelay()
	if !ok {
		w.logf("Failed to send control to %d. No relay.", peer.IP)
		return
	}

	enc = relay.EncryptDataPacket(peer.IP, enc, w.cBuf2)
	w.writeTo(enc, relay.DirectAddr)
}

func (w *ConnWriter) writeTo(pkt []byte, addr netip.AddrPort) {
	w.wLock.Lock()
	if _, err := w.writeToUDPAddrPort(pkt, addr); err != nil {
		w.logf("Failed to write to UDP port: %v", err)
	}
	w.wLock.Unlock()
}

func (w *ConnWriter) logf(s string, args ...any) {
	log.Printf("[ConnWriter] "+s, args...)
}