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

package main

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"log/slog"
	"math"
	"net/http"
	"sync"
	"time"

	"github.com/pion/webrtc/v4"
	"github.com/x0gp/server/internal/catalog"
	"github.com/x0gp/server/internal/control"
	"github.com/x0gp/server/internal/lobby"
	"github.com/x0gp/server/internal/transport"
)

type driverSession struct {
	driverID    string
	deviceID    uint8
	pc          *webrtc.PeerConnection
	videoDC     *webrtc.DataChannel
	dataDC      *webrtc.DataChannel
	lastCommand string
	cancelCtx   context.CancelFunc
}

type WebRTCService struct {
	logger     *slog.Logger
	engine     *control.Engine
	lobbySvc   *lobby.Service
	videoSvc   *UDPVideoService
	catalogSvc *catalog.Service
	sessions   map[string]*driverSession
	sessionsMu sync.RWMutex
}

func NewWebRTCService(logger *slog.Logger, e *control.Engine, lobbySvc *lobby.Service, videoSvc *UDPVideoService, catalogSvc *catalog.Service) *WebRTCService {
	return &WebRTCService{
		logger:     logger,
		engine:     e,
		lobbySvc:   lobbySvc,
		videoSvc:   videoSvc,
		catalogSvc: catalogSvc,
		sessions:   make(map[string]*driverSession),
	}
}

func (s *WebRTCService) Start(ctx context.Context, wg *sync.WaitGroup) {
	wg.Add(1)
	go func() {
		defer wg.Done()
		ticker := time.NewTicker(1 * time.Second)
		defer ticker.Stop()

		for {
			select {
			case <-ctx.Done():
				s.CloseAll()
				return
			case <-ticker.C:
				s.broadcastTelemetry()
			}
		}
	}()
}

// ConnectHandler godoc
// @Summary      Establish WebRTC connection for driver-car 1-to-1 proxying
// @Description  Accepts SDP Offer, negotiates peer connection, and returns SDP Answer. Routes video and control data via DataChannels.
// @Tags         webrtc
// @Accept       json
// @Produce      json
// @Param        body body transport.WebRTCConnectRequest true "SDP Offer details"
// @Success      200 {object} transport.WebRTCConnectResponse "SDP Answer negotiation success"
// @Router       /api/webrtc/connect [post]
func (s *WebRTCService) ConnectHandler() 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
		}

		var req transport.WebRTCConnectRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeError(w, http.StatusBadRequest, "bad_request", "Invalid request JSON")
			return
		}

		if req.DriverID == "" {
			writeError(w, http.StatusBadRequest, "bad_request", "driver_id is required")
			return
		}

		// Resolve device ID for this driver
		deviceID := uint8(1) // default fallback
		if driver, err := s.lobbySvc.GetDriver(req.DriverID); err == nil && driver.CarID != nil && *driver.CarID != "" {
			if car, ok := s.catalogSvc.GetCar(*driver.CarID); ok && car.DeviceID != nil {
				deviceID = uint8(*car.DeviceID)
			} else {
				s.logger.Warn("webrtc client connected but car has no physical device_id assigned in catalog. Defaulting to device 1", "driver_id", req.DriverID, "car_id", *driver.CarID)
			}
		} else {
			s.logger.Warn("webrtc client connected but no car_id assigned in lobbySvc. Defaulting to device 1", "driver_id", req.DriverID)
		}

		// Create PeerConnection
		config := webrtc.Configuration{
			ICEServers: []webrtc.ICEServer{
				{
					URLs: []string{"stun:stun.l.google.com:19302"},
				},
			},
		}

		pc, err := webrtc.NewPeerConnection(config)
		if err != nil {
			s.logger.Error("failed to create webrtc peer connection", "err", err)
			writeError(w, http.StatusInternalServerError, "internal_error", err.Error())
			return
		}

		// Prepare context for background video pumping
		connCtx, connCancel := context.WithCancel(context.Background())

		session := &driverSession{
			driverID:  req.DriverID,
			deviceID:  deviceID,
			pc:        pc,
			cancelCtx: connCancel,
		}

		// Handle Data Channels
		pc.OnDataChannel(func(d *webrtc.DataChannel) {
			s.logger.Info("received remote data channel", "label", d.Label(), "driver_id", req.DriverID)
			s.sessionsMu.Lock()
			if d.Label() == "video" {
				session.videoDC = d
			} else if d.Label() == "data" {
				session.dataDC = d
			}
			s.sessionsMu.Unlock()

			d.OnOpen(func() {
				s.logger.Info("data channel opened", "label", d.Label(), "driver_id", req.DriverID)
				if d.Label() == "video" {
					// Start pumping video frames from UDPVideoService to the data channel
					go s.pumpVideo(connCtx, session)
				}
			})

			d.OnMessage(func(msg webrtc.DataChannelMessage) {
				if d.Label() == "data" {
					s.handleDataChannelMessage(session, msg.Data)
				}
			})
		})

		pc.OnConnectionStateChange(func(state webrtc.PeerConnectionState) {
			s.logger.Info("webrtc connection state changed", "state", state.String(), "driver_id", req.DriverID)
			if state == webrtc.PeerConnectionStateClosed || state == webrtc.PeerConnectionStateFailed {
				s.CloseSession(req.DriverID)
			}
		})

		// Set the remote SessionDescription
		err = pc.SetRemoteDescription(webrtc.SessionDescription{
			Type: webrtc.SDPTypeOffer,
			SDP:  req.SDP,
		})
		if err != nil {
			pc.Close()
			connCancel()
			s.logger.Error("failed to set remote description", "err", err)
			writeError(w, http.StatusInternalServerError, "internal_error", err.Error())
			return
		}

		// Create Answer
		answer, err := pc.CreateAnswer(nil)
		if err != nil {
			pc.Close()
			connCancel()
			s.logger.Error("failed to create answer", "err", err)
			writeError(w, http.StatusInternalServerError, "internal_error", err.Error())
			return
		}

		// Set Local Description
		err = pc.SetLocalDescription(answer)
		if err != nil {
			pc.Close()
			connCancel()
			s.logger.Error("failed to set local description", "err", err)
			writeError(w, http.StatusInternalServerError, "internal_error", err.Error())
			return
		}

		// Wait for ICE Gathering to complete (Vanilla ICE)
		gatherComplete := webrtc.GatheringCompletePromise(pc)
		<-gatherComplete

		// Save Session
		s.sessionsMu.Lock()
		// Clean up existing session if any
		if existing, ok := s.sessions[req.DriverID]; ok {
			existing.pc.Close()
			existing.cancelCtx()
		}
		s.sessions[req.DriverID] = session
		s.sessionsMu.Unlock()

		localDesc := pc.LocalDescription()
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(transport.WebRTCConnectResponse{
			SDP:  localDesc.SDP,
			Type: "answer",
		})
	}
}

