aiforge/routers/repo/ai_model_manage.go

package repo

import (
	"archive/zip"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"path"
	"regexp"
	"strings"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/notification"
	"code.gitea.io/gitea/modules/setting"
	"code.gitea.io/gitea/modules/storage"
	"code.gitea.io/gitea/services/cloudbrain/resource"
	uuid "github.com/satori/go.uuid"
)

const (
	Attachment_model                 = "model"
	Model_prefix                     = "aimodels/"
	tplModelManageIndex              = "repo/modelmanage/index"
	tplModelManageDownload           = "repo/modelmanage/download"
	tplModelInfo                     = "repo/modelmanage/showinfo"
	tplCreateLocalModelInfo          = "repo/modelmanage/create_local_1"
	tplCreateLocalForUploadModelInfo = "repo/modelmanage/create_local_2"
	tplCreateOnlineModelInfo         = "repo/modelmanage/create_online"

	MODEL_LATEST      = 1
	MODEL_NOT_LATEST  = 0
	MODEL_MAX_SIZE    = 1024 * 1024 * 1024
	STATUS_COPY_MODEL = 1
	STATUS_FINISHED   = 0
	STATUS_ERROR      = 2

	MODEL_LOCAL_TYPE  = 1
	MODEL_ONLINE_TYPE = 0
)

func saveModelByParameters(jobId string, versionName string, name string, version string, label string, description string, engine int, ctx *context.Context) (string, error) {
	aiTask, err := models.GetCloudbrainByJobIDAndVersionName(jobId, versionName)
	if err != nil {
		aiTask, err = models.GetRepoCloudBrainByJobID(ctx.Repo.Repository.ID, jobId)
		if err != nil {
			log.Info("query task error." + err.Error())
			return "", err
		} else {
			log.Info("query gpu train task.")
		}
	}
	uuid := uuid.NewV4()
	id := uuid.String()
	modelPath := id
	var lastNewModelId string
	var modelSize int64

	log.Info("find task name:" + aiTask.JobName)
	aimodels := models.QueryModelByName(name, aiTask.RepoID)
	if len(aimodels) > 0 {
		for _, model := range aimodels {
			if model.Version == version {
				return "", errors.New(ctx.Tr("repo.model.manage.create_error"))
			}
			if model.New == MODEL_LATEST {
				lastNewModelId = model.ID
			}
		}
	}
	cloudType := aiTask.Type
	modelSelectedFile := ctx.Query("modelSelectedFile")
	//download model zip   //train type
	if aiTask.ComputeResource == models.NPUResource {
		cloudType = models.TypeCloudBrainTwo
	} else if aiTask.ComputeResource == models.GPUResource {
		cloudType = models.TypeCloudBrainOne
		spec, err := resource.GetCloudbrainSpec(aiTask.ID)
		if err == nil {
			flaverName := "GPU: " + fmt.Sprint(spec.AccCardsNum) + "*" + spec.AccCardType + ",CPU: " + fmt.Sprint(spec.CpuCores) + "," + ctx.Tr("cloudbrain.memory") + ": " + fmt.Sprint(spec.MemGiB) + "GB," + ctx.Tr("cloudbrain.shared_memory") + ": " + fmt.Sprint(spec.ShareMemGiB) + "GB"
			aiTask.FlavorName = flaverName
		}
	}

	accuracy := make(map[string]string)
	accuracy["F1"] = ""
	accuracy["Recall"] = ""
	accuracy["Accuracy"] = ""
	accuracy["Precision"] = ""
	accuracyJson, _ := json.Marshal(accuracy)
	log.Info("accuracyJson=" + string(accuracyJson))
	aiTask.ContainerIp = ""
	aiTaskJson, _ := json.Marshal(aiTask)

	model := &models.AiModelManage{
		ID:            id,
		Version:       version,
		VersionCount:  len(aimodels) + 1,
		Label:         label,
		Name:          name,
		Description:   description,
		New:           MODEL_LATEST,
		Type:          cloudType,
		Path:          modelPath,
		Size:          modelSize,
		AttachmentId:  aiTask.Uuid,
		RepoId:        aiTask.RepoID,
		UserId:        ctx.User.ID,
		CodeBranch:    aiTask.BranchName,
		CodeCommitID:  aiTask.CommitID,
		Engine:        int64(engine),
		TrainTaskInfo: string(aiTaskJson),
		Accuracy:      string(accuracyJson),
		Status:        STATUS_COPY_MODEL,
	}

	err = models.SaveModelToDb(model)
	if err != nil {
		return "", err
	}
	if len(lastNewModelId) > 0 {
		//udpate status and version count
		models.ModifyModelNewProperty(lastNewModelId, MODEL_NOT_LATEST, 0)
	}
	var units []models.RepoUnit
	var deleteUnitTypes []models.UnitType
	units = append(units, models.RepoUnit{
		RepoID: ctx.Repo.Repository.ID,
		Type:   models.UnitTypeModelManage,
		Config: &models.ModelManageConfig{
			EnableModelManage: true,
		},
	})
	deleteUnitTypes = append(deleteUnitTypes, models.UnitTypeModelManage)

	models.UpdateRepositoryUnits(ctx.Repo.Repository, units, deleteUnitTypes)

	go asyncToCopyModel(aiTask, id, modelSelectedFile)

	log.Info("save model end.")
	notification.NotifyOtherTask(ctx.User, ctx.Repo.Repository, id, name, models.ActionCreateNewModelTask)
	return id, nil
}

