package peer

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

	"golang.org/x/crypto/nacl/sign"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
)

const (
	// Beacon payload layout (after the 64-byte nacl/sign prefix):
	//   [0]      final octet of the sender's VPN IP
	//   [1:33]   WG public key
	//   [33:35]  WG listen port (big-endian uint16)
	//   [35:43]  send time, Unix seconds (big-endian int64) — freshness/replay gate
	mcBeaconLen          = 43
	mcSignedBeaconLen    = sign.Overhead + mcBeaconLen // 64-byte nacl/sign prefix + payload
	mcBroadcastInterval  = 32 * time.Second
	mcErrorRetryInterval = 16 * time.Second
	// mcBeaconMaxAge bounds how far a beacon's timestamp may be from now (in
	// either direction) before the reader drops it: it tolerates modest clock
	// skew between roughly-synchronized (NTP) hosts and caps the replay window.
	mcBeaconMaxAge = 60 * time.Second
)

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

// RunMCWriter broadcasts a signed 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, signKey *[64]byte) {
	for {
		runMCWriterInner(selfVPNIP, pubKey, wgPort, signKey)
		time.Sleep(mcErrorRetryInterval)
	}
}

func runMCWriterInner(selfVPNIP netip.Addr, pubKey wgtypes.Key, wgPort uint16, signKey *[64]byte) {
	conn, err := net.ListenMulticastUDP("udp", nil, mcAddr)
	if err != nil {
		log.Printf("[MCWriter] bind: %v", err)
		return
	}
	defer conn.Close()

	// Re-sign on each send so the timestamp is fresh; a stale timestamp would be
	// dropped by receivers' freshness gate.
	send := func() error {
		payload := buildBeacon(selfVPNIP, pubKey, wgPort, time.Now().Unix())
		signed := sign.Sign(nil, payload, signKey)
		_, err := conn.WriteToUDP(signed, mcAddr)
		return err
	}

	if err := send(); err != nil {
		log.Printf("[MCWriter] write: %v", err)
	}

	for range time.Tick(mcBroadcastInterval) {
		if err := send(); err != nil {
			log.Printf("[MCWriter] write: %v", err)
			return
		}
	}
}

func buildBeacon(selfVPNIP netip.Addr, pubKey wgtypes.Key, wgPort uint16, ts int64) []byte {
	beacon := make([]byte, mcBeaconLen)
	beacon[0] = selfVPNIP.As4()[3]
	copy(beacon[1:33], pubKey[:])
	binary.BigEndian.PutUint16(beacon[33:35], wgPort)
	binary.BigEndian.PutUint64(beacon[35:43], uint64(ts))
	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, mcSignedBeaconLen+1) // +1 to detect oversized packets
	netAddr := vpnNet.Addr().As4()

	for {
		conn.SetReadDeadline(time.Now().Add(32 * time.Second))
		n, src, err := conn.ReadFromUDPAddrPort(buf)
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Timeout() {
				continue
			}
			return fmt.Errorf("read: %w", err)
		}
		if n != mcSignedBeaconLen {
			continue
		}

		// Cheap pre-filters on the unverified payload, before the costly
		// signature check in the event loop: skip our own beacon, and drop
		// stale ones (replay/freshness gate). The timestamp is authenticated by
		// sign.Open later, so a forged-fresh timestamp still fails there — this
		// only spares us verifying old replays.
		octets := netAddr
		octets[3] = buf[sign.Overhead]
		vpnIP := netip.AddrFrom4(octets)
		if vpnIP == selfVPNIP {
			continue
		}

		if age := beaconAge(buf, time.Now()); age > mcBeaconMaxAge || age < -mcBeaconMaxAge {
			continue
		}

		signed := make([]byte, mcSignedBeaconLen)
		copy(signed, buf[:mcSignedBeaconLen])

		ch <- MulticastEvent{
			signed: signed,
			src:    src.Addr().Unmap(),
		}
	}
}