PilotVault/API Server/internal/hub/hub.go

// Package hub keeps the live, in-memory view of every connected device and
// fans telemetry out to dashboards over websockets. It is the drone-domain core
// of the API Server; the api package exposes it over HTTP.
package hub

import (
	"encoding/json"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

const (
	writeWait      = 10 * time.Second
	pongWait       = 60 * time.Second
	pingPeriod     = (pongWait * 9) / 10
	maxMessageSize = 1 << 20
	sendBuffer     = 256
	maxTrackPoints = 1000
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	// Dev default: accept any origin. Lock this down for production.
	CheckOrigin: func(r *http.Request) bool { return true },
}

type clientKind int

const (
	kindDevice clientKind = iota
	kindUI
)

// Client is a single websocket connection (either a device/app or a dashboard).
type Client struct {
	hub      *Hub
	conn     *websocket.Conn
	send     chan []byte
	kind     clientKind
	deviceID string
}

// Hub keeps track of all connections and the latest state per device.
type Hub struct {
	mu      sync.RWMutex
	uis     map[*Client]bool
	devices map[string]*Client      // currently-online device connections
	states  map[string]*DeviceState // last-known state, persists across reconnects
	tracks  map[string][]TrackPoint
}

// New constructs an empty Hub.
func New() *Hub {
	return &Hub{
		uis:     make(map[*Client]bool),
		devices: make(map[string]*Client),
		states:  make(map[string]*DeviceState),
		tracks:  make(map[string][]TrackPoint),
	}
}

func nowMs() int64 { return time.Now().UnixMilli() }

// ── Websocket entry points ───────────────────────────────────────────────────

// ServeDevice upgrades an incoming request into a device connection (the Fly
// App's telemetry uplink) bound to deviceID.
func (h *Hub) ServeDevice(w http.ResponseWriter, r *http.Request, deviceID string) {
	if deviceID == "" {
		deviceID = "default"
	}
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	c := &Client{hub: h, conn: conn, send: make(chan []byte, sendBuffer), kind: kindDevice, deviceID: deviceID}
	h.addDevice(c)
	log.Printf("device connected: %s", deviceID)
	go c.writePump()
	go c.readPump()
}

// ServeUI upgrades an incoming request into a dashboard connection.
func (h *Hub) ServeUI(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return
	}
	c := &Client{hub: h, conn: conn, send: make(chan []byte, sendBuffer), kind: kindUI}
	h.addUI(c)
	go c.writePump()
	go c.readPump()
}

// ── UI client lifecycle ──────────────────────────────────────────────────────

func (h *Hub) addUI(c *Client) {
	h.mu.Lock()
	h.uis[c] = true
	devices := make([]*DeviceState, 0, len(h.states))
	for _, s := range h.states {
		cp := *s
		devices = append(devices, &cp)
	}
	h.mu.Unlock()

	if msg, err := json.Marshal(ServerToUI{Type: "snapshot", Devices: devices, TS: nowMs()}); err == nil {
		c.send <- msg
	}
}

func (h *Hub) removeUI(c *Client) {
	h.mu.Lock()
	delete(h.uis, c)
	h.mu.Unlock()
}

// ── Device client lifecycle ──────────────────────────────────────────────────

func (h *Hub) addDevice(c *Client) {
	h.mu.Lock()
	h.devices[c.deviceID] = c
	s := h.states[c.deviceID]
	if s == nil {
		s = &DeviceState{DeviceID: c.deviceID}
		h.states[c.deviceID] = s
	}
	s.Online = true
	s.LastSeenMs = nowMs()
	snap := *s
	h.mu.Unlock()

	h.broadcastUI(ServerToUI{Type: "update", Device: &snap, TS: nowMs()})
}

func (h *Hub) removeDevice(c *Client) {
	h.mu.Lock()
	if h.devices[c.deviceID] == c {
		delete(h.devices, c.deviceID)
	}
	var snap *DeviceState
	if s := h.states[c.deviceID]; s != nil {
		s.Online = false
		s.Connected = false
		s.Telemetry = Telemetry{} // app stopped streaming: drop stale live telemetry
		s.LastSeenMs = nowMs()
		cp := *s
		snap = &cp
	}
	h.mu.Unlock()

	if snap != nil {
		h.broadcastUI(ServerToUI{Type: "update", Device: snap, TS: nowMs()})
	}
}

// ── Data flow ────────────────────────────────────────────────────────────────