func (s *WebRTCService) pumpVideo(ctx context.Context, session *driverSession) {
	s.logger.Info("starting WebRTC video pump for driver", "driver_id", session.driverID, "device_id", session.deviceID)

	frameCh := make(chan []byte, 10)
	s.videoSvc.registerClient(session.deviceID, frameCh)
	defer func() {
		s.videoSvc.unregisterClient(session.deviceID, frameCh)
		s.logger.Info("stopping WebRTC video pump for driver", "driver_id", session.driverID)
	}()

	for {
		select {
		case <-ctx.Done():
			return
		case frame, ok := <-frameCh:
			if !ok {
				return
			}
			s.sessionsMu.RLock()
			dc := session.videoDC
			s.sessionsMu.RUnlock()

			if dc != nil && dc.ReadyState() == webrtc.DataChannelStateOpen {
				if err := dc.Send(frame); err != nil {
					s.logger.Debug("failed to send video frame over datachannel", "driver_id", session.driverID, "err", err)
				}
			}
		}
	}
}

func (s *WebRTCService) handleDataChannelMessage(session *driverSession, data []byte) {
	var payload map[string]any
	if err := json.Unmarshal(data, &payload); err != nil {
		s.logger.Warn("received invalid JSON on data channel", "driver_id", session.driverID, "err", err)
		return
	}

	// 1. Steering & Throttle inputs
	in := transport.InputState{}
	hasInput := false
	if v, ok := payload["throttle"].(float64); ok {
		in.Throttle = v
		hasInput = true
	}
	if v, ok := payload["steering"].(float64); ok {
		in.Steering = v
		hasInput = true
	}
	if v, ok := payload["brake"].(float64); ok {
		in.Brake = v
		hasInput = true
	}
	if v, ok := payload["gear"].(float64); ok {
		in.Gear = int(v)
		hasInput = true
	}
	if v, ok := payload["buttons"].(float64); ok {
		in.Buttons = uint32(v)
		hasInput = true
	}

	if hasInput {
		s.logger.Debug("received inputs on webrtc data channel", "driver_id", session.driverID, "throttle", in.Throttle, "steering", in.Steering)
		car := s.engine.GetCarByDriver(session.driverID)
		if car != nil {
			car.ApplyInput(in, 1.0/60.0)
		}

		// Resolve destination device
		var destDeviceID *uint8
		if car != nil && car.DeviceID != nil {
			val := uint8(*car.DeviceID)
			destDeviceID = &val
		} else if session.deviceID != 0 {
			// Fallback to connection-time device
			destDeviceID = &session.deviceID
		}

		if destDeviceID != nil {
			var packet [5]byte

			// 1. Steer: mapping float64 [-1.0, 1.0] to [0..200] (100 is center)
			steerVal := int(math.Round((in.Steering + 1.0) * 100.0))
			if steerVal < 0 { steerVal = 0 }
			if steerVal > 200 { steerVal = 200 }
			packet[0] = uint8(steerVal)

			// 2. Throttle: mapping float64 [0.0, 1.0] to [0..100]
			throttleVal := int(math.Round(in.Throttle * 100.0))
			if throttleVal < 0 { throttleVal = 0 }
			if throttleVal > 100 { throttleVal = 100 }
			packet[1] = uint8(throttleVal)

			// 3. Brake: mapping float64 [0.0, 1.0] to [0..100]
			brakeVal := int(math.Round(in.Brake * 100.0))
			if brakeVal < 0 { brakeVal = 0 }
			if brakeVal > 100 { brakeVal = 100 }
			packet[2] = uint8(brakeVal)

			// 4. Gear: map -1 to 0 (R), 0 to 1 (N), >=1 to 2 (D)
			gearVal := 1 // Neutral by default
			if in.Gear < 0 {
				gearVal = 0 // R
			} else if in.Gear > 0 {
				gearVal = 2 // D
			}
			packet[3] = uint8(gearVal)

			// 5. Flags: bit 0 (DRS), bit 1 (Pit Limiter)
			packet[4] = uint8(in.Buttons & 0x03)

			hexCmd := hex.EncodeToString(packet[:])
			if hexCmd != session.lastCommand {
				s.logger.Info("webrtc control state transition", "driver_id", session.driverID, "old_hex", session.lastCommand, "new_hex", hexCmd)
				session.lastCommand = hexCmd
				err := s.videoSvc.SendControlPacket(*destDeviceID, packet)
				if err != nil {
					s.logger.Error("failed to forward control packet to ESP32", "device_id", *destDeviceID, "packet", packet, "err", err)
				} else {
					s.logger.Info("forwarded control packet to ESP32 via UDP", "device_id", *destDeviceID, "packet", packet)
				}
			}
		}
	}

	// Echo back message / send ACK
	s.sessionsMu.RLock()
	dc := session.dataDC
	s.sessionsMu.RUnlock()
	if dc != nil && dc.ReadyState() == webrtc.DataChannelStateOpen {
		ack, _ := json.Marshal(map[string]any{
			"type":      "ack",
			"applied":   true,
			"timestamp": time.Now().UnixMilli(),
		})
		_ = dc.SendText(string(ack))
	}
}

