visionA/local-agent/server/internal/api/ws/hub.go

package ws

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/gorilla/websocket"
)

// replayPrefix 決定哪些 room 啟用「late-join replay 緩存」。
//
// 只對推論結果 room（"inference:<serial>"）緩存最近 N 筆結果，理由：
//   - 推論結果面走雲端 tunnel（ADR-019 混合路徑）、比上傳的 localhost 慢，
//     上傳完成到 WS join room 之間有時間窗口（見 video-inference-stuck-rootcause.md §2）。
//   - flash / system / device-events / server-logs 等 room 沒有這種「早期訊息會被丟」的問題，
//     不緩存以免無謂佔記憶體。
const replayPrefix = "inference:"

// replayBufferSize 是每個 inference room 緩存的最近結果筆數上限。
//
// 30 筆 ≈ 15fps 影片的 2 秒；足夠覆蓋「上傳完成 → tunnel WS join room」的窗口，
// 又不會讓記憶體膨脹（單筆 InferenceResult JSON 通常 < 幾 KB）。緩存在
// ClearRoomReplay（pipeline 停止 / 切換時）被清掉，跟 stopActivePipeline 生命週期對齊。
const replayBufferSize = 30

type Client struct {
	Conn *websocket.Conn
	Send chan []byte
}

type Subscription struct {
	Client *Client
	Room   string
	done   chan struct{} // used by RegisterSync to wait for completion
}

type RoomMessage struct {
	Room    string
	Message []byte
}

// roomWaiter 由 WaitForRoomClient 註冊、在 Run() loop 內處理，
// 確保「檢查 room 是否有 client」與「register 事件」在同一 goroutine 序列化、無 race。
//
// 若註冊當下 room 已有 client → 立即 close(ch)；否則存進 h.waiters，
// 待該 room 有 client register 時 close(ch) 喚醒。
type roomWaiter struct {
	room string
	ch   chan struct{}
}

// Hub 管理 WebSocket client 訂閱與訊息廣播。
//
// M8-4b：Hub 額外負責「第一個 client 連上時寫 sentinel file」，
// 讓 Wails 端的 StartupPipeline 知道階段 6（Wait for Web UI WebSocket）已完成。
// 詳細設計見 .autoflow/04-architecture/v2/startup-pipeline.md §3。
//
// dataDir 由 main.go 在初始化 Hub 後透過 SetStartupSentinel(dataDir) 注入。
// 若 dataDir 為空，sentinel 寫入會被跳過（單元測試或缺少資料目錄時的安全行為）。
type Hub struct {
	rooms       map[string]map[*Client]bool
	register    chan *Subscription
	unregister  chan *Subscription
	broadcast   chan *RoomMessage
	waitReq     chan *roomWaiter // WaitForRoomClient 的等待請求（在 Run() 內序列化處理）
	clearReplay chan string      // ClearRoomReplay 的清除請求
	mu          sync.RWMutex

	// M8-4b: 啟動 sentinel file
	sentinelDataDir string    // <dataDir>，由 SetStartupSentinel 設定
	sentinelOnce    sync.Once // 確保只在「第一個」client 連上時寫一次
	bootID          string    // 寫入 sentinel 內容供 debug

	// video-inference-stuck 修法 A2/B（均在 Run() goroutine 內存取、無需額外鎖）：
	//   waiters — WaitForRoomClient 尚未被喚醒的等待者，key=room
	//   replay  — 每個 inference room 的最近 replayBufferSize 筆訊息 ring（late-join replay）
	waiters map[string][]*roomWaiter
	replay  map[string][][]byte
}

func NewHub() *Hub {
	return &Hub{
		rooms:       make(map[string]map[*Client]bool),
		register:    make(chan *Subscription, 10),
		unregister:  make(chan *Subscription, 10),
		broadcast:   make(chan *RoomMessage, 100),
		waitReq:     make(chan *roomWaiter, 10),
		clearReplay: make(chan string, 10),
		bootID:      fmt.Sprintf("boot-%d", time.Now().UnixNano()),
		waiters:     make(map[string][]*roomWaiter),
		replay:      make(map[string][][]byte),
	}
}