func asyncToCopyModel(aiTask *models.Cloudbrain, id string, modelSelectedFile string) {
	if aiTask.ComputeResource == models.NPUResource {
		modelPath, modelSize, err := downloadModelFromCloudBrainTwo(id, aiTask.JobName, "", aiTask.TrainUrl, modelSelectedFile)
		if err != nil {
			updateStatus(id, 0, STATUS_ERROR, modelPath, err.Error())
			log.Info("download model from CloudBrainTwo faild." + err.Error())
		} else {
			updateStatus(id, modelSize, STATUS_FINISHED, modelPath, "")
		}
	} else if aiTask.ComputeResource == models.GPUResource {

		modelPath, modelSize, err := downloadModelFromCloudBrainOne(id, aiTask.JobName, "", aiTask.TrainUrl, modelSelectedFile)
		if err != nil {
			updateStatus(id, 0, STATUS_ERROR, modelPath, err.Error())
			log.Info("download model from CloudBrainOne faild." + err.Error())
		} else {
			updateStatus(id, modelSize, STATUS_FINISHED, modelPath, "")
		}
	}
}

func updateStatus(id string, modelSize int64, status int, modelPath string, statusDesc string) {
	if len(statusDesc) > 400 {
		statusDesc = statusDesc[0:400]
	}
	err := models.ModifyModelStatus(id, modelSize, status, modelPath, statusDesc)
	if err != nil {
		log.Info("update status error." + err.Error())
	}
}

func SaveNewNameModel(ctx *context.Context) {
	if !ctx.Repo.CanWrite(models.UnitTypeModelManage) {
		ctx.Error(403, ctx.Tr("repo.model_noright"))
		return
	}
	name := ctx.Query("name")
	if name == "" {
		ctx.Error(500, fmt.Sprintf("name or version is null."))
		return
	}

	aimodels := models.QueryModelByName(name, ctx.Repo.Repository.ID)
	if len(aimodels) > 0 {
		ctx.Error(500, ctx.Tr("repo.model_rename"))
		return
	}
	SaveModel(ctx)
	ctx.Status(200)
	log.Info("save model end.")
}

func SaveLocalModel(ctx *context.Context) {
	if !ctx.Repo.CanWrite(models.UnitTypeModelManage) {
		ctx.Error(403, ctx.Tr("repo.model_noright"))
		return
	}
	re := map[string]string{
		"code": "-1",
	}
	log.Info("save SaveLocalModel start.")
	uuid := uuid.NewV4()
	id := uuid.String()
	name := ctx.Query("name")
	version := ctx.Query("version")
	if version == "" {
		version = "0.0.1"
	}
	label := ctx.Query("label")
	description := ctx.Query("description")
	engine := ctx.QueryInt("engine")
	taskType := ctx.QueryInt("type")
	modelActualPath := ""
	if taskType == models.TypeCloudBrainOne {
		destKeyNamePrefix := Model_prefix + models.AttachmentRelativePath(id) + "/"
		modelActualPath = setting.Attachment.Minio.Bucket + "/" + destKeyNamePrefix
	} else if taskType == models.TypeCloudBrainTwo {
		destKeyNamePrefix := Model_prefix + models.AttachmentRelativePath(id) + "/"
		modelActualPath = setting.Bucket + "/" + destKeyNamePrefix
	} else {
		re["msg"] = "type is error."
		ctx.JSON(200, re)
		return
	}
	var lastNewModelId string
	repoId := ctx.Repo.Repository.ID
	aimodels := models.QueryModelByName(name, repoId)
	if len(aimodels) > 0 {
		for _, model := range aimodels {
			if model.Version == version {
				re["msg"] = ctx.Tr("repo.model.manage.create_error")
				ctx.JSON(200, re)
				return
			}
			if model.New == MODEL_LATEST {
				lastNewModelId = model.ID
			}
		}
	}
	model := &models.AiModelManage{
		ID:            id,
		Version:       version,
		ModelType:     1,
		VersionCount:  len(aimodels) + 1,
		Label:         label,
		Name:          name,
		Description:   description,
		New:           MODEL_LATEST,
		Type:          taskType,
		Path:          modelActualPath,
		Size:          0,
		AttachmentId:  "",
		RepoId:        repoId,
		UserId:        ctx.User.ID,
		Engine:        int64(engine),
		TrainTaskInfo: "",
		Accuracy:      "",
		Status:        STATUS_FINISHED,
	}

	err := models.SaveModelToDb(model)
	if err != nil {
		re["msg"] = err.Error()
		ctx.JSON(200, re)
		return
	}
	if len(lastNewModelId) > 0 {
		//udpate status and version count
		models.ModifyModelNewProperty(lastNewModelId, MODEL_NOT_LATEST, 0)
	}
	var units []models.RepoUnit
	var deleteUnitTypes []models.UnitType
	units = append(units, models.RepoUnit{
		RepoID: ctx.Repo.Repository.ID,
		Type:   models.UnitTypeModelManage,
		Config: &models.ModelManageConfig{
			EnableModelManage: true,
		},
	})
	deleteUnitTypes = append(deleteUnitTypes, models.UnitTypeModelManage)

	models.UpdateRepositoryUnits(ctx.Repo.Repository, units, deleteUnitTypes)

	log.Info("save model end.")
	notification.NotifyOtherTask(ctx.User, ctx.Repo.Repository, id, name, models.ActionCreateNewModelTask)
	re["code"] = "0"
	re["id"] = id
	ctx.JSON(200, re)
}