// Ingest applies a raw event from a device and fans it out to the dashboards.
func (h *Hub) Ingest(deviceID string, raw map[string]any) {
	h.mu.Lock()
	s := h.states[deviceID]
	if s == nil {
		s = &DeviceState{DeviceID: deviceID}
		h.states[deviceID] = s
	}
	s.Online = true
	s.LastSeenMs = nowMs()

	switch raw["type"] {
	case "registration":
		if st, ok := raw["state"].(string); ok {
			s.Registration = st
		}
	case "connection":
		if c, ok := raw["connected"].(bool); ok {
			s.Connected = c
			if !c {
				s.Telemetry = Telemetry{} // drone unlinked: live telemetry is no longer valid
			}
		}
		if m, ok := raw["model"].(string); ok {
			s.Model = m
		}
	case "battery":
		if p, ok := toInt(raw["percent"]); ok {
			s.Telemetry.BatteryPercent = &p
		}
	case "telemetry":
		applyTelemetry(&s.Telemetry, raw)
		lat, okLat := toFloat(raw["latitude"])
		lng, okLng := toFloat(raw["longitude"])
		if okLat && okLng && (lat != 0 || lng != 0) {
			alt, _ := toFloat(raw["altitude"])
			h.appendTrackLocked(deviceID, TrackPoint{Lat: lat, Lng: lng, Alt: alt, TS: nowMs()})
		}
	}
	snap := *s
	h.mu.Unlock()

	h.broadcastUI(ServerToUI{Type: "update", Device: &snap, Event: raw, TS: nowMs()})
}

// appendTrackLocked must be called with h.mu held.
func (h *Hub) appendTrackLocked(deviceID string, p TrackPoint) {
	t := append(h.tracks[deviceID], p)
	if len(t) > maxTrackPoints {
		t = t[len(t)-maxTrackPoints:]
	}
	h.tracks[deviceID] = t
}

// SendCommand routes a command from the server (or a dashboard) to a device.
// It returns false if the device is not currently connected.
func (h *Hub) SendCommand(deviceID, command string, payload map[string]any) bool {
	cmd := Command{Type: "command", Command: command, Payload: payload, TS: nowMs()}
	msg, err := json.Marshal(cmd)
	if err != nil {
		return false
	}
	h.mu.RLock()
	c := h.devices[deviceID]
	h.mu.RUnlock()
	if c == nil {
		return false
	}
	select {
	case c.send <- msg:
		return true
	default:
		return false
	}
}

func (h *Hub) broadcastUI(m ServerToUI) {
	msg, err := json.Marshal(m)
	if err != nil {
		return
	}
	h.mu.RLock()
	for c := range h.uis {
		select {
		case c.send <- msg:
		default: // drop messages for a slow/stuck dashboard rather than block
		}
	}
	h.mu.RUnlock()
}

// Forget drops a device's stored state and track. Intended for clearing
// stale/offline entries; a still-online device will simply repopulate.
func (h *Hub) Forget(deviceID string) bool {
	h.mu.Lock()
	_, existed := h.states[deviceID]
	delete(h.states, deviceID)
	delete(h.tracks, deviceID)
	h.mu.Unlock()
	if existed {
		h.broadcastUI(ServerToUI{Type: "removed", DeviceID: deviceID, TS: nowMs()})
	}
	return existed
}

// OnlineCount returns the number of devices with a live websocket connection
// right now (offline/last-known states are not counted).
func (h *Hub) OnlineCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.devices)
}

// Snapshot returns a copy of every known device's last state.
func (h *Hub) Snapshot() []*DeviceState {
	h.mu.RLock()
	defer h.mu.RUnlock()
	out := make([]*DeviceState, 0, len(h.states))
	for _, s := range h.states {
		cp := *s
		out = append(out, &cp)
	}
	return out
}

// Track returns a copy of a device's GPS track.
func (h *Hub) Track(deviceID string) []TrackPoint {
	h.mu.RLock()
	defer h.mu.RUnlock()
	src := h.tracks[deviceID]
	out := make([]TrackPoint, len(src))
	copy(out, src)
	return out
}

// ── Pumps ────────────────────────────────────────────────────────────────────

func (c *Client) readPump() {
	defer func() {
		if c.kind == kindDevice {
			c.hub.removeDevice(c)
		} else {
			c.hub.removeUI(c)
		}
		c.conn.Close()
	}()

	c.conn.SetReadLimit(maxMessageSize)
	_ = c.conn.SetReadDeadline(time.Now().Add(pongWait))
	c.conn.SetPongHandler(func(string) error {
		return c.conn.SetReadDeadline(time.Now().Add(pongWait))
	})

	for {
		_, data, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("ws read error (%s): %v", c.deviceID, err)
			}
			return
		}

		if c.kind == kindDevice {
			var raw map[string]any
			if err := json.Unmarshal(data, &raw); err != nil {
				continue
			}
			c.hub.Ingest(c.deviceID, raw)
			continue
		}

		// UI -> server: command requests
		var req struct {
			Action   string         `json:"action"`
			DeviceID string         `json:"deviceId"`
			Command  string         `json:"command"`
			Payload  map[string]any `json:"payload"`
		}
		if err := json.Unmarshal(data, &req); err != nil {
			continue
		}
		if req.Action == "command" && req.Command != "" {
			c.hub.SendCommand(req.DeviceID, req.Command, req.Payload)
		}
	}
}

func (c *Client) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()

	for {
		select {
		case msg, ok := <-c.send:
			_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				_ = c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}
			if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
				return
			}
		case <-ticker.C:
			_ = c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}