package peer

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

type connReader struct {
	conn    udpReader
	iface   ifWriter
	sender  encryptedPacketSender
	super   controlMsgHandler
	localIP byte
	peers   [256]*atomic.Pointer[remotePeer]

	buf    []byte
	decBuf []byte
}

func newConnReader(
	conn udpReader,
	ifWriter ifWriter,
	sender encryptedPacketSender,
	super controlMsgHandler,
	localIP byte,
	peers [256]*atomic.Pointer[remotePeer],
) *connReader {
	return &connReader{
		conn:    conn,
		iface:   ifWriter,
		sender:  sender,
		super:   super,
		localIP: localIP,
		peers:   peers,
		buf:     make([]byte, bufferSize),
		decBuf:  make([]byte, bufferSize),
	}
}

func (r *connReader) Run() {
	for {
		r.handleNextPacket()
	}
}

func (r *connReader) logf(s string, args ...any) {
	log.Printf("[ConnReader] "+s, args...)
}

func (r *connReader) handleNextPacket() {
	buf := r.buf[:bufferSize]
	n, remoteAddr, err := r.conn.ReadFromUDPAddrPort(buf)
	if err != nil {
		log.Fatalf("Failed to read from UDP port: %v", err)
	}

	if n < headerSize {
		return
	}

	remoteAddr = netip.AddrPortFrom(remoteAddr.Addr().Unmap(), remoteAddr.Port())

	buf = buf[:n]
	h := parseHeader(buf)

	peer := r.peers[h.SourceIP].Load()

	switch h.StreamID {
	case controlStreamID:
		r.handleControlPacket(peer, remoteAddr, h, buf)

	case dataStreamID:
		r.handleDataPacket(peer, h, buf)

	default:
		r.logf("Unknown stream ID: %d", h.StreamID)
	}
}

func (r *connReader) handleControlPacket(
	peer *remotePeer,
	addr netip.AddrPort,
	h header,
	enc []byte,
) {
	if peer.ControlCipher == nil {
		return
	}

	if h.DestIP != r.localIP {
		r.logf("Incorrect destination IP on control packet: %d", h.DestIP)
		return
	}

	msg, err := decryptControlPacket(peer, addr, h, enc, r.decBuf)
	if err != nil {
		r.logf("Failed to decrypt control packet: %v", err)
		return
	}

	r.super.HandleControlMsg(msg)
}

func (r *connReader) handleDataPacket(peer *remotePeer, h header, enc []byte) {
	if !peer.Up {
		r.logf("Not connected (recv).")
		return
	}

	data, err := decryptDataPacket(peer, h, enc, r.decBuf)
	if err != nil {
		r.logf("Failed to decrypt data packet: %v", err)
		return
	}

	if h.DestIP == r.localIP {
		if _, err := r.iface.Write(data); err != nil {
			log.Fatalf("Failed to write to interface: %v", err)
		}
		return
	}

	destPeer := r.peers[h.DestIP].Load()
	if !destPeer.Up {
		r.logf("Not connected (relay): %d", destPeer.IP)
		return
	}

	r.sender.SendEncryptedDataPacket(data, destPeer)
}