func (s *WebRTCService) broadcastTelemetry() {
	s.sessionsMu.RLock()
	defer s.sessionsMu.RUnlock()

	for _, session := range s.sessions {
		dc := session.dataDC
		if dc != nil && dc.ReadyState() == webrtc.DataChannelStateOpen {
			// Pull some stats/telemetry
			car := s.engine.GetCarByDriver(session.driverID)
			var speed float64
			var x, y float64
			var carID string
			if car != nil {
				speed = car.Speed
				x = car.X
				y = car.Y
				carID = car.ID
			}

			telemetry := map[string]any{
				"type":      "telemetry",
				"car_id":    carID,
				"driver_id": session.driverID,
				"speed":     speed,
				"x":         x,
				"y":         y,
				"timestamp": time.Now().UnixMilli(),
			}

			payload, _ := json.Marshal(telemetry)
			_ = dc.SendText(string(payload))
		}
	}
}

func (s *WebRTCService) CloseSession(driverID string) {
	s.sessionsMu.Lock()
	defer s.sessionsMu.Unlock()

	if session, ok := s.sessions[driverID]; ok {
		s.logger.Info("closing WebRTC session for driver", "driver_id", driverID)
		session.pc.Close()
		session.cancelCtx()
		delete(s.sessions, driverID)
	}
}

func (s *WebRTCService) CloseAll() {
	s.sessionsMu.Lock()
	defer s.sessionsMu.Unlock()

	for driverID, session := range s.sessions {
		s.logger.Info("closing WebRTC session (shutdown)", "driver_id", driverID)
		session.pc.Close()
		session.cancelCtx()
	}
	s.sessions = make(map[string]*driverSession)
}