artmap/artnet/discovery.go

package artnet

import (
	"log"
	"net"
	"sync"
	"time"
)

// Node represents a discovered ArtNet node
type Node struct {
	IP          net.IP
	Port        uint16
	ShortName   string
	LongName    string
	Universes   []Universe
	LastSeen    time.Time
	CanTransmit bool
}

// Discovery handles ArtNet node discovery
type Discovery struct {
	sender       *Sender
	receiver     *Receiver
	nodes        map[string]*Node // keyed by IP string
	nodesMu      sync.RWMutex
	localIP      [4]byte
	localMAC     [6]byte
	broadcast    net.IP
	shortName    string
	longName     string
	inputUnivs   []Universe // universes we transmit TO (SwIn)
	outputUnivs  []Universe // universes we receive FROM (SwOut)
	pollTargets  []*net.UDPAddr
	done         chan struct{}
}

// NewDiscovery creates a new discovery handler
func NewDiscovery(sender *Sender, shortName, longName string, inputUnivs, outputUnivs []Universe, pollTargets []*net.UDPAddr) *Discovery {
	return &Discovery{
		sender:      sender,
		nodes:       make(map[string]*Node),
		shortName:   shortName,
		longName:    longName,
		inputUnivs:  inputUnivs,
		outputUnivs: outputUnivs,
		pollTargets: pollTargets,
		done:        make(chan struct{}),
	}
}

// Start begins periodic discovery
func (d *Discovery) Start() {
	d.detectInterface()
	go d.pollLoop()
}

// Stop stops discovery
func (d *Discovery) Stop() {
	close(d.done)
}

func (d *Discovery) pollLoop() {
	// Send initial poll to all targets
	d.sendPolls()

	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	cleanupTicker := time.NewTicker(30 * time.Second)
	defer cleanupTicker.Stop()

	for {
		select {
		case <-d.done:
			return
		case <-ticker.C:
			d.sendPolls()
		case <-cleanupTicker.C:
			d.cleanup()
		}
	}
}

func (d *Discovery) sendPolls() {
	for _, target := range d.pollTargets {
		if err := d.sender.SendPoll(target); err != nil {
			log.Printf("[->artnet] poll error: dst=%s err=%v", target.IP, err)
		}
	}
}

func (d *Discovery) cleanup() {
	d.nodesMu.Lock()
	defer d.nodesMu.Unlock()

	cutoff := time.Now().Add(-60 * time.Second)
	for ip, node := range d.nodes {
		if node.LastSeen.Before(cutoff) {
			log.Printf("[artnet] node timeout ip=%s name=%s", ip, node.ShortName)
			delete(d.nodes, ip)
		}
	}
}

// HandlePollReply processes an incoming ArtPollReply
func (d *Discovery) HandlePollReply(src *net.UDPAddr, pkt *PollReplyPacket) {
	d.nodesMu.Lock()
	defer d.nodesMu.Unlock()

	ip := src.IP.String()

	// Skip our own replies
	localIP := net.IP(d.localIP[:])
	if src.IP.Equal(localIP) {
		return
	}

	// Parse universes from SwOut
	var universes []Universe
	numPorts := int(pkt.NumPortsLo)
	if numPorts > 4 {
		numPorts = 4
	}

	for i := 0; i < numPorts; i++ {
		// Check if port can output DMX
		if pkt.PortTypes[i]&0x80 != 0 {
			u := NewUniverse(pkt.NetSwitch, pkt.SubSwitch, pkt.SwOut[i])
			universes = append(universes, u)
		}
	}

	shortName := string(pkt.ShortName[:])
	// Trim null bytes
	for i, b := range pkt.ShortName {
		if b == 0 {
			shortName = string(pkt.ShortName[:i])
			break
		}
	}

	longName := string(pkt.LongName[:])
	for i, b := range pkt.LongName {
		if b == 0 {
			longName = string(pkt.LongName[:i])
			break
		}
	}

	node, exists := d.nodes[ip]
	if !exists {
		node = &Node{
			IP:   src.IP,
			Port: pkt.Port, // Use port from packet, not UDP source port
		}
		d.nodes[ip] = node
	}

	node.ShortName = shortName
	node.LongName = longName
	node.LastSeen = time.Now()
	node.CanTransmit = true

	// Accumulate universes from multiple ArtPollReply packets
	// (multi-port devices send separate replies for each group of 4 ports)
	prevLen := len(node.Universes)
	for _, u := range universes {
		found := false
		for _, existing := range node.Universes {
			if existing == u {
				found = true
				break
			}
		}
		if !found {
			node.Universes = append(node.Universes, u)
		}
	}

	if !exists {
		log.Printf("[artnet] discovered ip=%s name=%s universes=%v", ip, shortName, node.Universes)
	} else if len(node.Universes) != prevLen {
		log.Printf("[artnet] updated ip=%s name=%s universes=%v", ip, shortName, node.Universes)
	}
}

