MARTe-Integrated-Components/Common/Client/go/wshub/sources.go

package wshub

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"marte2/common/udpsprotocol"
)

// ─── Source configuration ─────────────────────────────────────────────────────

// SourceConfig is the serialisable description of one data source (for file save/load).
type SourceConfig struct {
	Label          string `json:"label"`
	Addr           string `json:"addr"`
	MulticastGroup string `json:"multicastGroup,omitempty"`
	DataPort       int    `json:"dataPort,omitempty"`
}

// managedSource is the SourceManager's view of one running source.
type managedSource struct {
	id             string
	label          string
	addr           string
	multicastGroup string
	dataPort       int
	client         *UDPClient
}

// SourceManager owns the lifecycle of all active data sources.
type SourceManager struct {
	mu       sync.RWMutex
	sources  map[string]*managedSource
	hub      *Hub
	filePath string
	nextID   atomic.Int32
}

// NewSourceManager creates a SourceManager bound to the given hub.
func NewSourceManager(hub *Hub, filePath string) *SourceManager {
	return &SourceManager{
		sources:  make(map[string]*managedSource),
		hub:      hub,
		filePath: filePath,
	}
}

func (sm *SourceManager) genID() string {
	return fmt.Sprintf("s%d", sm.nextID.Add(1))
}

// Add creates a new source and starts connecting.
func (sm *SourceManager) Add(label, addr, multicastGroup string, dataPort int) string {
	if label == "" {
		label = addr
	}
	id := sm.genID()
	c := NewUDPClient(addr, id, sm.hub, multicastGroup, dataPort)
	ms := &managedSource{
		id: id, label: label, addr: addr,
		multicastGroup: multicastGroup, dataPort: dataPort,
		client: c,
	}

	sm.mu.Lock()
	sm.sources[id] = ms
	sm.mu.Unlock()

	sm.hub.AddSource(id, label, addr)
	go c.Run()
	return id
}

// Remove stops the source and removes it from the hub.
func (sm *SourceManager) Remove(id string) {
	sm.mu.Lock()
	ms, ok := sm.sources[id]
	if ok {
		delete(sm.sources, id)
	}
	sm.mu.Unlock()

	if ok {
		ms.client.Stop()
		sm.hub.RemoveSource(id)
	}
}