func getSize(files []storage.FileInfo) int64 {
	var size int64
	for _, file := range files {
		size += file.Size
	}
	return size
}

func UpdateModelSize(modeluuid string) {
	model, err := models.QueryModelById(modeluuid)
	if err == nil {
		if model.Type == models.TypeCloudBrainOne {
			if strings.HasPrefix(model.Path, setting.Attachment.Minio.Bucket+"/"+Model_prefix) {
				files, err := storage.GetAllObjectByBucketAndPrefixMinio(setting.Attachment.Minio.Bucket, model.Path[len(setting.Attachment.Minio.Bucket)+1:])
				if err != nil {
					log.Info("Failed to query model size from minio. id=" + modeluuid)
				}
				size := getSize(files)
				models.ModifyModelSize(modeluuid, size)
			}
		} else if model.Type == models.TypeCloudBrainTwo {
			if strings.HasPrefix(model.Path, setting.Bucket+"/"+Model_prefix) {
				files, err := storage.GetAllObjectByBucketAndPrefix(setting.Bucket, model.Path[len(setting.Bucket)+1:])
				if err != nil {
					log.Info("Failed to query model size from obs. id=" + modeluuid)
				}
				size := getSize(files)
				models.ModifyModelSize(modeluuid, size)
			}
		}
	} else {
		log.Info("not found model,uuid=" + modeluuid)
	}
}

func SaveModel(ctx *context.Context) {
	if !ctx.Repo.CanWrite(models.UnitTypeModelManage) {
		ctx.Error(403, ctx.Tr("repo.model_noright"))
		return
	}
	log.Info("save model start.")
	JobId := ctx.Query("jobId")
	VersionName := ctx.Query("versionName")
	name := ctx.Query("name")
	version := ctx.Query("version")
	label := ctx.Query("label")
	description := ctx.Query("description")
	engine := ctx.QueryInt("engine")
	modelSelectedFile := ctx.Query("modelSelectedFile")
	log.Info("engine=" + fmt.Sprint(engine) + " modelSelectedFile=" + modelSelectedFile)
	re := map[string]string{
		"code": "-1",
	}
	if JobId == "" || VersionName == "" {
		re["msg"] = "JobId or VersionName is null."
		ctx.JSON(200, re)
		return
	}
	if modelSelectedFile == "" {
		re["msg"] = "Not selected model file."
		ctx.JSON(200, re)
		return
	}
	if name == "" || version == "" {
		re["msg"] = "name or version is null."
		ctx.JSON(200, re)
		return
	}
	id, err := saveModelByParameters(JobId, VersionName, name, version, label, description, engine, ctx)
	if err != nil {
		log.Info("save model error." + err.Error())
		re["msg"] = err.Error()
	} else {
		re["code"] = "0"
		re["id"] = id
	}
	ctx.JSON(200, re)
	log.Info("save model end.")
}

func downloadModelFromCloudBrainTwo(modelUUID string, jobName string, parentDir string, trainUrl string, modelSelectedFile string) (string, int64, error) {
	objectkey := strings.TrimPrefix(path.Join(setting.TrainJobModelPath, jobName, setting.OutPutPath, parentDir), "/")
	if trainUrl != "" {
		objectkey = strings.Trim(trainUrl[len(setting.Bucket)+1:], "/")
	}
	prefix := objectkey + "/"

	filterFiles := strings.Split(modelSelectedFile, ";")
	Files := make([]string, 0)
	for _, shortFile := range filterFiles {
		Files = append(Files, prefix+shortFile)
	}
	totalSize := storage.ObsGetFilesSize(setting.Bucket, Files)
	if float64(totalSize) > setting.MaxModelSize*MODEL_MAX_SIZE {
		return "", 0, errors.New("Cannot create model, as model is exceed " + fmt.Sprint(setting.MaxModelSize) + "G.")
	}

	modelDbResult, err := storage.GetOneLevelAllObjectUnderDir(setting.Bucket, objectkey, "")
	log.Info("bucket=" + setting.Bucket + "  objectkey=" + objectkey)
	if err != nil {
		log.Info("get TrainJobListModel failed:", err)
		return "", 0, err
	}
	if len(modelDbResult) == 0 {
		return "", 0, errors.New("Cannot create model, as model is empty.")
	}

	destKeyNamePrefix := Model_prefix + models.AttachmentRelativePath(modelUUID) + "/"
	size, err := storage.ObsCopyManyFile(setting.Bucket, prefix, setting.Bucket, destKeyNamePrefix, filterFiles)

	dataActualPath := setting.Bucket + "/" + destKeyNamePrefix
	return dataActualPath, size, nil
}

func downloadModelFromCloudBrainOne(modelUUID string, jobName string, parentDir string, trainUrl string, modelSelectedFile string) (string, int64, error) {
	modelActualPath := storage.GetMinioPath(jobName, "/model/")
	log.Info("modelActualPath=" + modelActualPath)
	modelSrcPrefix := setting.CBCodePathPrefix + jobName + "/model/"
	destKeyNamePrefix := Model_prefix + models.AttachmentRelativePath(modelUUID) + "/"
	bucketName := setting.Attachment.Minio.Bucket
	log.Info("destKeyNamePrefix=" + destKeyNamePrefix + "   modelSrcPrefix=" + modelSrcPrefix + "  bucket=" + bucketName)
	filterFiles := strings.Split(modelSelectedFile, ";")
	Files := make([]string, 0)
	for _, shortFile := range filterFiles {
		Files = append(Files, modelSrcPrefix+shortFile)
	}
	totalSize := storage.MinioGetFilesSize(bucketName, Files)
	if float64(totalSize) > setting.MaxModelSize*MODEL_MAX_SIZE {
		return "", 0, errors.New("Cannot create model, as model is exceed " + fmt.Sprint(setting.MaxModelSize) + "G.")
	}
	size, err := storage.MinioCopyFiles(bucketName, modelSrcPrefix, destKeyNamePrefix, filterFiles)
	if err == nil {
		dataActualPath := bucketName + "/" + destKeyNamePrefix
		return dataActualPath, size, nil
	} else {
		return "", 0, nil
	}
}

