3588AdminBackend/internal/service/model_management.go

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"`
}

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,
		}
		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 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
		}
		row := ModelStatusRow{
			DeviceID:   device.DeviceID,
			DeviceName: device.DisplayName(),
			Online:     device.Online,
			Cells:      make([]ModelStatusCell, 0, len(standardModels)),
		}
		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)
		}
		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
}