path: root/transport/mesh/discovery.go

// Package mesh - Peer discovery for military mesh networking
package mesh

import (
	"encoding/json"
	"fmt"
	"net"
	"sync"
	"time"
)

// PeerDiscovery manages automatic peer discovery via mDNS and UDP broadcast
type PeerDiscovery struct {
	localPeer    *PeerInfo
	discoveredPeers map[string]*PeerInfo // address -> peer info
	mu           sync.RWMutex
	udpConn      *net.UDPConn
	running      bool
	onPeerFound  func(*PeerInfo) // Callback when new peer found
	onPeerLost   func(string)    // Callback when peer lost
}

// PeerInfo contains information about a discovered peer
type PeerInfo struct {
	Address    string    `json:"address"`     // IP:Port
	PublicKey  string    `json:"public_key"`  // Ed25519 public key hex
	Nickname   string    `json:"nickname"`    // Optional display name
	Timestamp  time.Time `json:"timestamp"`   // Last seen
	IsRelay    bool      `json:"is_relay"`    // Can relay messages?
	Hops       int       `json:"hops"`        // Distance in hops
}

// DiscoveryConfig holds peer discovery configuration
type DiscoveryConfig struct {
	BroadcastPort    int           // UDP broadcast port
	BroadcastInterval time.Duration // How often to announce
	PeerTimeout      time.Duration // Remove peer if not seen
	EnableMDNS       bool          // Use mDNS (Bonjour)
	EnableBroadcast  bool          // Use UDP broadcast
}

// NewPeerDiscovery creates a new peer discovery instance
func NewPeerDiscovery(localPeer *PeerInfo, config *DiscoveryConfig) (*PeerDiscovery, error) {
	if config == nil {
		config = &DiscoveryConfig{
			BroadcastPort:     9876,
			BroadcastInterval: 10 * time.Second,
			PeerTimeout:       60 * time.Second,
			EnableBroadcast:   true,
			EnableMDNS:        false, // mDNS requires external library
		}
	}

	pd := &PeerDiscovery{
		localPeer:       localPeer,
		discoveredPeers: make(map[string]*PeerInfo),
	}

	// Setup UDP broadcast listener
	if config.EnableBroadcast {
		addr := &net.UDPAddr{
			IP:   net.IPv4zero,
			Port: config.BroadcastPort,
		}

		conn, err := net.ListenUDP("udp4", addr)
		if err != nil {
			return nil, fmt.Errorf("failed to create UDP listener: %w", err)
		}

		pd.udpConn = conn
	}

	return pd, nil
}

// Start starts peer discovery
func (pd *PeerDiscovery) Start(config *DiscoveryConfig) error {
	pd.mu.Lock()
	if pd.running {
		pd.mu.Unlock()
		return fmt.Errorf("discovery already running")
	}
	pd.running = true
	pd.mu.Unlock()

	if config == nil {
		config = &DiscoveryConfig{
			BroadcastPort:     9876,
			BroadcastInterval: 10 * time.Second,
			PeerTimeout:       60 * time.Second,
			EnableBroadcast:   true,
		}
	}

	// Start UDP broadcast announcer
	if config.EnableBroadcast {
		go pd.broadcastAnnouncer(config)
		go pd.broadcastListener(config)
	}

	// Start peer timeout checker
	go pd.peerTimeoutChecker(config)

	return nil
}

// Stop stops peer discovery
func (pd *PeerDiscovery) Stop() error {
	pd.mu.Lock()
	defer pd.mu.Unlock()

	if !pd.running {
		return nil
	}

	pd.running = false

	if pd.udpConn != nil {
		pd.udpConn.Close()
	}

	return nil
}

// broadcastAnnouncer periodically broadcasts our presence
func (pd *PeerDiscovery) broadcastAnnouncer(config *DiscoveryConfig) {
	ticker := time.NewTicker(config.BroadcastInterval)
	defer ticker.Stop()

	for {
		pd.mu.RLock()
		if !pd.running {
			pd.mu.RUnlock()
			return
		}
		pd.mu.RUnlock()

		<-ticker.C

		// Create announcement message
		announcement := map[string]interface{}{
			"type":       "peer_announcement",
			"address":    pd.localPeer.Address,
			"public_key": pd.localPeer.PublicKey,
			"nickname":   pd.localPeer.Nickname,
			"is_relay":   pd.localPeer.IsRelay,
			"timestamp":  time.Now().Unix(),
		}

		data, err := json.Marshal(announcement)
		if err != nil {
			continue
		}

		// Broadcast to LAN
		broadcastAddr := &net.UDPAddr{
			IP:   net.IPv4bcast,
			Port: config.BroadcastPort,
		}

		conn, err := net.DialUDP("udp4", nil, broadcastAddr)
		if err != nil {
			continue
		}

		conn.Write(data)
		conn.Close()
	}
}