// SetStartupSentinel 設定 sentinel file 的根目錄。
// main.go 在 NewHub() 之後、Run() 之前呼叫一次，dataDir 應為完整路徑。
//
// 寫入路徑：<dataDir>/.first-ws-connected
// 內容：boot-id + timestamp（用於 debug，內容對 Wails 端的判斷沒有意義，存在即可）
//
// dataDir 為空字串時 sentinel 機制完全停用。
func (h *Hub) SetStartupSentinel(dataDir string) {
	h.mu.Lock()
	h.sentinelDataDir = dataDir
	h.mu.Unlock()
}

// writeStartupSentinel 在第一個 WebSocket client 連上時呼叫一次。
// 由 sentinelOnce 確保只執行一次；後續連線完全 no-op。
//
// 寫入失敗不會 panic 也不會回 error：sentinel 是 best-effort 機制，
// 若 disk 滿/權限錯，Wails 端會走 hard timeout 路徑進 Error state。
func (h *Hub) writeStartupSentinel() {
	h.sentinelOnce.Do(func() {
		h.mu.RLock()
		dir := h.sentinelDataDir
		bootID := h.bootID
		h.mu.RUnlock()
		if dir == "" {
			return
		}
		path := filepath.Join(dir, ".first-ws-connected")
		// 確保父目錄存在（dataDir 通常已存在，但保險起見）
		_ = os.MkdirAll(dir, 0o755)
		f, err := os.Create(path)
		if err != nil {
			return
		}
		_, _ = fmt.Fprintf(f, "bootId=%s\nts=%d\n", bootID, time.Now().UnixMilli())
		_ = f.Close()
	})
}

func (h *Hub) Run() {
	for {
		select {
		case sub := <-h.register:
			h.mu.Lock()
			if h.rooms[sub.Room] == nil {
				h.rooms[sub.Room] = make(map[*Client]bool)
			}
			h.rooms[sub.Room][sub.Client] = true
			// B（late-join replay）：inference room 若有緩存的早期結果，
			// 在此把它們補送給剛 join 的 client，解「WS 稍慢也不丟第一筆」。
			// 在鎖內取出 replay 快照、鎖外送出，避免 client.Send 阻塞時卡住 Run() 。
			var pending [][]byte
			if buffered, ok := h.replay[sub.Room]; ok && len(buffered) > 0 {
				pending = make([][]byte, len(buffered))
				copy(pending, buffered)
			}
			h.mu.Unlock()
			for _, msg := range pending {
				select {
				case sub.Client.Send <- msg:
				default:
					// client buffer 已滿（極罕見：replay 30 筆 > send buffer 20）→ 停止補送，
					// 後續 live 廣播仍會照常送達，不因 replay 溢出而 drop client。
				}
			}
			// A2：喚醒等待「此 room 有 client」的 waiter（WaitForRoomClient）。
			if ws := h.waiters[sub.Room]; len(ws) > 0 {
				for _, w := range ws {
					close(w.ch)
				}
				delete(h.waiters, sub.Room)
			}
			// M8-4b：第一次有 client 加入任何 room → 寫 sentinel file
			// （sync.Once 保證後續呼叫 no-op）
			h.writeStartupSentinel()
			if sub.done != nil {
				close(sub.done)
			}

		case w := <-h.waitReq:
			// A2：WaitForRoomClient 的請求。若 room 當下已有 client → 立即喚醒；
			// 否則存進 waiters，待 register 時喚醒。與 register 在同一 goroutine
			// 序列化處理，故「檢查 + 掛等待」對 register 事件無 race。
			h.mu.RLock()
			hasClient := len(h.rooms[w.room]) > 0
			h.mu.RUnlock()
			if hasClient {
				close(w.ch)
			} else {
				h.waiters[w.room] = append(h.waiters[w.room], w)
			}

		case room := <-h.clearReplay:
			// pipeline 停止 / 切換時清掉該 room 的 replay 緩存（跟 stopActivePipeline 對齊），
			// 防止上一支影片的結果殘留給下一次 join 的 client。
			h.mu.Lock()
			delete(h.replay, room)
			h.mu.Unlock()

		case sub := <-h.unregister:
			h.mu.Lock()
			if clients, ok := h.rooms[sub.Room]; ok {
				if _, exists := clients[sub.Client]; exists {
					delete(clients, sub.Client)
					close(sub.Client.Send)
				}
			}
			h.mu.Unlock()

		case msg := <-h.broadcast:
			h.mu.Lock()
			// B：inference room 的訊息先進 replay ring（不論當下有無 client）。
			if strings.HasPrefix(msg.Room, replayPrefix) {
				buf := append(h.replay[msg.Room], msg.Message)
				if len(buf) > replayBufferSize {
					buf = buf[len(buf)-replayBufferSize:]
				}
				h.replay[msg.Room] = buf
			}
			if clients, ok := h.rooms[msg.Room]; ok {
				for client := range clients {
					select {
					case client.Send <- msg.Message:
					default:
						close(client.Send)
						delete(clients, client)
					}
				}
			}
			h.mu.Unlock()
		}
	}
}

