package peer

import (
	"encoding/binary"
	"fmt"
	"log"
	"net"
	"net/netip"
	"time"

	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
)

const (
	mcBeaconLen          = 35 // 1 VPN IP byte + 32 WG pubkey + 2 WG listen port
	mcBroadcastInterval  = 32 * time.Second
	mcErrorRetryInterval = 16 * time.Second
)

var mcAddr = net.UDPAddrFromAddrPort(netip.AddrPortFrom(
	netip.AddrFrom4([4]byte{224, 0, 0, 157}),
	4560))

// RunMCWriter broadcasts a beacon on the local multicast group every
// mcBroadcastInterval so that LAN peers can discover our WireGuard endpoint.
func RunMCWriter(selfVPNIP netip.Addr, pubKey wgtypes.Key, wgPort uint16) {
	conn, err := net.ListenMulticastUDP("udp", nil, mcAddr)
	if err != nil {
		log.Fatalf("[MCWriter] bind: %v", err)
	}

	beacon := buildBeacon(selfVPNIP, pubKey, wgPort)

	for range time.Tick(mcBroadcastInterval) {
		if _, err := conn.WriteToUDP(beacon, mcAddr); err != nil {
			log.Printf("[MCWriter] write: %v", err)
		}
	}
}

func buildBeacon(selfVPNIP netip.Addr, pubKey wgtypes.Key, wgPort uint16) []byte {
	beacon := make([]byte, mcBeaconLen)
	beacon[0] = selfVPNIP.As4()[3]
	copy(beacon[1:33], pubKey[:])
	binary.BigEndian.PutUint16(beacon[33:35], wgPort)
	return beacon
}

// RunMCReader listens for multicast beacons from LAN peers and feeds
// MulticastEvents to ch.
func RunMCReader(vpnNet netip.Prefix, selfVPNIP netip.Addr, ch chan<- MulticastEvent) {
	for {
		if err := runMCReaderInner(vpnNet, selfVPNIP, ch); err != nil {
			log.Printf("[MCReader] %v", err)
		}
		time.Sleep(mcErrorRetryInterval)
	}
}

func runMCReaderInner(vpnNet netip.Prefix, selfVPNIP netip.Addr, ch chan<- MulticastEvent) error {
	conn, err := net.ListenMulticastUDP("udp", nil, mcAddr)
	if err != nil {
		return fmt.Errorf("bind: %w", err)
	}
	defer conn.Close()

	buf := make([]byte, 64)
	netAddr := vpnNet.Addr().As4()

	for {
		conn.SetReadDeadline(time.Now().Add(32 * time.Second))
		n, src, err := conn.ReadFromUDPAddrPort(buf)
		if err != nil {
			return fmt.Errorf("read: %w", err)
		}
		if n != mcBeaconLen {
			continue
		}

		octets := netAddr
		octets[3] = buf[0]
		vpnIP := netip.AddrFrom4(octets)
		if vpnIP == selfVPNIP {
			continue
		}

		pubKey, err := wgtypes.NewKey(buf[1:33])
		if err != nil {
			continue
		}

		wgPort := binary.BigEndian.Uint16(buf[33:35])
		endpoint := netip.AddrPortFrom(src.Addr().Unmap(), wgPort)

		ch <- MulticastEvent{
			pubKey:   pubKey,
			vpnIP:    vpnIP,
			endpoint: endpoint,
		}
	}
}