func DeleteModel(ctx *context.Context) {
	log.Info("delete model start.")
	id := ctx.Query("id")
	err := deleteModelByID(ctx, id)
	if err != nil {
		re := map[string]string{
			"code": "-1",
		}
		re["msg"] = err.Error()
		ctx.JSON(200, re)
	} else {
		ctx.JSON(200, map[string]string{
			"code": "0",
		})
	}
}

func deleteModelByID(ctx *context.Context, id string) error {
	log.Info("delete model start. id=" + id)
	model, err := models.QueryModelById(id)
	if !isCanDelete(ctx, model.UserId) {
		return errors.New(ctx.Tr("repo.model_noright"))
	}
	if err == nil {
		log.Info("bucket=" + setting.Bucket + " path=" + model.Path)
		if model.Type == models.TypeCloudBrainOne {
			if strings.HasPrefix(model.Path, setting.Bucket+"/"+Model_prefix) {
				err := storage.Attachments.DeleteDir(model.Path[len(setting.Bucket)+1:])
				if err != nil {
					log.Info("Failed to delete model. id=" + id)
					return err
				}
			}
		} else if model.Type == models.TypeCloudBrainTwo {
			if strings.HasPrefix(model.Path, setting.Bucket+"/"+Model_prefix) {
				err := storage.ObsRemoveObject(setting.Bucket, model.Path[len(setting.Bucket)+1:])
				if err != nil {
					log.Info("Failed to delete model. id=" + id)
					return err
				}
			}
		}

		err = models.DeleteModelById(id)
		if err == nil { //find a model to change new
			aimodels := models.QueryModelByName(model.Name, model.RepoId)
			if model.New == MODEL_LATEST {
				if len(aimodels) > 0 {
					//udpate status and version count
					models.ModifyModelNewProperty(aimodels[0].ID, MODEL_LATEST, len(aimodels))
				}
			} else {
				for _, tmpModel := range aimodels {
					if tmpModel.New == MODEL_LATEST {
						models.ModifyModelNewProperty(tmpModel.ID, MODEL_LATEST, len(aimodels))
						break
					}
				}
			}
		}
	}
	return err
}

func QueryModelByParameters(repoId int64, page int) ([]*models.AiModelManage, int64, error) {

	return models.QueryModel(&models.AiModelQueryOptions{
		ListOptions: models.ListOptions{
			Page:     page,
			PageSize: setting.UI.IssuePagingNum,
		},
		RepoID: repoId,
		Type:   -1,
		New:    MODEL_LATEST,
		Status: -1,
	})
}

func DownloadMultiModelFile(ctx *context.Context) {
	log.Info("DownloadMultiModelFile start.")
	id := ctx.Query("id")
	log.Info("id=" + id)
	task, err := models.QueryModelById(id)
	if err != nil {
		log.Error("no such model!", err.Error())
		ctx.ServerError("no such model:", err)
		return
	}
	if !isOper(ctx, task.UserId) {
		ctx.NotFound(ctx.Req.URL.RequestURI(), nil)
		return
	}

	path := Model_prefix + models.AttachmentRelativePath(id) + "/"
	if task.Type == models.TypeCloudBrainTwo {
		downloadFromCloudBrainTwo(path, task, ctx, id)
	} else if task.Type == models.TypeCloudBrainOne {
		downloadFromCloudBrainOne(path, task, ctx, id)
	}
}

func MinioDownloadManyFile(path string, ctx *context.Context, returnFileName string, allFile []storage.FileInfo) {
	ctx.Resp.Header().Set("Content-Disposition", "attachment; filename="+url.QueryEscape(returnFileName))
	ctx.Resp.Header().Set("Content-Type", "application/octet-stream")
	w := zip.NewWriter(ctx.Resp)
	defer w.Close()
	for _, oneFile := range allFile {
		if oneFile.IsDir {
			log.Info("zip dir name:" + oneFile.FileName)
		} else {
			log.Info("zip file name:" + oneFile.FileName)
			fDest, err := w.Create(oneFile.FileName)
			if err != nil {
				log.Info("create zip entry error, download file failed: %s\n", err.Error())
				ctx.ServerError("download file failed:", err)
				return
			}
			log.Info("minio file path=" + (path + oneFile.FileName))
			body, err := storage.Attachments.DownloadAFile(setting.Attachment.Minio.Bucket, path+oneFile.FileName)
			if err != nil {
				log.Info("download file failed: %s\n", err.Error())
				ctx.ServerError("download file failed:", err)
				return
			} else {
				defer body.Close()
				p := make([]byte, 1024)
				var readErr error
				var readCount int
				// 读取对象内容
				for {
					readCount, readErr = body.Read(p)
					if readCount > 0 {
						fDest.Write(p[:readCount])
					}
					if readErr != nil {
						break
					}
				}
			}
		}
	}

}