// broadcastListener listens for peer announcements
func (pd *PeerDiscovery) broadcastListener(config *DiscoveryConfig) {
	buf := make([]byte, 4096)

	for {
		pd.mu.RLock()
		if !pd.running {
			pd.mu.RUnlock()
			return
		}
		pd.mu.RUnlock()

		// Set read deadline to allow periodic checks
		pd.udpConn.SetReadDeadline(time.Now().Add(1 * time.Second))

		n, remoteAddr, err := pd.udpConn.ReadFromUDP(buf)
		if err != nil {
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				continue // Timeout is expected
			}
			continue
		}

		// Parse announcement
		var announcement map[string]interface{}
		if err := json.Unmarshal(buf[:n], &announcement); err != nil {
			continue
		}

		// Ignore our own announcements
		if announcement["address"] == pd.localPeer.Address {
			continue
		}

		// Extract peer info
		peerInfo := &PeerInfo{
			Address:   fmt.Sprintf("%v", announcement["address"]),
			PublicKey: fmt.Sprintf("%v", announcement["public_key"]),
			Nickname:  fmt.Sprintf("%v", announcement["nickname"]),
			Timestamp: time.Now(),
			IsRelay:   announcement["is_relay"].(bool),
			Hops:      1, // Direct peer
		}

		// If address is missing, use source address
		if peerInfo.Address == "" || peerInfo.Address == "<nil>" {
			peerInfo.Address = remoteAddr.IP.String()
		}

		pd.addPeer(peerInfo)
	}
}

// addPeer adds or updates a discovered peer
func (pd *PeerDiscovery) addPeer(peer *PeerInfo) {
	pd.mu.Lock()
	defer pd.mu.Unlock()

	isNew := false
	if _, exists := pd.discoveredPeers[peer.Address]; !exists {
		isNew = true
	}

	pd.discoveredPeers[peer.Address] = peer

	// Call callback if new peer
	if isNew && pd.onPeerFound != nil {
		go pd.onPeerFound(peer)
	}
}

// peerTimeoutChecker removes stale peers
func (pd *PeerDiscovery) peerTimeoutChecker(config *DiscoveryConfig) {
	ticker := time.NewTicker(config.PeerTimeout / 2)
	defer ticker.Stop()

	for {
		pd.mu.RLock()
		if !pd.running {
			pd.mu.RUnlock()
			return
		}
		pd.mu.RUnlock()

		<-ticker.C

		pd.mu.Lock()
		now := time.Now()
		for addr, peer := range pd.discoveredPeers {
			if now.Sub(peer.Timestamp) > config.PeerTimeout {
				delete(pd.discoveredPeers, addr)

				// Call callback
				if pd.onPeerLost != nil {
					go pd.onPeerLost(addr)
				}
			}
		}
		pd.mu.Unlock()
	}
}

// GetPeers returns all discovered peers
func (pd *PeerDiscovery) GetPeers() []*PeerInfo {
	pd.mu.RLock()
	defer pd.mu.RUnlock()

	peers := make([]*PeerInfo, 0, len(pd.discoveredPeers))
	for _, peer := range pd.discoveredPeers {
		peers = append(peers, peer)
	}
	return peers
}

// GetPeerCount returns the number of discovered peers
func (pd *PeerDiscovery) GetPeerCount() int {
	pd.mu.RLock()
	defer pd.mu.RUnlock()
	return len(pd.discoveredPeers)
}

// SetOnPeerFound sets callback for when new peer is discovered
func (pd *PeerDiscovery) SetOnPeerFound(callback func(*PeerInfo)) {
	pd.mu.Lock()
	defer pd.mu.Unlock()
	pd.onPeerFound = callback
}

// SetOnPeerLost sets callback for when peer is lost
func (pd *PeerDiscovery) SetOnPeerLost(callback func(string)) {
	pd.mu.Lock()
	defer pd.mu.Unlock()
	pd.onPeerLost = callback
}

// FindPeerByPublicKey finds a peer by their public key
func (pd *PeerDiscovery) FindPeerByPublicKey(publicKey string) (*PeerInfo, bool) {
	pd.mu.RLock()
	defer pd.mu.RUnlock()

	for _, peer := range pd.discoveredPeers {
		if peer.PublicKey == publicKey {
			return peer, true
		}
	}
	return nil, false
}

// String returns a string representation
func (pd *PeerDiscovery) String() string {
	pd.mu.RLock()
	defer pd.mu.RUnlock()

	status := "stopped"
	if pd.running {
		status = "running"
	}

	return fmt.Sprintf("PeerDiscovery(status=%s, peers=%d)",
		status, len(pd.discoveredPeers))
}