// Save writes the current source list to filePath.
func (sm *SourceManager) Save() error {
	if sm.filePath == "" {
		return fmt.Errorf("no sources-file configured")
	}
	sm.mu.RLock()
	cfgs := make([]SourceConfig, 0, len(sm.sources))
	for _, ms := range sm.sources {
		cfgs = append(cfgs, SourceConfig{
			Label:          ms.label,
			Addr:           ms.addr,
			MulticastGroup: ms.multicastGroup,
			DataPort:       ms.dataPort,
		})
	}
	sm.mu.RUnlock()

	data, err := json.MarshalIndent(cfgs, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(sm.filePath, data, 0644)
}

// Load reads sources from path and adds them.
func (sm *SourceManager) Load(path string) error {
	data, err := os.ReadFile(path)
	if err != nil {
		return err
	}
	var cfgs []SourceConfig
	if err := json.Unmarshal(data, &cfgs); err != nil {
		return err
	}
	sm.filePath = path
	for _, cfg := range cfgs {
		sm.Add(cfg.Label, cfg.Addr, cfg.MulticastGroup, cfg.DataPort)
	}
	return nil
}

// ParseSourceArg parses "label@host:port" or "host:port".
func ParseSourceArg(s string) (label, addr string) {
	label, addr, _, _ = ParseSourceArgFull(s)
	return
}

// ParseSourceArgFull parses "[label@]host:port[/multicastGroup:dataPort]".
func ParseSourceArgFull(s string) (label, addr, multicastGroup string, dataPort int) {
	s = strings.TrimSpace(s)
	rest := s
	if idx := strings.Index(s, "@"); idx >= 0 {
		label = strings.TrimSpace(s[:idx])
		rest = strings.TrimSpace(s[idx+1:])
	}
	if idx := strings.Index(rest, "/"); idx >= 0 {
		addr = strings.TrimSpace(rest[:idx])
		mcastPart := strings.TrimSpace(rest[idx+1:])
		if lastColon := strings.LastIndex(mcastPart, ":"); lastColon >= 0 {
			multicastGroup = strings.TrimSpace(mcastPart[:lastColon])
			dataPort, _ = strconv.Atoi(strings.TrimSpace(mcastPart[lastColon+1:]))
		} else {
			multicastGroup = mcastPart
		}
	} else {
		addr = rest
	}
	return
}

// ─── UDPClient ────────────────────────────────────────────────────────────────

const (
	silenceTimeout = 5 * time.Second
	reconnectDelay = 2 * time.Second
	readBufSize    = 65536
	udpRcvBufSize  = 8 * 1024 * 1024
)

// UDPClient manages the connection to one MARTe2 streamer source.
type UDPClient struct {
	serverAddr     string
	sourceID       string
	hub            *Hub
	multicastGroup string
	dataPort       int
	stopCh         chan struct{}
}

// NewUDPClient creates a UDPClient bound to a specific source ID.
func NewUDPClient(serverAddr, sourceID string, hub *Hub, multicastGroup string, dataPort int) *UDPClient {
	return &UDPClient{
		serverAddr:     serverAddr,
		sourceID:       sourceID,
		hub:            hub,
		multicastGroup: multicastGroup,
		dataPort:       dataPort,
		stopCh:         make(chan struct{}),
	}
}

// Stop asks the client to disconnect and exit.
func (u *UDPClient) Stop() {
	close(u.stopCh)
}

// Run is the main loop; it reconnects automatically if the server goes silent.
func (u *UDPClient) Run() {
	for {
		select {
		case <-u.stopCh:
			return
		default:
		}

		u.hub.SetSourceState(u.sourceID, "connecting")
		log.Printf("[%s] connecting to %s", u.sourceID, u.serverAddr)

		var err error
		if u.multicastGroup != "" {
			err = u.runMulticastSession()
		} else {
			err = u.runSession()
		}
		if err != nil {
			log.Printf("[%s] session ended: %v", u.sourceID, err)
		}
		u.hub.SetSourceState(u.sourceID, "disconnected")

		select {
		case <-u.stopCh:
			return
		case <-time.After(reconnectDelay):
		}
	}
}

// runSession opens a UDP socket, sends CONNECT, reads data until silent or error.
func (u *UDPClient) runSession() error {
	conn, err := net.ListenUDP("udp4", &net.UDPAddr{})
	if err != nil {
		return err
	}
	defer conn.Close()

	if err := conn.SetReadBuffer(udpRcvBufSize); err != nil {
		log.Printf("[%s] udp: SetReadBuffer: %v", u.sourceID, err)
	}

	serverAddr, err := net.ResolveUDPAddr("udp4", u.serverAddr)
	if err != nil {
		return err
	}

	if _, err := conn.WriteToUDP(udpsprotocol.BuildConnectPacket(), serverAddr); err != nil {
		return err
	}
	log.Printf("[%s] udp: sent CONNECT", u.sourceID)

	reassembler := udpsprotocol.NewReassembler(2 * time.Second)
	buf := make([]byte, readBufSize)
	var currentSigs []udpsprotocol.SignalInfo
	var currentPublishMode uint8

	for {
		conn.SetReadDeadline(time.Now().Add(silenceTimeout))

		n, _, err := conn.ReadFromUDP(buf)
		arrivalTime := time.Now()
		if err != nil {
			conn.WriteToUDP(udpsprotocol.BuildDisconnectPacket(), serverAddr)
			return err
		}

		if n < udpsprotocol.HeaderSize {
			log.Printf("[%s] udp: short datagram (%d bytes), skipping", u.sourceID, n)
			continue
		}

		hdr, err := udpsprotocol.ParseHeader(buf[:n])
		if err != nil {
			log.Printf("[%s] udp: parse header: %v", u.sourceID, err)
			continue
		}

		payload := make([]byte, n-udpsprotocol.HeaderSize)
		copy(payload, buf[udpsprotocol.HeaderSize:n])

		complete, ok := reassembler.AddFragment(hdr, payload)
		if hdr.Type == udpsprotocol.PktData {
			u.hub.RecordDataFragment(u.sourceID, hdr.Counter, n, arrivalTime.UnixNano(), ok)
		}
		if !ok {
			continue
		}

		switch hdr.Type {
		case udpsprotocol.PktConfig:
			sigs, pm, err := udpsprotocol.ParseConfig(complete)
			if err != nil {
				log.Printf("[%s] udp: parse config: %v", u.sourceID, err)
				continue
			}
			currentSigs = sigs
			currentPublishMode = pm
			log.Printf("[%s] udp: received CONFIG (%d signals, publishMode=%d)", u.sourceID, len(sigs), pm)
			u.hub.SetSourceState(u.sourceID, "connected")
			u.hub.UpdateConfigForSource(u.sourceID, sigs)

		case udpsprotocol.PktData:
			if len(currentSigs) == 0 {
				continue
			}
			samples, err := udpsprotocol.ParseData(complete, currentSigs, currentPublishMode, arrivalTime)
			if err != nil {
				log.Printf("[%s] udp: parse data: %v", u.sourceID, err)
				continue
			}
			for _, s := range samples {
				u.hub.PushDataForSource(u.sourceID, s)
			}

		case udpsprotocol.PktACK:
			log.Printf("[%s] udp: received ACK (counter=%d)", u.sourceID, hdr.Counter)

		case udpsprotocol.PktDisconnect:
			log.Printf("[%s] udp: server sent DISCONNECT", u.sourceID)
			return nil

		default:
			log.Printf("[%s] udp: unknown packet type %d", u.sourceID, hdr.Type)
		}

		select {
		case <-u.stopCh:
			conn.WriteToUDP(udpsprotocol.BuildDisconnectPacket(), serverAddr)
			return nil
		default:
		}
	}
}

// runMulticastSession handles the multicast mode session.
func (u *UDPClient) runMulticastSession() error {
	tcpAddr, err := net.ResolveTCPAddr("tcp4", u.serverAddr)
	if err != nil {
		return err
	}

	tcpConn, err := net.DialTCP("tcp4", nil, tcpAddr)
	if err != nil {
		return err
	}
	defer tcpConn.Close()

	if _, err := tcpConn.Write(udpsprotocol.BuildConnectPacket()); err != nil {
		return err
	}
	log.Printf("[%s] tcp: sent CONNECT to %s", u.sourceID, u.serverAddr)

	hdrBuf := make([]byte, udpsprotocol.HeaderSize)
	if _, err := io.ReadFull(tcpConn, hdrBuf); err != nil {
		return err
	}
	cfgHdr, err := udpsprotocol.ParseHeader(hdrBuf)
	if err != nil {
		return err
	}
	if cfgHdr.Type != udpsprotocol.PktConfig {
		return net.ErrClosed
	}
	cfgPayload := make([]byte, cfgHdr.PayloadBytes)
	if cfgHdr.PayloadBytes > 0 {
		if _, err := io.ReadFull(tcpConn, cfgPayload); err != nil {
			return err
		}
	}
	currentSigs, currentPublishMode, err := udpsprotocol.ParseConfig(cfgPayload)
	if err != nil {
		return err
	}
	log.Printf("[%s] tcp: received CONFIG (%d signals, publishMode=%d)", u.sourceID, len(currentSigs), currentPublishMode)
	u.hub.SetSourceState(u.sourceID, "connected")
	u.hub.UpdateConfigForSource(u.sourceID, currentSigs)

	mcastPort := u.dataPort
	if mcastPort == 0 {
		mcastPort = tcpAddr.Port + 1
	}

	mcastIP := net.ParseIP(u.multicastGroup)
	if mcastIP == nil {
		return &net.AddrError{Err: "invalid multicast group IP", Addr: u.multicastGroup}
	}
	mcastAddr := &net.UDPAddr{IP: mcastIP, Port: mcastPort}
	mcastConn, err := net.ListenMulticastUDP("udp4", nil, mcastAddr)
	if err != nil {
		return err
	}
	defer mcastConn.Close()
	if err := mcastConn.SetReadBuffer(udpRcvBufSize); err != nil {
		log.Printf("[%s] multicast SetReadBuffer: %v", u.sourceID, err)
	}
	log.Printf("[%s] joined multicast %s:%s", u.sourceID, u.multicastGroup, strconv.Itoa(mcastPort))

	tcpDone := make(chan error, 1)
	go func() {
		buf := make([]byte, udpsprotocol.HeaderSize+64)
		for {
			n, readErr := tcpConn.Read(buf)
			if readErr != nil {
				tcpDone <- readErr
				return
			}
			if n >= udpsprotocol.HeaderSize {
				hdr, parseErr := udpsprotocol.ParseHeader(buf[:n])
				if parseErr == nil && hdr.Type == udpsprotocol.PktDisconnect {
					tcpDone <- nil
					return
				}
			}
		}
	}()

	reassembler := udpsprotocol.NewReassembler(2 * time.Second)
	buf := make([]byte, readBufSize)

	for {
		mcastConn.SetReadDeadline(time.Now().Add(silenceTimeout))
		n, _, readErr := mcastConn.ReadFromUDP(buf)
		arrivalTime := time.Now()
		if readErr != nil {
			select {
			case <-tcpDone:
				return nil
			default:
			}
			tcpConn.Write(udpsprotocol.BuildDisconnectPacket())
			return readErr
		}

		if n < udpsprotocol.HeaderSize {
			continue
		}
		hdr, parseErr := udpsprotocol.ParseHeader(buf[:n])
		if parseErr != nil {
			log.Printf("[%s] multicast: parse header: %v", u.sourceID, parseErr)
			continue
		}

		payload := make([]byte, n-udpsprotocol.HeaderSize)
		copy(payload, buf[udpsprotocol.HeaderSize:n])

		complete, ok := reassembler.AddFragment(hdr, payload)
		if hdr.Type == udpsprotocol.PktData {
			u.hub.RecordDataFragment(u.sourceID, hdr.Counter, n, arrivalTime.UnixNano(), ok)
		}
		if !ok {
			continue
		}

		if hdr.Type == udpsprotocol.PktData {
			if len(currentSigs) == 0 {
				continue
			}
			samples, parseErr := udpsprotocol.ParseData(complete, currentSigs, currentPublishMode, arrivalTime)
			if parseErr != nil {
				log.Printf("[%s] multicast: parse data: %v", u.sourceID, parseErr)
				continue
			}
			for _, s := range samples {
				u.hub.PushDataForSource(u.sourceID, s)
			}
		}

		select {
		case <-u.stopCh:
			tcpConn.Write(udpsprotocol.BuildDisconnectPacket())
			return nil
		case tcpErr := <-tcpDone:
			log.Printf("[%s] tcp control closed: %v", u.sourceID, tcpErr)
			return nil
		default:
		}
	}
}