func downloadFromCloudBrainOne(path string, task *models.AiModelManage, ctx *context.Context, id string) {
	allFile, err := storage.GetAllObjectByBucketAndPrefixMinio(setting.Attachment.Minio.Bucket, path)
	if err == nil {
		//count++
		models.ModifyModelDownloadCount(id)
		returnFileName := task.Name + "_" + task.Version + ".zip"
		MinioDownloadManyFile(path, ctx, returnFileName, allFile)
	} else {
		log.Info("error,msg=" + err.Error())
		ctx.ServerError("no file to download.", err)
	}
}

func ObsDownloadManyFile(path string, ctx *context.Context, returnFileName string, allFile []storage.FileInfo) {
	ctx.Resp.Header().Set("Content-Disposition", "attachment; filename="+url.QueryEscape(returnFileName))
	ctx.Resp.Header().Set("Content-Type", "application/octet-stream")
	w := zip.NewWriter(ctx.Resp)
	defer w.Close()
	for _, oneFile := range allFile {
		if oneFile.IsDir {
			log.Info("zip dir name:" + oneFile.FileName)
		} else {
			log.Info("zip file name:" + oneFile.FileName)
			fDest, err := w.Create(oneFile.FileName)
			if err != nil {
				log.Info("create zip entry error, download file failed: %s\n", err.Error())
				ctx.ServerError("download file failed:", err)
				return
			}
			body, err := storage.ObsDownloadAFile(setting.Bucket, path+oneFile.FileName)
			if err != nil {
				log.Info("download file failed: %s\n", err.Error())
				ctx.ServerError("download file failed:", err)
				return
			} else {
				defer body.Close()
				p := make([]byte, 1024)
				var readErr error
				var readCount int
				// 读取对象内容
				for {
					readCount, readErr = body.Read(p)
					if readCount > 0 {
						fDest.Write(p[:readCount])
					}
					if readErr != nil {
						break
					}
				}
			}
		}
	}
}

func downloadFromCloudBrainTwo(path string, task *models.AiModelManage, ctx *context.Context, id string) {
	allFile, err := storage.GetAllObjectByBucketAndPrefix(setting.Bucket, path)
	if err == nil {
		//count++
		models.ModifyModelDownloadCount(id)
		returnFileName := task.Name + "_" + task.Version + ".zip"
		ObsDownloadManyFile(path, ctx, returnFileName, allFile)
	} else {
		log.Info("error,msg=" + err.Error())
		ctx.ServerError("no file to download.", err)
	}
}

func QueryTrainJobVersionList(ctx *context.Context) {
	log.Info("query train job version list. start.")
	JobID := ctx.Query("jobId")
	if JobID == "" {
		JobID = ctx.Query("JobId")
	}

	VersionListTasks, count, err := models.QueryModelTrainJobVersionList(JobID)

	log.Info("query return count=" + fmt.Sprint(count))

	if err != nil {
		ctx.ServerError("QueryTrainJobList:", err)
	} else {
		ctx.JSON(200, VersionListTasks)
	}
}

func QueryTrainJobList(ctx *context.Context) {
	log.Info("query train job list. start.")
	repoId := ctx.QueryInt64("repoId")

	VersionListTasks, count, err := models.QueryModelTrainJobList(repoId)
	log.Info("query return count=" + fmt.Sprint(count))

	if err != nil {
		ctx.ServerError("QueryTrainJobList:", err)
	} else {
		ctx.JSON(200, VersionListTasks)
	}

}

func QueryTrainModelFileById(ctx *context.Context) ([]storage.FileInfo, error) {
	JobID := ctx.Query("jobId")
	VersionListTasks, count, err := models.QueryModelTrainJobVersionList(JobID)
	if err == nil {
		if count == 1 {
			task := VersionListTasks[0]
			jobName := task.JobName
			taskType := task.Type
			VersionName := task.VersionName
			modelDbResult, err := getModelFromObjectSave(jobName, taskType, VersionName)
			return modelDbResult, err
		}
	}
	log.Info("get TypeCloudBrainTwo TrainJobListModel failed:", err)
	return nil, errors.New("Not found task.")
}

func getModelFromObjectSave(jobName string, taskType int, VersionName string) ([]storage.FileInfo, error) {
	if taskType == models.TypeCloudBrainTwo {
		objectkey := path.Join(setting.TrainJobModelPath, jobName, setting.OutPutPath, VersionName) + "/"
		modelDbResult, err := storage.GetAllObjectByBucketAndPrefix(setting.Bucket, objectkey)
		log.Info("bucket=" + setting.Bucket + "  objectkey=" + objectkey)
		if err != nil {
			log.Info("get TypeCloudBrainTwo TrainJobListModel failed:", err)
			return nil, err
		} else {
			return modelDbResult, nil
		}
	} else if taskType == models.TypeCloudBrainOne {
		modelSrcPrefix := setting.CBCodePathPrefix + jobName + "/model/"
		bucketName := setting.Attachment.Minio.Bucket
		modelDbResult, err := storage.GetAllObjectByBucketAndPrefixMinio(bucketName, modelSrcPrefix)
		if err != nil {
			log.Info("get TypeCloudBrainOne TrainJobListModel failed:", err)
			return nil, err
		} else {
			return modelDbResult, nil
		}
	}
	return nil, errors.New("Not support.")
}

