ai/routers/datasets.go

package routers

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"main/configs"
	"main/models"
	"main/utils"
	"net/http"
	"os"

	"github.com/gorilla/mux"
)

// 獲取數據集列表
func DatasetsGet(w http.ResponseWriter, r *http.Request) {
	var listview models.ListView
	listview.Page = utils.ParamInt(r.URL.Query().Get("page"), 1)
	listview.PageSize = utils.ParamInt(r.URL.Query().Get("pageSize"), 10)
	var dataset_list []models.Dataset
	db := configs.ORMDB()
	db.Offset((listview.Page - 1) * listview.PageSize).Limit(listview.PageSize).Find(&dataset_list)
	for _, dataset := range dataset_list {
		if dataset.Images == nil {
			dataset.Images = models.ImageList{}
		}
		listview.List = append(listview.List, dataset)
	}
	db.Model(&models.Dataset{}).Count(&listview.Total)
	listview.Next = listview.Page*listview.PageSize < int(listview.Total)
	listview.WriteJSON(w)
}

// 創建數據集
func DatasetsPost(w http.ResponseWriter, r *http.Request) {
	models.AccountRead(w, r, func(account *models.Account) {
		var dataset models.Dataset
		body, err := ioutil.ReadAll(r.Body)
		if err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte("500 - Internal Server Error"))
			return
		}
		defer r.Body.Close()
		if err = json.Unmarshal(body, &dataset); err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte("500 - Internal Server Error"))
			return
		}
		if dataset.Images == nil {
			dataset.Images = models.ImageList{}
		}
		dataset.UserID = account.ID
		if err := configs.ORMDB().Create(&dataset).Error; err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte("500 - Internal Server Error"))
			return
		}
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		w.Write(utils.ToJSON(dataset))
	})
}

// 上传图片文件
func DatasetsUpload(w http.ResponseWriter, r *http.Request) {
	models.AccountRead(w, r, func(account *models.Account) {
		// 獲取數據集
		dataset := models.Dataset{ID: utils.ParamInt(mux.Vars(r)["dataset_id"], 0)}
		if err := configs.ORMDB().Find(&dataset).Error; err != nil {
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte("404 - Not Found"))
			return
		}
		// 只能修改自己的數據集, 除非是管理員
		if dataset.UserID != account.ID && !account.Admin {
			w.WriteHeader(http.StatusForbidden)
			w.Write([]byte("403 - Forbidden"))
			return
		}
		// 解析 HTTP 请求中的多个文件 (限制上传文件的大小为 32MB)
		if err := r.ParseMultipartForm(32 << 20); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		// 遍历所有上传的文件
		for _, fileHeaders := range r.MultipartForm.File {
			for _, fileHeader := range fileHeaders {
				// 打开上传的文件
				file, err := fileHeader.Open()
				if err != nil {
					http.Error(w, err.Error(), http.StatusInternalServerError)
					return
				}
				defer file.Close()

				// 在本地文件夹中创建一个新文件
				localFile, err := os.Create(fmt.Sprintf("data/dataset/%d/%s", dataset.ID, fileHeader.Filename))
				if err != nil {
					http.Error(w, err.Error(), http.StatusInternalServerError)
					return
				}
				defer localFile.Close()

				// 将上传文件的内容复制到本地文件
				_, err = io.Copy(localFile, file)
				if err != nil {
					http.Error(w, err.Error(), http.StatusInternalServerError)
					return
				}

				// 将文件名添加到數據集中
				dataset.Images = append(dataset.Images, fileHeader.Filename)
			}
		}

		// 去除重复项
		uniqueImages := make(map[string]bool)
		for _, image := range dataset.Images {
			uniqueImages[image] = true
		}

		// 转换为切片
		dataset.Images = []string{}
		for image := range uniqueImages {
			dataset.Images = append(dataset.Images, image)
		}

		// 保存數據集
		if err := configs.ORMDB().Save(&dataset).Error; err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte("500 - Internal Server Error"))
			return
		}
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		w.Write(utils.ToJSON(dataset))
	})
}

// 獲取數據集
func DatasetsItemGet(w http.ResponseWriter, r *http.Request) {
	dataset := models.Dataset{ID: utils.ParamInt(mux.Vars(r)["dataset_id"], 0)}
	if err := configs.ORMDB().Find(&dataset).Error; err != nil {
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte("404 - Not Found"))
		return
	}
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.Write(utils.ToJSON(dataset))
}

// 修改數據集
func DatasetsItemPatch(w http.ResponseWriter, r *http.Request) {
	models.AccountRead(w, r, func(account *models.Account) {
		var dataset models.Dataset = models.Dataset{ID: utils.ParamInt(mux.Vars(r)["dataset_id"], 0)}
		if err := configs.ORMDB().Find(&dataset).Error; err != nil {
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte("404 - Not Found"))
			return
		}
		if dataset.UserID != account.ID && !account.Admin {
			w.WriteHeader(http.StatusForbidden)
			w.Write([]byte("403 - Forbidden"))
			return
		}
		var form map[string]interface{} = utils.BodyRead(r)
		if name, ok := form["name"].(string); ok {
			dataset.Name = name
		}
		if info, ok := form["info"].(string); ok {
			dataset.Info = info
		}
		if images, ok := form["images"].([]interface{}); ok {
			var image_list models.ImageList
			for _, image := range images {
				if image, ok := image.(string); ok {
					image_list = append(image_list, image)
				}
			}
			dataset.Images = image_list
		}
		if err := configs.ORMDB().Save(&dataset).Error; err != nil {
			w.WriteHeader(http.StatusInternalServerError)
			w.Write([]byte("500 - Internal Server Error"))
			return
		}
		w.Header().Set("Content-Type", "application/json; charset=utf-8")
		w.Write(utils.ToJSON(dataset))
	})
}

// 刪除數據集
func DatasetsItemDelete(w http.ResponseWriter, r *http.Request) {
	models.AccountRead(w, r, func(account *models.Account) {
		// 獲取數據集
		dataset := models.Dataset{ID: utils.ParamInt(mux.Vars(r)["id"], 0)}
		if err := configs.ORMDB().Find(&dataset).Error; err != nil {
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte("404 - Not Found"))
			return
		}
		// 只能刪除自己的數據集, 除非是管理員
		if dataset.UserID != account.ID && !account.Admin {
			w.WriteHeader(http.StatusForbidden)
			w.Write([]byte("403 - Forbidden"))
			return
		}
		// 刪除數據集
		if err := configs.ORMDB().Delete(&dataset).Error; err != nil {
			w.WriteHeader(http.StatusNotFound)
			w.Write([]byte("404 - Not Found"))
			return
		}
		w.WriteHeader(http.StatusNoContent)
	})
}