visionA/local-tool/server/internal/api/handlers/device_handler.go

package handlers

import (
	"context"
	"fmt"
	"time"

	"visiona-local/server/internal/api/ws"
	"visiona-local/server/internal/device"
	"visiona-local/server/internal/driver"
	"visiona-local/server/internal/flash"
	"visiona-local/server/internal/inference"

	"github.com/gin-gonic/gin"
)

type DeviceHandler struct {
	deviceMgr    *device.Manager
	flashSvc     *flash.Service
	inferenceSvc *inference.Service
	wsHub        *ws.Hub
}

func NewDeviceHandler(
	deviceMgr *device.Manager,
	flashSvc *flash.Service,
	inferenceSvc *inference.Service,
	wsHub *ws.Hub,
) *DeviceHandler {
	return &DeviceHandler{
		deviceMgr:    deviceMgr,
		flashSvc:     flashSvc,
		inferenceSvc: inferenceSvc,
		wsHub:        wsHub,
	}
}

func (h *DeviceHandler) ScanDevices(c *gin.Context) {
	devices := h.deviceMgr.Rescan()
	c.JSON(200, gin.H{
		"success": true,
		"data": gin.H{
			"devices": devices,
		},
	})
}

func (h *DeviceHandler) ListDevices(c *gin.Context) {
	devices := h.deviceMgr.ListDevices()
	c.JSON(200, gin.H{
		"success": true,
		"data": gin.H{
			"devices": devices,
		},
	})
}

func (h *DeviceHandler) GetDevice(c *gin.Context) {
	id := c.Param("id")
	session, err := h.deviceMgr.GetDevice(id)
	if err != nil {
		c.JSON(404, gin.H{
			"success": false,
			"error":   gin.H{"code": "DEVICE_NOT_FOUND", "message": err.Error()},
		})
		return
	}
	c.JSON(200, gin.H{"success": true, "data": session.Driver.Info()})
}

func (h *DeviceHandler) ConnectDevice(c *gin.Context) {
	id := c.Param("id")

	// KL520 USB Boot flow can take ~40s: retry connect (3x2s) + firmware
	// load + 5s reboot wait + reconnect retry (3x3s). Use 60s timeout.
	ctx, cancel := context.WithTimeout(c.Request.Context(), 60*time.Second)
	defer cancel()

	errCh := make(chan error, 1)
	go func() {
		errCh <- h.deviceMgr.Connect(id)
	}()

	select {
	case err := <-errCh:
		if err != nil {
			c.JSON(400, gin.H{
				"success": false,
				"error":   gin.H{"code": "CONNECT_FAILED", "message": err.Error()},
			})
			return
		}
		c.JSON(200, gin.H{"success": true})
	case <-ctx.Done():
		c.JSON(504, gin.H{
			"success": false,
			"error":   gin.H{"code": "CONNECT_TIMEOUT", "message": fmt.Sprintf("device connect timed out after 60s for %s", id)},
		})
	}
}

func (h *DeviceHandler) DisconnectDevice(c *gin.Context) {
	id := c.Param("id")
	if err := h.deviceMgr.Disconnect(id); err != nil {
		c.JSON(400, gin.H{
			"success": false,
			"error":   gin.H{"code": "DISCONNECT_FAILED", "message": err.Error()},
		})
		return
	}
	c.JSON(200, gin.H{"success": true})
}

func (h *DeviceHandler) FlashDevice(c *gin.Context) {
	id := c.Param("id")
	var req struct {
		ModelID string `json:"modelId"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(400, gin.H{
			"success": false,
			"error":   gin.H{"code": "BAD_REQUEST", "message": "modelId is required"},
		})
		return
	}

	taskID, progressCh, err := h.flashSvc.StartFlash(id, req.ModelID)
	if err != nil {
		c.JSON(400, gin.H{
			"success": false,
			"error":   gin.H{"code": "FLASH_FAILED", "message": err.Error()},
		})
		return
	}

	// Forward progress to WebSocket, then cleanup task (M2 fix)
	go func() {
		room := "flash:" + id
		for progress := range progressCh {
			h.wsHub.BroadcastToRoom(room, progress)
		}
		h.flashSvc.CleanupTask(taskID)
	}()

	c.JSON(200, gin.H{"success": true, "data": gin.H{"taskId": taskID}})
}

func (h *DeviceHandler) StartInference(c *gin.Context) {
	id := c.Param("id")
	resultCh := make(chan *driver.InferenceResult, 10)

	if err := h.inferenceSvc.Start(id, resultCh); err != nil {
		c.JSON(400, gin.H{
			"success": false,
			"error":   gin.H{"code": "INFERENCE_ERROR", "message": err.Error()},
		})
		return
	}

	// Forward results to WebSocket, enriching with device ID
	go func() {
		room := "inference:" + id
		for result := range resultCh {
			result.DeviceID = id
			h.wsHub.BroadcastToRoom(room, result)
		}
	}()

	c.JSON(200, gin.H{"success": true})
}

func (h *DeviceHandler) StopInference(c *gin.Context) {
	id := c.Param("id")
	if err := h.inferenceSvc.Stop(id); err != nil {
		c.JSON(400, gin.H{
			"success": false,
			"error":   gin.H{"code": "INFERENCE_ERROR", "message": err.Error()},
		})
		return
	}
	c.JSON(200, gin.H{"success": true})
}