func QueryTrainModelList(ctx *context.Context) {
	log.Info("query train job list. start.")
	jobName := ctx.Query("jobName")
	taskType := ctx.QueryInt("type")
	VersionName := ctx.Query("versionName")
	if VersionName == "" {
		VersionName = ctx.Query("VersionName")
	}
	modelDbResult, err := getModelFromObjectSave(jobName, taskType, VersionName)
	if err != nil {
		log.Info("get TypeCloudBrainTwo TrainJobListModel failed:", err)
		ctx.JSON(200, "")
	} else {
		ctx.JSON(200, modelDbResult)
		return
	}
}

func DownloadSingleModelFile(ctx *context.Context) {
	log.Info("DownloadSingleModelFile start.")
	id := ctx.Params(":ID")
	parentDir := ctx.Query("parentDir")
	fileName := ctx.Query("fileName")
	path := Model_prefix + models.AttachmentRelativePath(id) + "/" + parentDir + fileName
	task, err := models.QueryModelById(id)
	if err != nil {
		log.Error("no such model!", err.Error())
		ctx.ServerError("no such model:", err)
		return
	}
	if !isOper(ctx, task.UserId) {
		ctx.NotFound(ctx.Req.URL.RequestURI(), nil)
		return
	}
	if task.Type == models.TypeCloudBrainTwo {
		if setting.PROXYURL != "" {
			body, err := storage.ObsDownloadAFile(setting.Bucket, path)
			if err != nil {
				log.Info("download error.")
			} else {
				//count++
				models.ModifyModelDownloadCount(id)
				defer body.Close()
				ctx.Resp.Header().Set("Content-Disposition", "attachment; filename="+fileName)
				ctx.Resp.Header().Set("Content-Type", "application/octet-stream")
				p := make([]byte, 1024)
				var readErr error
				var readCount int
				// 读取对象内容
				for {
					readCount, readErr = body.Read(p)
					if readCount > 0 {
						ctx.Resp.Write(p[:readCount])
						//fmt.Printf("%s", p[:readCount])
					}
					if readErr != nil {
						break
					}
				}
			}
		} else {
			url, err := storage.GetObsCreateSignedUrlByBucketAndKey(setting.Bucket, path)
			if err != nil {
				log.Error("GetObsCreateSignedUrl failed: %v", err.Error(), ctx.Data["msgID"])
				ctx.ServerError("GetObsCreateSignedUrl", err)
				return
			}
			//count++
			models.ModifyModelDownloadCount(id)
			http.Redirect(ctx.Resp, ctx.Req.Request, url, http.StatusMovedPermanently)
		}
	} else if task.Type == models.TypeCloudBrainOne {
		log.Info("start to down load minio file.")
		url, err := storage.Attachments.PresignedGetURL(path, fileName)
		if err != nil {
			log.Error("Get minio get SignedUrl failed: %v", err.Error(), ctx.Data["msgID"])
			ctx.ServerError("Get minio get SignedUrl failed", err)
			return
		}
		models.ModifyModelDownloadCount(id)
		http.Redirect(ctx.Resp, ctx.Req.Request, url, http.StatusMovedPermanently)
	}

}

func ShowModelInfo(ctx *context.Context) {
	ctx.Data["ID"] = ctx.Query("id")
	ctx.Data["name"] = ctx.Query("name")
	ctx.Data["isModelManage"] = true
	ctx.Data["ModelManageAccess"] = ctx.Repo.CanWrite(models.UnitTypeModelManage)

	ctx.HTML(200, tplModelInfo)
}

func QueryModelById(ctx *context.Context) {
	id := ctx.Query("id")
	model, err := models.QueryModelById(id)
	if err == nil {
		model.IsCanOper = isOper(ctx, model.UserId)
		model.IsCanDelete = isCanDelete(ctx, model.UserId)
		removeIpInfo(model)
		ctx.JSON(http.StatusOK, model)
	} else {
		ctx.JSON(http.StatusNotFound, nil)
	}
}

func ShowSingleModel(ctx *context.Context) {
	name := ctx.Query("name")

	log.Info("Show single ModelInfo start.name=" + name)
	models := models.QueryModelByName(name, ctx.Repo.Repository.ID)

	userIds := make([]int64, len(models))
	for i, model := range models {
		model.IsCanOper = isOper(ctx, model.UserId)
		model.IsCanDelete = isCanDelete(ctx, model.UserId)
		userIds[i] = model.UserId
	}
	userNameMap := queryUserName(userIds)

	for _, model := range models {
		removeIpInfo(model)
		value := userNameMap[model.UserId]
		if value != nil {
			model.UserName = value.Name
			model.UserRelAvatarLink = value.RelAvatarLink()
		}
	}
	ctx.JSON(http.StatusOK, models)
}

func removeIpInfo(model *models.AiModelManage) {
	reg, _ := regexp.Compile(`[[:digit:]]{1,3}\.[[:digit:]]{1,3}\.[[:digit:]]{1,3}\.[[:digit:]]{1,3}`)
	taskInfo := model.TrainTaskInfo
	taskInfo = reg.ReplaceAllString(taskInfo, "")
	model.TrainTaskInfo = taskInfo
}

func queryUserName(intSlice []int64) map[int64]*models.User {
	keys := make(map[int64]string)
	uniqueElements := []int64{}
	for _, entry := range intSlice {
		if _, value := keys[entry]; !value {
			keys[entry] = ""
			uniqueElements = append(uniqueElements, entry)
		}
	}
	result := make(map[int64]*models.User)
	userLists, err := models.GetUsersByIDs(uniqueElements)
	if err == nil {
		for _, user := range userLists {
			result[user.ID] = user
		}
	}
	return result
}

