package service import ( "crypto/sha256" "encoding/json" "encoding/hex" "fmt" "io" "os" "path/filepath" "sort" "strings" "3588AdminBackend/internal/models" "3588AdminBackend/internal/storage" ) type ModelManagementService struct { models *storage.ModelsRepo } type InstalledModelStatus struct { Name string `json:"name"` FileName string `json:"file_name"` SHA256 string `json:"sha256"` SizeBytes int64 `json:"size_bytes"` UpdatedAt int64 `json:"updated_at"` } type ModelStatusCell struct { ModelName string `json:"model_name"` FileName string `json:"file_name"` Status string `json:"status"` Version string `json:"version"` } type ModelStatusRow struct { DeviceID string `json:"device_id"` DeviceName string `json:"device_name"` Online bool `json:"online"` Cells []ModelStatusCell `json:"cells"` ExtraModelCount int `json:"extra_model_count"` ExtraModels []InstalledModelStatus `json:"extra_models"` } type ModelStatusSummary struct { StandardModels int `json:"standard_models"` Devices int `json:"devices"` CompleteDevices int `json:"complete_devices"` MissingDevices int `json:"missing_devices"` MismatchDevices int `json:"mismatch_devices"` } type ModelStatusBoard struct { Summary ModelStatusSummary `json:"summary"` Rows []ModelStatusRow `json:"rows"` } func NewModelManagementService(models *storage.ModelsRepo) *ModelManagementService { return &ModelManagementService{models: models} } func (s *ModelManagementService) SyncStandardModelsFromDirectory(dir string) error { if s == nil || s.models == nil { return fmt.Errorf("models repo is not configured") } dir = filepath.Clean(strings.TrimSpace(dir)) if dir == "" { return fmt.Errorf("standard models dir is empty") } if err := os.MkdirAll(dir, 0o755); err != nil { return err } entries, err := os.ReadDir(dir) if err != nil { return err } for _, entry := range entries { if entry.IsDir() || strings.ToLower(filepath.Ext(entry.Name())) != ".rknn" { continue } fullPath := filepath.Join(dir, entry.Name()) sum, size, err := hashFile(fullPath) if err != nil { return err } record := storage.StandardModelRecord{ Name: strings.TrimSuffix(entry.Name(), filepath.Ext(entry.Name())), FileName: entry.Name(), Version: "auto", SHA256: sum, SizeBytes: size, ModelType: inferModelType(entry.Name()), } if err := s.models.Save(record); err != nil { return err } } return nil } func hashFile(path string) (string, int64, error) { file, err := os.Open(path) if err != nil { return "", 0, err } defer file.Close() hasher := sha256.New() size, err := io.Copy(hasher, file) if err != nil { return "", 0, err } return hex.EncodeToString(hasher.Sum(nil)), size, nil } func inferModelType(fileName string) string { name := strings.ToLower(strings.TrimSpace(fileName)) switch { case strings.Contains(name, "face_det"), strings.Contains(name, "retinaface"), strings.Contains(name, "scrfd"): return "face_detection" case strings.Contains(name, "face_recog"), strings.Contains(name, "mobilefacenet"), strings.Contains(name, "arcface"): return "face_recognition" case strings.Contains(name, "ppe"): return "ppe_detection" case strings.Contains(name, "shoe"): return "shoe_detection" case strings.Contains(name, "object_det"), strings.Contains(name, "yolo"): return "object_detection" default: return "other" } } func BuildModelStatusBoard(standardModels []storage.StandardModelRecord, devices []*models.Device, installed map[string][]InstalledModelStatus) ModelStatusBoard { board := ModelStatusBoard{ Summary: ModelStatusSummary{ StandardModels: len(standardModels), Devices: len(devices), }, Rows: make([]ModelStatusRow, 0, len(devices)), } for _, device := range devices { if device == nil { continue } index := make(map[string]InstalledModelStatus, len(installed[device.DeviceID])) for _, item := range installed[device.DeviceID] { index[item.Name] = item } standardIndex := make(map[string]struct{}, len(standardModels)) for _, model := range standardModels { standardIndex[model.Name] = struct{}{} } row := ModelStatusRow{ DeviceID: device.DeviceID, DeviceName: device.DisplayName(), Online: device.Online, Cells: make([]ModelStatusCell, 0, len(standardModels)), ExtraModels: make([]InstalledModelStatus, 0), } hasMissing := false hasMismatch := false for _, model := range standardModels { cell := ModelStatusCell{ ModelName: model.Name, FileName: model.FileName, Version: model.Version, Status: "missing", } if item, ok := index[model.Name]; ok { cell.FileName = item.FileName if strings.EqualFold(strings.TrimSpace(item.SHA256), strings.TrimSpace(model.SHA256)) { cell.Status = "ok" } else { cell.Status = "mismatch" hasMismatch = true } } else { hasMissing = true } if cell.Status == "missing" { hasMissing = true } row.Cells = append(row.Cells, cell) } for _, item := range installed[device.DeviceID] { if _, ok := standardIndex[item.Name]; ok { continue } row.ExtraModels = append(row.ExtraModels, item) } sort.Slice(row.ExtraModels, func(i, j int) bool { return row.ExtraModels[i].Name < row.ExtraModels[j].Name }) row.ExtraModelCount = len(row.ExtraModels) switch { case hasMismatch: board.Summary.MismatchDevices++ case hasMissing: board.Summary.MissingDevices++ default: board.Summary.CompleteDevices++ } board.Rows = append(board.Rows, row) } sort.SliceStable(board.Rows, func(i, j int) bool { return board.Rows[i].DeviceName < board.Rows[j].DeviceName }) return board } func FetchInstalledModelStatuses(agent *AgentClient, device *models.Device) ([]InstalledModelStatus, error) { if agent == nil || device == nil || strings.TrimSpace(device.IP) == "" || device.AgentPort <= 0 { return nil, nil } body, status, err := agent.Do("GET", device.IP, device.AgentPort, "/v1/models/status", nil) if err != nil { return nil, err } if status != 200 { return nil, fmt.Errorf("agent returned status %d", status) } var resp struct { Models []InstalledModelStatus `json:"models"` } if err := json.Unmarshal(body, &resp); err != nil { return nil, err } return resp.Models, nil }