package models

import (
	"fmt"
	"main/configs"
	"time"
)

type Model struct {
	ID           int       `json:"id" gorm:"primary_key"`
	Name         string    `json:"name"`
	Type         string    `json:"type"`                     // (lora|ckp|hyper|ti)
	TriggerWords string    `json:"trigger_words"`            // 觸發詞
	BaseModel    string    `json:"base_model"`               // (SD1.5|SD2)
	ModelPath    string    `json:"model_path"`               // 模型路徑
	Status       string    `json:"status" default:"initial"` // (initial|ready|waiting|running|success|error)
	Progress     int       `json:"progress"`                 // (0-100)
	Image        string    `json:"image"`                    // 封面圖片實際地址
	Hash         string    `json:"hash"`                     // 模型哈希值
	Tags         string    `json:"tags"`
	UserID       int       `json:"user_id"`
	CreatedAt    time.Time `json:"created_at" gorm:"autoCreateTime"`
	UpdatedAt    time.Time `json:"updated_at" gorm:"autoUpdateTime"`
}

func init() {
	configs.ORMDB().AutoMigrate(&Model{})
}

func (model *Model) Train() (err error) {
	if model.Type == "lora" {
		fmt.Println("lora")
		return
	}
	if model.Type == "ckp" {
		fmt.Println("ckp")
		return
	}
	if model.Type == "hyper" {
		fmt.Println("hyper")
		return
	}
	if model.Type == "ti" {
		fmt.Println("ti")
		return
	}
	return
}