func (h *Hub) Register(sub *Subscription) {
	h.register <- sub
}

// RegisterSync registers a subscription and blocks until the Hub has processed it,
// ensuring the client is in the room before returning.
func (h *Hub) RegisterSync(sub *Subscription) {
	sub.done = make(chan struct{})
	h.register <- sub
	<-sub.done
}

func (h *Hub) Unregister(sub *Subscription) {
	h.unregister <- sub
}

func (h *Hub) BroadcastToRoom(room string, data interface{}) {
	jsonData, err := json.Marshal(data)
	if err != nil {
		return
	}
	h.broadcast <- &RoomMessage{Room: room, Message: jsonData}
}

// HasClients 回報指定 room 當下是否有至少一個 client。
// 用於快速判斷、不阻塞（WaitForRoomClient 才是「等到有 client」的阻塞版）。
func (h *Hub) HasClients(room string) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.rooms[room]) > 0
}

// WaitForRoomClient 阻塞直到 room 有至少一個 client join、或 ctx 被取消。
//
// 回傳 true 代表「room 已有 client」；false 代表 ctx 先結束（timeout / 上游取消）。
//
// A2 主修用途：UploadVideo 存檔後即回 200，但 pipeline 廣播延到「結果 WS 已 join
// inference room」才開跑——避免上傳（localhost，快）與結果訂閱（tunnel WS，慢）
// 時序解耦後，早期結果在 Hub 因 room 無 client 被靜默丟棄（root cause §2）。
//
// 實作：把等待請求丟進 Run() goroutine 序列化處理（waitReq case），
// 確保「檢查 room 是否有 client」與「register 事件」無 race。
// ctx 先結束時仍會殘留一個 waiter 在 h.waiters，但 register 喚醒它只是 close 一個
// 沒人收的 channel（無害）；replay/room 生命週期短，不會累積。
func (h *Hub) WaitForRoomClient(ctx context.Context, room string) bool {
	w := &roomWaiter{room: room, ch: make(chan struct{})}
	select {
	case h.waitReq <- w:
	case <-ctx.Done():
		return false
	}
	select {
	case <-w.ch:
		return true
	case <-ctx.Done():
		return false
	}
}

// ClearRoomReplay 清除指定 room 的 late-join replay 緩存。
// 由 CameraHandler 在切換 / 停止 pipeline 時呼叫，跟 stopActivePipeline 生命週期對齊，
// 防止上一支影片的結果殘留給下一次 join 的 client。非 inference room 呼叫也安全（no-op）。
func (h *Hub) ClearRoomReplay(room string) {
	h.clearReplay <- room
}