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

package handlers

import (
	"encoding/json"
	"fmt"
	"path/filepath"
	"strings"
	"time"

	"visiona-agent/server/internal/model"

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

type ModelUploadHandler struct {
	repo  *model.Repository
	store *model.ModelStore
}

func NewModelUploadHandler(repo *model.Repository, store *model.ModelStore) *ModelUploadHandler {
	return &ModelUploadHandler{repo: repo, store: store}
}

func (h *ModelUploadHandler) UploadModel(c *gin.Context) {
	// Get uploaded file
	file, header, err := c.Request.FormFile("file")
	if err != nil {
		c.JSON(400, gin.H{
			"success": false,
			"error":   gin.H{"code": "BAD_REQUEST", "message": "file is required"},
		})
		return
	}
	defer file.Close()

	// Validate file extension
	ext := strings.ToLower(filepath.Ext(header.Filename))
	if ext != ".nef" {
		c.JSON(400, gin.H{
			"success": false,
			"error":   gin.H{"code": "BAD_REQUEST", "message": "only .nef files are supported"},
		})
		return
	}

	// Get required fields
	name := c.PostForm("name")
	taskType := c.PostForm("taskType")
	labelsStr := c.PostForm("labels")

	if name == "" || taskType == "" || labelsStr == "" {
		c.JSON(400, gin.H{
			"success": false,
			"error":   gin.H{"code": "BAD_REQUEST", "message": "name, taskType, and labels are required"},
		})
		return
	}

	// Parse labels
	var labels []string
	if err := json.Unmarshal([]byte(labelsStr), &labels); err != nil {
		c.JSON(400, gin.H{
			"success": false,
			"error":   gin.H{"code": "BAD_REQUEST", "message": "labels must be a JSON array"},
		})
		return
	}

	// Optional fields
	description := c.PostForm("description")
	quantization := c.PostForm("quantization")
	if quantization == "" {
		quantization = "INT8"
	}
	inputWidth := 640
	inputHeight := 640
	if w := c.PostForm("inputWidth"); w != "" {
		fmt.Sscanf(w, "%d", &inputWidth)
	}
	if h := c.PostForm("inputHeight"); h != "" {
		fmt.Sscanf(h, "%d", &inputHeight)
	}

	// Generate ID
	id := uuid.New().String()
	now := time.Now().UTC().Format(time.RFC3339)

	// Save file
	nefPath, err := h.store.SaveModel(id, file)
	if err != nil {
		c.JSON(500, gin.H{
			"success": false,
			"error":   gin.H{"code": "STORAGE_ERROR", "message": err.Error()},
		})
		return
	}

	// Build model
	m := model.Model{
		ID:                id,
		Name:              name,
		Description:       description,
		TaskType:          taskType,
		Categories:        []string{"custom"},
		Framework:         "NEF",
		InputSize:         model.InputSize{Width: inputWidth, Height: inputHeight},
		ModelSize:         header.Size,
		Quantization:      quantization,
		SupportedHardware: []string{"KL520", "KL720"},
		Labels:            labels,
		Version:           "1.0.0",
		Author:            "Custom",
		License:           "Custom",
		CreatedAt:         now,
		UpdatedAt:         now,
		IsCustom:          true,
		FilePath:          nefPath,
	}

	// Save metadata
	if err := h.store.SaveMetadata(id, m); err != nil {
		c.JSON(500, gin.H{
			"success": false,
			"error":   gin.H{"code": "STORAGE_ERROR", "message": err.Error()},
		})
		return
	}

	// Add to repository
	h.repo.Add(m)

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

func (h *ModelUploadHandler) DeleteModel(c *gin.Context) {
	id := c.Param("id")

	// Remove from repository (validates it's custom)
	if err := h.repo.Remove(id); err != nil {
		c.JSON(400, gin.H{
			"success": false,
			"error":   gin.H{"code": "DELETE_FAILED", "message": err.Error()},
		})
		return
	}

	// Remove files
	if err := h.store.DeleteModel(id); err != nil {
		fmt.Printf("[WARN] Failed to delete model files for %s: %v\n", id, err)
	}

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