visionA/local-tool/server/internal/flash/service.go

package flash

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"visiona-local/server/internal/device"
	"visiona-local/server/internal/driver"
	"visiona-local/server/internal/model"
)

// 可指定的推論種類。與 models.json / 前端同一組值。
//
// models.json 另有 segmentation / pose_estimation，但 Python bridge 與前端都還
// 沒有對應的解析路徑，所以只開放這兩種真的能產出結果的種類。
//
// 燒錄時不再讓使用者選推論種類（改由推論期的
// POST /devices/:id/inference/options 即時切換、不必重燒），但這組常數與
// IsValidTaskTypeOverride 仍是「解析方式」的值域來源，由該 endpoint 沿用，
// 讓 wire 上永遠只有一組合法命名。
const (
	TaskTypeClassification  = "classification"
	TaskTypeObjectDetection = "object_detection"
)

// IsValidTaskTypeOverride 回報 taskType 是否為合法的解析方式覆寫值。
//
// 空字串（未指定）不算合法覆寫 —— 呼叫端要自己先判斷「有沒有要覆寫」，
// 這樣「未指定」與「指定了但打錯字」不會被混為一談。
func IsValidTaskTypeOverride(taskType string) bool {
	return taskType == TaskTypeClassification || taskType == TaskTypeObjectDetection
}

func isCompatible(modelHardware []string, deviceType string) bool {
	dt := strings.ToUpper(deviceType)
	for _, hw := range modelHardware {
		if strings.ToUpper(hw) == dt || strings.Contains(dt, strings.ToUpper(hw)) {
			return true
		}
	}
	return false
}

func resolveModelPath(filePath string, deviceType string) string {
	if filePath == "" {
		return filePath
	}

	targetChip := ""
	if strings.Contains(strings.ToLower(deviceType), "kl720") {
		targetChip = "kl720"
	} else if strings.Contains(strings.ToLower(deviceType), "kl520") {
		targetChip = "kl520"
	}
	if targetChip == "" {
		return filePath
	}

	if strings.Contains(filePath, "/"+targetChip+"/") {
		return filePath
	}

	dir := filepath.Dir(filePath)
	base := filepath.Base(filePath)

	sourceChip := ""
	if strings.Contains(dir, "kl520") {
		sourceChip = "kl520"
	} else if strings.Contains(dir, "kl720") {
		sourceChip = "kl720"
	}

	if sourceChip != "" && sourceChip != targetChip {
		newDir := strings.Replace(dir, sourceChip, targetChip, 1)
		newBase := strings.Replace(base, sourceChip, targetChip, 1)
		candidate := filepath.Join(newDir, newBase)
		if _, err := os.Stat(candidate); err == nil {
			return candidate
		}
	}

	return filePath
}

type Service struct {
	deviceMgr *device.Manager
	modelRepo *model.Repository
	dataDir   string
	tracker   *ProgressTracker
}

func NewService(deviceMgr *device.Manager, modelRepo *model.Repository, dataDir string) *Service {
	return &Service{
		deviceMgr: deviceMgr,
		modelRepo: modelRepo,
		dataDir:   dataDir,
		tracker:   NewProgressTracker(),
	}
}

// CleanupTask 清除已完成的 flash task（由 handler goroutine 在讀取完 progressCh 後呼叫）。
func (s *Service) CleanupTask(taskID string) {
	s.tracker.Remove(taskID)
}

// StartFlash 把 model 載入到裝置。
//
// 推論種類一律用 models.json 宣告的值。使用者若要改解析方式，走推論期的
// POST /devices/:id/inference/options —— 那條路徑不必重燒、可即時切換，
// 功能完全涵蓋燒錄時再選一次的舊做法。
func (s *Service) StartFlash(deviceID, modelID string) (string, <-chan driver.FlashProgress, error) {
	session, err := s.deviceMgr.GetDevice(deviceID)
	if err != nil {
		return "", nil, fmt.Errorf("device not found: %w", err)
	}
	if !session.Driver.IsConnected() {
		return "", nil, fmt.Errorf("device not connected")
	}

	m, err := s.modelRepo.GetByID(modelID)
	if err != nil {
		return "", nil, fmt.Errorf("model not found: %w", err)
	}

	deviceInfo := session.Driver.Info()
	if !isCompatible(m.SupportedHardware, deviceInfo.Type) {
		return "", nil, fmt.Errorf("model not compatible with device type %s", deviceInfo.Type)
	}

	modelPath := m.FilePath
	if modelPath == "" {
		return "", nil, fmt.Errorf("model %s has no .nef file path", modelID)
	}

	// models.json 的 filePath 是相對路徑（例如 "data/nef/kl520/xxx.nef"）。
	// 如果不是絕對路徑，用 dataDir 解析：
	//   "data/nef/..." → 去掉 "data/" 前綴 → dataDir + "/nef/..."
	//   其他相對路徑 → dataDir + "/" + filePath
	if !filepath.IsAbs(modelPath) {
		if strings.HasPrefix(modelPath, "data/") || strings.HasPrefix(modelPath, "data\\") {
			modelPath = filepath.Join(s.dataDir, modelPath[len("data/"):])
		} else {
			modelPath = filepath.Join(s.dataDir, modelPath)
		}
	}

	modelPath = resolveModelPath(modelPath, deviceInfo.Type)

	taskID := fmt.Sprintf("flash-%s-%s", deviceID, modelID)

	// M3 fix: 防止同裝置同模型重複 flash
	task := s.tracker.Create(taskID, deviceID, modelID)
	if task == nil {
		return "", nil, fmt.Errorf("flash already in progress for device %s model %s", deviceID, modelID)
	}

	go func() {
		// M1 fix: 先跑 driver.Flash，收集 error，最後才寫 error message + close channel。
		// driver.Flash 內部會多次寫入 task.ProgressCh（進度更新），我們不能在它還在寫的時候 close。
		// driver.Flash 返回時保證不會再寫入 progressCh。

		time.Sleep(500 * time.Millisecond)

		// 把 models.json 宣告的 metadata 一起帶下去 —— bridge 端有 taskType
		// 就不再靠檔名猜 model type（自訂模型存成 model.nef、檔名沒有關鍵字，
		// 猜測必定落到 detection 分支）。labels 純顯示層、沒有也能跑。
		//
		// inputSize 是宣告值、**優先序最低**：bridge 端會先問 SDK 模型自己
		// 宣告的 input shape，只有問不到才用這裡的值（這欄是人填的，可能亂填）。
		flashErr := session.Driver.Flash(modelPath, driver.FlashOptions{
			TaskType:    m.TaskType,
			Labels:      m.Labels,
			InputWidth:  m.InputSize.Width,
			InputHeight: m.InputSize.Height,
		}, task.ProgressCh)

		// Flash 完成或失敗後，driver 不會再寫 progressCh，安全地寫 error 訊息然後 close。
		if flashErr != nil {
			task.ProgressCh <- driver.FlashProgress{
				Percent: -1,
				Stage:   "error",
				Error:   flashErr.Error(),
			}
		}

		task.Done = true
		close(task.ProgressCh)
		// M2 note: 不在這裡 Remove — 讓 handler 讀完 progressCh 後呼叫 CleanupTask
	}()

	return taskID, task.ProgressCh, nil
}