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

package handlers

import (
	"context"
	"fmt"
	"os"
	"runtime"
	"time"

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

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

// udevRuleInstalled checks if the Kneron udev rule is installed on Linux.
func udevRuleInstalled() bool {
	_, err := os.Stat("/etc/udev/rules.d/99-kneron.rules")
	return err == nil
}

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()
	resp := gin.H{
		"devices": devices,
	}
	// Linux: 0 裝置 + udev rule 不存在 → 提示使用者安裝 USB 權限
	if runtime.GOOS == "linux" && len(devices) == 0 && !udevRuleInstalled() {
		resp["udevHint"] = true
	}
	c.JSON(200, gin.H{"success": true, "data": resp})
}

func (h *DeviceHandler) ListDevices(c *gin.Context) {
	devices := h.deviceMgr.ListDevices()
	resp := gin.H{
		"devices": devices,
	}
	if runtime.GOOS == "linux" && len(devices) == 0 && !udevRuleInstalled() {
		resp["udevHint"] = true
	}
	c.JSON(200, gin.H{"success": true, "data": resp})
}

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 now includes mandatory reset + firmware reload on
	// first connect (required for inference to work — see kl720_driver.go
	// needsReset block). Worst-case path on Windows: Loader-mode reconnect
	// retry (16s) + firmware load (~31s) + reboot wait + second reconnect
	// (~13s) = ~60-65s. Use 120s to leave headroom and avoid spurious 504s.
	ctx, cancel := context.WithTimeout(c.Request.Context(), 120*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})
}