x0gp/server/cmd/poc-server/video_udp.go

package main

import (
	"bufio"
	"context"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

type deviceState struct {
	addr         *net.UDPAddr
	latestFrame  []byte
	currentFrame []byte
	lastSeen     time.Time
}

// UDPVideoService handles UDP video frame collection and command forwarding.
type UDPVideoService struct {
	logger     *slog.Logger
	listenAddr string
	espCmdPort int

	mu      sync.RWMutex
	devices map[uint8]*deviceState

	clientsMu sync.RWMutex
	clients   map[uint8]map[chan []byte]struct{}

	conn *net.UDPConn
}

// NewUDPVideoService initializes a new UDPVideoService.
func NewUDPVideoService(logger *slog.Logger, listenAddr string, espCmdPort int) *UDPVideoService {
	return &UDPVideoService{
		logger:     logger,
		listenAddr: listenAddr,
		espCmdPort: espCmdPort,
		devices:    make(map[uint8]*deviceState),
		clients:    make(map[uint8]map[chan []byte]struct{}),
	}
}

// Start runs the UDP receiver loop (gracefully shut down by ctx).
// Stdin command reader runs in a separate un-tracked goroutine so it doesn't block shutdown.
func (s *UDPVideoService) Start(ctx context.Context, wg *sync.WaitGroup) {
	addr, err := net.ResolveUDPAddr("udp", s.listenAddr)
	if err != nil {
		s.logger.Error("failed to resolve UDP address", "listenAddr", s.listenAddr, "err", err)
		return
	}

	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		s.logger.Error("failed to listen UDP", "listenAddr", s.listenAddr, "err", err)
		return
	}
	s.conn = conn

	s.logger.Info("UDP video server started. Listening for video frames...", "listenAddr", s.listenAddr)

	// Goroutine for handling UDP receive loop
	wg.Add(1)
	go func() {
		defer wg.Done()
		defer conn.Close()

		buf := make([]byte, 65535)

		for {
			select {
			case <-ctx.Done():
				s.logger.Info("stopping UDP video receiver loop")
				return
			default:
			}

			// Use a read deadline to check for context cancellation periodically
			_ = conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
			n, remoteAddr, err := conn.ReadFromUDP(buf)
			if err != nil {
				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
					continue
				}
				s.logger.Debug("UDP read error", "err", err)
				continue
			}

			if n < 1 {
				continue
			}
			deviceID := buf[0] // first byte is DEVICE_ID (0..127)
			payload := buf[1:n]

			s.mu.Lock()
			dev, ok := s.devices[deviceID]
			if !ok {
				dev = &deviceState{
					addr:     remoteAddr,
					lastSeen: time.Now(),
				}
				s.devices[deviceID] = dev
				s.logger.Info("ESP32 video source discovered", "device_id", deviceID, "ip", remoteAddr.IP.String())
			} else {
				if dev.addr == nil || dev.addr.IP.String() != remoteAddr.IP.String() {
					dev.addr = remoteAddr
					s.logger.Info("ESP32 video source IP updated", "device_id", deviceID, "ip", remoteAddr.IP.String())
				}
				dev.lastSeen = time.Now()
			}

			// Reassemble JPEG frame for this device
			if len(payload) >= 2 && payload[0] == 0xFF && payload[1] == 0xD8 {
				dev.currentFrame = make([]byte, 0, 65535)
			}
			dev.currentFrame = append(dev.currentFrame, payload...)

			if len(payload) >= 2 && payload[len(payload)-2] == 0xFF && payload[len(payload)-1] == 0xD9 {
				// Frame fully assembled
				s.broadcastFrame(deviceID, dev.currentFrame)
			}
			s.mu.Unlock()
		}
	}()

	// Stdin command reader (run in background, not tracked by wg to avoid blocking shutdown)
	go func() {
		reader := bufio.NewReader(os.Stdin)
		s.logger.Info("Console command reader active. Enter commands (e.g., '1 left', '1 start', or 'left') in console:")

		for {
			text, err := reader.ReadString('\n')
			if err != nil {
				// Stdin closed or EOF (e.g. non-interactive/background mode)
				s.logger.Debug("stdin closed, stopping console command reader", "err", err)
				return
			}

			text = strings.TrimSpace(text)
			if text == "" {
				continue
			}

			// Format: "ID command" (e.g. "1 left") or "command" (defaults to the first discovered/available device)
			parts := strings.SplitN(text, " ", 2)
			var deviceID uint8
			var cmd string
			hasDevice := false

			if len(parts) == 2 {
				if idVal, err := strconv.Atoi(parts[0]); err == nil && idVal >= 0 && idVal <= 127 {
					deviceID = uint8(idVal)
					cmd = parts[1]
					hasDevice = true
				}
			}

			if !hasDevice {
				cmd = text
				// Find the first discovered/active device
				s.mu.RLock()
				for id := range s.devices {
					deviceID = id
					hasDevice = true
					break
				}
				s.mu.RUnlock()
			}

			if !hasDevice {
				s.logger.Warn("cannot send command: no ESP32 devices discovered yet")
				continue
			}

			err = s.SendCommand(deviceID, cmd)
			if err != nil {
				s.logger.Error("failed to send command to ESP32", "device_id", deviceID, "cmd", cmd, "err", err)
			} else {
				s.logger.Info("command sent to ESP32", "device_id", deviceID, "cmd", cmd)
			}
		}
	}()
}