func ShowOneVersionOtherModel(ctx *context.Context) {
	repoId := ctx.Repo.Repository.ID
	name := ctx.Query("name")
	aimodels := models.QueryModelByName(name, repoId)

	userIds := make([]int64, len(aimodels))
	for i, model := range aimodels {
		model.IsCanOper = isOper(ctx, model.UserId)
		model.IsCanDelete = isCanDelete(ctx, model.UserId)
		userIds[i] = model.UserId
	}
	userNameMap := queryUserName(userIds)

	for _, model := range aimodels {
		removeIpInfo(model)
		value := userNameMap[model.UserId]
		if value != nil {
			model.UserName = value.Name
			model.UserRelAvatarLink = value.RelAvatarLink()
		}
	}

	if len(aimodels) > 0 {
		ctx.JSON(200, aimodels[1:])
	} else {
		ctx.JSON(200, aimodels)
	}
}

func SetModelCount(ctx *context.Context) {
	repoId := ctx.Repo.Repository.ID
	Type := -1
	_, count, _ := models.QueryModel(&models.AiModelQueryOptions{
		ListOptions: models.ListOptions{
			Page:     1,
			PageSize: 2,
		},
		RepoID: repoId,
		Type:   Type,
		New:    MODEL_LATEST,
		Status: -1,
	})
	ctx.Data["MODEL_COUNT"] = count
}

func ShowModelTemplate(ctx *context.Context) {
	ctx.Data["isModelManage"] = true
	repoId := ctx.Repo.Repository.ID
	SetModelCount(ctx)
	ctx.Data["ModelManageAccess"] = ctx.Repo.CanWrite(models.UnitTypeModelManage)
	_, trainCount, _ := models.QueryModelTrainJobList(repoId)
	log.Info("query train count=" + fmt.Sprint(trainCount))
	ctx.Data["TRAIN_COUNT"] = trainCount
	ctx.HTML(200, tplModelManageIndex)
}

func isQueryRight(ctx *context.Context) bool {
	if ctx.Repo.Repository.IsPrivate {
		if ctx.Repo.CanRead(models.UnitTypeModelManage) || ctx.User.IsAdmin || ctx.Repo.IsAdmin() || ctx.Repo.IsOwner() {
			return true
		}
		return false
	} else {
		return true
	}
}

func isCanDelete(ctx *context.Context, modelUserId int64) bool {
	if ctx.User == nil {
		return false
	}
	if ctx.User.IsAdmin || ctx.User.ID == modelUserId {
		return true
	}
	if ctx.Repo.IsOwner() {
		return true
	}
	return false
}

func isOper(ctx *context.Context, modelUserId int64) bool {
	if ctx.User == nil {
		return false
	}
	if ctx.User.IsAdmin || ctx.User.ID == modelUserId {
		return true
	}
	return false
}

func ShowModelPageInfo(ctx *context.Context) {
	log.Info("ShowModelInfo start.")
	if !isQueryRight(ctx) {
		ctx.NotFound(ctx.Req.URL.RequestURI(), nil)
		return
	}
	page := ctx.QueryInt("page")
	if page <= 0 {
		page = 1
	}
	pageSize := ctx.QueryInt("pageSize")
	if pageSize <= 0 {
		pageSize = setting.UI.IssuePagingNum
	}
	repoId := ctx.Repo.Repository.ID
	Type := -1
	modelResult, count, err := models.QueryModel(&models.AiModelQueryOptions{
		ListOptions: models.ListOptions{
			Page:     page,
			PageSize: pageSize,
		},
		RepoID: repoId,
		Type:   Type,
		New:    MODEL_LATEST,
		Status: -1,
	})
	if err != nil {
		ctx.ServerError("Cloudbrain", err)
		return
	}

	userIds := make([]int64, len(modelResult))
	for i, model := range modelResult {
		model.IsCanOper = isOper(ctx, model.UserId)
		model.IsCanDelete = isCanDelete(ctx, model.UserId)
		userIds[i] = model.UserId
	}

	userNameMap := queryUserName(userIds)

	for _, model := range modelResult {
		removeIpInfo(model)
		value := userNameMap[model.UserId]
		if value != nil {
			model.UserName = value.Name
			model.UserRelAvatarLink = value.RelAvatarLink()
		}
	}

	mapInterface := make(map[string]interface{})
	mapInterface["data"] = modelResult
	mapInterface["count"] = count
	ctx.JSON(http.StatusOK, mapInterface)
}

func ModifyModel(id string, description string) error {
	err := models.ModifyModelDescription(id, description)
	if err == nil {
		log.Info("modify success.")
	} else {
		log.Info("Failed to modify.id=" + id + "  desc=" + description + " error:" + err.Error())
	}
	return err
}