// HandlePoll processes an incoming ArtPoll and responds
func (d *Discovery) HandlePoll(src *net.UDPAddr) {
	dst := &net.UDPAddr{IP: d.broadcast, Port: Port}
	d.sendPollReplies(dst, d.inputUnivs, true)
	d.sendPollReplies(dst, d.outputUnivs, false)
}

func (d *Discovery) sendPollReplies(dst *net.UDPAddr, universes []Universe, isInput bool) {
	groups := make(map[uint16][]Universe)
	for _, u := range universes {
		key := uint16(u.Net())<<8 | uint16(u.SubNet())<<4
		groups[key] = append(groups[key], u)
	}

	for _, univs := range groups {
		for i := 0; i < len(univs); i += 4 {
			end := i + 4
			if end > len(univs) {
				end = len(univs)
			}
			chunk := univs[i:end]

			pkt := BuildPollReplyPacket(d.localIP, d.localMAC, d.shortName, d.longName, chunk, isInput)
			err := d.receiver.SendTo(pkt, dst)
			if err != nil {
				log.Printf("[->artnet] pollreply error: dst=%s err=%v", dst.IP, err)
			}
		}
	}
}

// GetNodesForUniverse returns nodes that support a given universe
func (d *Discovery) GetNodesForUniverse(universe Universe) []*Node {
	d.nodesMu.RLock()
	defer d.nodesMu.RUnlock()

	var result []*Node
	for _, node := range d.nodes {
		for _, u := range node.Universes {
			if u == universe {
				result = append(result, node)
				break
			}
		}
	}

	if len(result) == 0 && len(d.nodes) > 0 {
		log.Printf("[artnet] no nodes for universe=%s, have %d nodes", universe, len(d.nodes))
		for ip, node := range d.nodes {
			log.Printf("[artnet]   node ip=%s universes=%v", ip, node.Universes)
		}
	}

	return result
}

// GetAllNodes returns all discovered nodes
func (d *Discovery) GetAllNodes() []*Node {
	d.nodesMu.RLock()
	defer d.nodesMu.RUnlock()

	result := make([]*Node, 0, len(d.nodes))
	for _, node := range d.nodes {
		result = append(result, node)
	}
	return result
}

func (d *Discovery) detectInterface() {
	d.broadcast = net.IPv4bcast

	ifaces, err := net.Interfaces()
	if err != nil {
		return
	}

	for _, iface := range ifaces {
		if iface.Flags&net.FlagLoopback != 0 || iface.Flags&net.FlagUp == 0 {
			continue
		}

		addrs, err := iface.Addrs()
		if err != nil {
			continue
		}

		for _, addr := range addrs {
			ipnet, ok := addr.(*net.IPNet)
			if !ok {
				continue
			}

			ip4 := ipnet.IP.To4()
			if ip4 == nil {
				continue
			}

			copy(d.localIP[:], ip4)

			if len(iface.HardwareAddr) == 6 {
				copy(d.localMAC[:], iface.HardwareAddr)
			}

			bcast := make(net.IP, 4)
			for i := 0; i < 4; i++ {
				bcast[i] = ip4[i] | ^ipnet.Mask[i]
			}
			d.broadcast = bcast

			return
		}
	}
}

// SetLocalIP sets the local IP for PollReply responses
func (d *Discovery) SetLocalIP(ip net.IP) {
	if ip4 := ip.To4(); ip4 != nil {
		copy(d.localIP[:], ip4)
	}
}

// SetReceiver sets the receiver for sending poll replies from port 6454
func (d *Discovery) SetReceiver(r *Receiver) {
	d.receiver = r
}