// SendCommand sends a command string to the ESP32 representing deviceID.
func (s *UDPVideoService) SendCommand(deviceID uint8, cmd string) error {
	s.mu.RLock()
	dev, ok := s.devices[deviceID]
	conn := s.conn
	s.mu.RUnlock()

	if !ok || dev.addr == nil {
		return fmt.Errorf("ESP32 device %d address unknown (no video packets received yet)", deviceID)
	}
	if conn == nil {
		return fmt.Errorf("UDP service not started")
	}

	cmdAddr := &net.UDPAddr{
		IP:   dev.addr.IP,
		Port: s.espCmdPort,
	}

	_, err := conn.WriteToUDP([]byte(cmd), cmdAddr)
	return err
}

func (s *UDPVideoService) broadcastFrame(deviceID uint8, frame []byte) {
	// dev.latestFrame is updated inside the caller's lock (Start loop) to avoid races
	if dev, ok := s.devices[deviceID]; ok {
		dev.latestFrame = make([]byte, len(frame))
		copy(dev.latestFrame, frame)
	}

	s.clientsMu.RLock()
	defer s.clientsMu.RUnlock()
	deviceClients, ok := s.clients[deviceID]
	if ok {
		for ch := range deviceClients {
			select {
			case ch <- frame:
			default:
				// Drop frame if client is slow to read
			}
		}
	}
}

func (s *UDPVideoService) registerClient(deviceID uint8, ch chan []byte) {
	s.clientsMu.Lock()
	defer s.clientsMu.Unlock()
	if s.clients[deviceID] == nil {
		s.clients[deviceID] = make(map[chan []byte]struct{})
	}
	s.clients[deviceID][ch] = struct{}{}
}

func (s *UDPVideoService) unregisterClient(deviceID uint8, ch chan []byte) {
	s.clientsMu.Lock()
	defer s.clientsMu.Unlock()
	if s.clients[deviceID] != nil {
		delete(s.clients[deviceID], ch)
		if len(s.clients[deviceID]) == 0 {
			delete(s.clients, deviceID)
		}
	}
}