func ModifyModelInfo(ctx *context.Context) {
	log.Info("modify model start.")
	id := ctx.Query("id")
	description := ctx.Query("description")
	re := map[string]string{
		"code": "-1",
	}
	task, err := models.QueryModelById(id)
	if err != nil {
		re["msg"] = err.Error()
		log.Error("no such model!", err.Error())
		ctx.JSON(200, re)
		return
	}
	if !isOper(ctx, task.UserId) {
		re["msg"] = "No right to operation."
		ctx.JSON(200, re)
		return
	}
	if task.ModelType == MODEL_LOCAL_TYPE {
		name := ctx.Query("name")
		label := ctx.Query("label")
		description := ctx.Query("description")
		engine := ctx.QueryInt("engine")
		aimodels := models.QueryModelByName(name, task.RepoId)
		if aimodels != nil && len(aimodels) > 0 {
			if len(aimodels) == 1 {
				if aimodels[0].ID != task.ID {
					re["msg"] = ctx.Tr("repo.model.manage.create_error")
					ctx.JSON(200, re)
					return
				}
			} else {
				re["msg"] = ctx.Tr("repo.model.manage.create_error")
				ctx.JSON(200, re)
				return
			}
		}
		err = models.ModifyLocalModel(id, name, label, description, engine)

	} else {
		err = ModifyModel(id, description)
	}

	if err != nil {
		re["msg"] = err.Error()
		ctx.JSON(200, re)
		return
	} else {
		re["code"] = "0"
		ctx.JSON(200, re)
	}
}

func QueryModelListForPredict(ctx *context.Context) {
	repoId := ctx.Repo.Repository.ID
	modelResult, count, err := models.QueryModel(&models.AiModelQueryOptions{
		ListOptions: models.ListOptions{
			Page:     -1,
			PageSize: -1,
		},
		RepoID: repoId,
		Type:   ctx.QueryInt("type"),
		New:    -1,
		Status: 0,
	})
	if err != nil {
		ctx.ServerError("Cloudbrain", err)
		return
	}
	log.Info("query return count=" + fmt.Sprint(count))

	nameList := make([]string, 0)

	nameMap := make(map[string][]*models.AiModelManage)
	for _, model := range modelResult {
		removeIpInfo(model)
		if _, value := nameMap[model.Name]; !value {
			models := make([]*models.AiModelManage, 0)
			models = append(models, model)
			nameMap[model.Name] = models
			nameList = append(nameList, model.Name)
		} else {
			nameMap[model.Name] = append(nameMap[model.Name], model)
		}
	}

	mapInterface := make(map[string]interface{})
	mapInterface["nameList"] = nameList
	mapInterface["nameMap"] = nameMap
	ctx.JSON(http.StatusOK, mapInterface)
}

func QueryModelFileForPredict(ctx *context.Context) {
	id := ctx.Query("id")
	if id == "" {
		id = ctx.Query("ID")
	}
	ctx.JSON(http.StatusOK, QueryModelFileByID(id))
}

func QueryModelFileByID(id string) []storage.FileInfo {
	model, err := models.QueryModelById(id)
	if err == nil {
		if model.Type == models.TypeCloudBrainTwo {
			prefix := model.Path[len(setting.Bucket)+1:]
			fileinfos, _ := storage.GetAllObjectByBucketAndPrefix(setting.Bucket, prefix)
			return fileinfos
		} else if model.Type == models.TypeCloudBrainOne {
			prefix := model.Path[len(setting.Attachment.Minio.Bucket)+1:]
			fileinfos, _ := storage.GetAllObjectByBucketAndPrefixMinio(setting.Attachment.Minio.Bucket, prefix)
			return fileinfos
		}
	} else {
		log.Error("no such model!", err.Error())
	}
	return nil
}

func QueryOneLevelModelFile(ctx *context.Context) {
	id := ctx.Query("id")
	if id == "" {
		id = ctx.Query("ID")
	}
	parentDir := ctx.Query("parentDir")
	model, err := models.QueryModelById(id)
	if err != nil {
		log.Error("no such model!", err.Error())
		ctx.ServerError("no such model:", err)
		return
	}
	if model.Type == models.TypeCloudBrainTwo {
		log.Info("TypeCloudBrainTwo list model file.")
		prefix := model.Path[len(setting.Bucket)+1:]
		fileinfos, _ := storage.GetOneLevelAllObjectUnderDir(setting.Bucket, prefix, parentDir)
		if fileinfos == nil {
			fileinfos = make([]storage.FileInfo, 0)
		}
		ctx.JSON(http.StatusOK, fileinfos)
	} else if model.Type == models.TypeCloudBrainOne {
		log.Info("TypeCloudBrainOne list model file.")
		prefix := model.Path[len(setting.Attachment.Minio.Bucket)+1:]
		fileinfos, _ := storage.GetOneLevelAllObjectUnderDirMinio(setting.Attachment.Minio.Bucket, prefix, parentDir)
		if fileinfos == nil {
			fileinfos = make([]storage.FileInfo, 0)
		}
		ctx.JSON(http.StatusOK, fileinfos)
	}
}

func CreateLocalModel(ctx *context.Context) {
	ctx.Data["isModelManage"] = true
	ctx.Data["ModelManageAccess"] = ctx.Repo.CanWrite(models.UnitTypeModelManage)

	ctx.HTML(200, tplCreateLocalModelInfo)
}

func CreateLocalModelForUpload(ctx *context.Context) {
	ctx.Data["uuid"] = ctx.Query("uuid")
	ctx.Data["isModelManage"] = true
	ctx.Data["ModelManageAccess"] = ctx.Repo.CanWrite(models.UnitTypeModelManage)

	ctx.HTML(200, tplCreateLocalForUploadModelInfo)
}

func CreateOnlineModel(ctx *context.Context) {
	ctx.Data["isModelManage"] = true
	ctx.Data["ModelManageAccess"] = ctx.Repo.CanWrite(models.UnitTypeModelManage)

	ctx.HTML(200, tplCreateOnlineModelInfo)
}