// StreamHandler godoc
// @Summary      Get MJPEG video stream from ESP32
// @Description  Streams video frames received via UDP as multipart/x-mixed-replace. Suitable for raw img tags in browser.
// @Tags         video
// @Produce      multipart/x-mixed-replace
// @Param        id query int false "Device ID (0..127)"
// @Success      200 {string} string "MJPEG stream"
// @Router       /api/video/stream [get]
func (s *UDPVideoService) StreamHandler() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		idVal := r.URL.Query().Get("id")
		deviceID := uint8(1) // default device ID is 1 as per client.ino
		if idVal != "" {
			if idInt, err := strconv.Atoi(idVal); err == nil && idInt >= 0 && idInt <= 127 {
				deviceID = uint8(idInt)
			} else {
				http.Error(w, "invalid device id", http.StatusBadRequest)
				return
			}
		}

		w.Header().Set("Content-Type", "multipart/x-mixed-replace; boundary=frame")
		w.Header().Set("Cache-Control", "no-cache, private, max-age=0, no-store, must-revalidate")
		w.Header().Set("Connection", "keep-alive")
		w.Header().Set("Pragma", "no-cache")

		frameCh := make(chan []byte, 10)
		s.registerClient(deviceID, frameCh)
		defer s.unregisterClient(deviceID, frameCh)

		// Send the latest frame immediately if we have one
		s.mu.RLock()
		var initFrame []byte
		if dev, ok := s.devices[deviceID]; ok {
			initFrame = dev.latestFrame
		}
		s.mu.RUnlock()
		if len(initFrame) > 0 {
			_, _ = fmt.Fprintf(w, "--frame\r\nContent-Type: image/jpeg\r\nContent-Length: %d\r\n\r\n", len(initFrame))
			_, _ = w.Write(initFrame)
			_, _ = w.Write([]byte("\r\n"))
			if flusher, ok := w.(http.Flusher); ok {
				flusher.Flush()
			}
		}

		for {
			select {
			case <-r.Context().Done():
				return
			case frame, ok := <-frameCh:
				if !ok {
					return
				}
				_, err := fmt.Fprintf(w, "--frame\r\nContent-Type: image/jpeg\r\nContent-Length: %d\r\n\r\n", len(frame))
				if err != nil {
					return
				}
				_, err = w.Write(frame)
				if err != nil {
					return
				}
				_, err = w.Write([]byte("\r\n"))
				if err != nil {
					return
				}
				if flusher, ok := w.(http.Flusher); ok {
					flusher.Flush()
				}
			}
		}
	}
}

// ControlHandler godoc
// @Summary      Send command to ESP32
// @Description  Sends a text command (e.g. start, stop, left, right) to the ESP32 command port (UDP 8888).
// @Tags         video
// @Param        id query int false "Device ID (0..127)"
// @Param        cmd query string true "Command text"
// @Produce      json
// @Success      200 {object} map[string]interface{} "Command sent"
// @Failure      400 {string} string "Missing cmd parameter"
// @Failure      500 {string} string "Internal error"
// @Router       /api/video/control [post]
func (s *UDPVideoService) ControlHandler() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "Only POST is allowed")
			return
		}
		cmd := r.URL.Query().Get("cmd")
		if cmd == "" {
			cmd = r.FormValue("cmd")
		}
		if cmd == "" {
			writeError(w, http.StatusBadRequest, "bad_request", "Missing 'cmd' parameter")
			return
		}

		idVal := r.URL.Query().Get("id")
		if idVal == "" {
			idVal = r.FormValue("id")
		}
		deviceID := uint8(1)
		if idVal != "" {
			if idInt, err := strconv.Atoi(idVal); err == nil && idInt >= 0 && idInt <= 127 {
				deviceID = uint8(idInt)
			} else {
				writeError(w, http.StatusBadRequest, "bad_request", "invalid device id")
				return
			}
		}

		err := s.SendCommand(deviceID, cmd)
		if err != nil {
			writeError(w, http.StatusInternalServerError, "internal_error", err.Error())
			return
		}

		writeJSON(w, http.StatusOK, map[string]any{
			"status":    "success",
			"device_id": deviceID,
			"command":   cmd,
		})
	}
}