package main import ( "bytes" "fmt" "image" "image/color" _ "image/png" "io" "net/http" "os" "path" "strconv" ) const ( FAILED_PREPARING = -1 PREPARING = iota CONFIRM_PRE_TRAINING ) type BaseModel struct { Name string Status int Id string } func modelUpdateStatus(handle *Handle, id string, status int) { _, err := handle.db.Exec("update models set status = $1 where id = $2", status, id) if err != nil { fmt.Println("Failed to update model status") fmt.Println(err) panic("TODO handle better") } } func loadBaseImage(handle *Handle, id string) { // TODO handle more types than png infile, err := os.Open(path.Join("savedData", id, "baseimage.png")) if err != nil { // TODO better logging fmt.Println(err) fmt.Printf("Failed to read image for model with id %s\n", id) modelUpdateStatus(handle, id, -1) return } defer infile.Close() src, format, err := image.Decode(infile) if err != nil { // TODO better logging fmt.Println(err) fmt.Printf("Failed to load image for model with id %s\n", id) modelUpdateStatus(handle, id, -1) return } if format != "png" { // TODO better logging fmt.Printf("Found unkown format '%s'\n", format) panic("Handle diferent files than .png") } var model_color string bounds := src.Bounds() width, height := bounds.Max.X, bounds.Max.Y switch src.ColorModel() { case color.Gray16Model: fallthrough case color.GrayModel: model_color = "greyscale" default: fmt.Println("Do not know how to handle this color model") if src.ColorModel() == color.RGBA64Model { fmt.Println("Color is rgb") } else if src.ColorModel() == color.NRGBAModel { fmt.Println("Color is nrgb") } else if src.ColorModel() == color.YCbCrModel { fmt.Println("Color is ycbcr") } else if src.ColorModel() == color.AlphaModel { fmt.Println("Color is alpha") } else if src.ColorModel() == color.CMYKModel { fmt.Println("Color is cmyk") } else { fmt.Println("Other so assuming color") } modelUpdateStatus(handle, id, -1) return } // Note: this also updates the status to 2 _, err = handle.db.Exec("update models set width=$1, height=$2, color_mode=$3, status=$4 where id=$5", width, height, model_color, CONFIRM_PRE_TRAINING, id) if err != nil { // TODO better logging fmt.Println(err) fmt.Printf("Could not update model\n") modelUpdateStatus(handle, id, -1) return } } func deleteModel(handle *Handle, id string, w http.ResponseWriter, c *Context, model BaseModel) { fmt.Printf("Removing model with id: %s\n", id) _, err := handle.db.Exec("delete from models where id=$1;", id) if err != nil { fmt.Println(err) panic("TODO handle better deleteModel failed delete database query") } model_path := path.Join("./savedData", id) fmt.Printf("Removing folder of model with id: %s at %s\n", id, model_path) err = os.RemoveAll(model_path) if err != nil { fmt.Println(err) panic("TODO handle better deleteModel failed to delete folder") } if c.Mode == HTML { // TODO move this to a constant so i don't forget w.WriteHeader(309) c.Mode = HTMLFULL } LoadBasedOnAnswer(c.Mode, w, "/models/delete.html", c.addMap(AnyMap{ "Model": model, })) } func handleModelsEndpoints(handle *Handle) { // TODO json handle.GetHTML("/models", func(w http.ResponseWriter, r *http.Request, c *Context) *Error { if c.Mode == JSON { panic("TODO JSON") } if !checkAuthLevel(1, w, r, c) { return nil } rows, err := handle.db.Query("select id, name from models where user_id=$1;", c.User.id) if err != nil { return error500(err) } type row struct { Name string Id string } got := []row{} for rows.Next() { var r row err = rows.Scan(&r.Id, &r.Name) if err != nil { return error500(err) } got = append(got, r) } LoadBasedOnAnswer(c.Mode, w, "/models/list.html", c.addMap(AnyMap{ "List": got, })) return nil }) handle.GetHTML("/models/add", AnswerTemplate("models/add.html", nil, 1)) // TODO json handle.Post("/models/add", func(w http.ResponseWriter, r *http.Request, c *Context) *Error { if c.Mode == JSON { panic("TODO JSON") } if !checkAuthLevel(1, w, r, c) { return nil } read_form, err := r.MultipartReader() if err != nil { LoadBasedOnAnswer(c.Mode, w, "models/add.html", c.addMap(nil)) return nil } var name string var file []byte for { part, err_part := read_form.NextPart() if err_part == io.EOF { break } else if err_part != nil { return &Error{code: http.StatusBadRequest} } if part.FormName() == "name" { buf := new(bytes.Buffer) buf.ReadFrom(part) name = buf.String() } if part.FormName() == "file" { buf := new(bytes.Buffer) buf.ReadFrom(part) file = buf.Bytes() } } if name == "" || len(file) == 0 { LoadBasedOnAnswer(c.Mode, w, "models/add.html", c.addMap(nil)) return nil } row, err := handle.db.Query("select id from models where name=$1 and user_id=$2;", name, c.User.id) if err != nil { return error500(err) } if row.Next() { LoadBasedOnAnswer(c.Mode, w, "models/add.html", c.addMap(AnyMap{ "NameFoundError": true, "Name": name, })) return nil } _, err = handle.db.Exec("insert into models (user_id, name) values ($1, $2)", c.User.id, name) if err != nil { return error500(err) } row, err = handle.db.Query("select id from models where name=$1 and user_id=$2;", name, c.User.id) if err != nil { return error500(err) } if !row.Next() { return &Error{code: http.StatusInternalServerError} } var id string err = row.Scan(&id) if err != nil { return error500(err) } // TODO mk this path configurable dir_path := path.Join("savedData", id) err = os.Mkdir(dir_path, os.ModePerm) if err != nil { return error500(err) } f, err := os.Create(path.Join(dir_path, "baseimage.png")) if err != nil { return error500(err) } defer f.Close() f.Write(file) fmt.Printf("Created model with id %s! Started to proccess image!\n", id) go loadBaseImage(handle, id) redirect("/models/edit?id="+id, c.Mode, w, r) return nil }) handle.GetHTML("/models/edit", func(w http.ResponseWriter, r *http.Request, c *Context) *Error { if !checkAuthLevel(1, w, r, c) { return nil } id, err := getIdFromUrl(r, "id") if err != nil { return errorCode(nil, http.StatusNotFound, AnyMap{ "NotFoundMessage": "Model not found", "GoBackLink": "/models", }) } // TODO handle admin users rows, err := handle.db.Query("select name, status, width, height, color_mode from models where id=$1 and user_id=$2;", id, c.User.id) if err != nil { return error500(err) } if !rows.Next() { return errorCode(nil, http.StatusNotFound, AnyMap{ "NotFoundMessage": "Model not found", "GoBackLink": "/models", }) } type rowmodel struct { Name string Status int Id string Width *int Height *int Color_mode *string } var model rowmodel = rowmodel{} model.Id = id err = rows.Scan(&model.Name, &model.Status, &model.Width, &model.Height, &model.Color_mode) if err != nil { return error500(err) } // Handle errors // All errors will be negative if model.Status < 0 { LoadBasedOnAnswer(c.Mode, w, "/models/edit.html", c.addMap(AnyMap{ "Model": model, })) return nil } switch model.Status { case PREPARING: LoadBasedOnAnswer(c.Mode, w, "/models/edit.html", c.addMap(AnyMap{ "Model": model, })) case CONFIRM_PRE_TRAINING: LoadBasedOnAnswer(c.Mode, w, "/models/edit.html", c.addMap(AnyMap{ "Model": model, })) default: fmt.Printf("Unkown Status: %d\n", model.Status) return error500(nil) } return nil }) handle.Delete("/models/delete", func(w http.ResponseWriter, r *http.Request, c *Context) *Error { if c.Mode == JSON { panic("TODO handle json on models/delete") } f, err := MyParseForm(r) if err != nil { return errorCode(err, 400, nil) } if !checkId(f, "id") { return errorCode(nil, http.StatusNotFound, AnyMap{ "NotFoundMessage": "Model not found", "GoBackLink": "/models", }) } id := f.Get("id") // TODO handle admin users rows, err := handle.db.Query("select name, status from models where id=$1 and user_id=$2;", id, c.User.id) if err != nil { return error500(err) } if !rows.Next() { return errorCode(nil, http.StatusNotFound, AnyMap{ "NotFoundMessage": "Model not found", "GoBackLink": "/models", }) } var model BaseModel = BaseModel{} model.Id = id err = rows.Scan(&model.Name, &model.Status) if err != nil { return error500(err) } switch model.Status { case FAILED_PREPARING: deleteModel(handle, id, w, c, model) return nil case CONFIRM_PRE_TRAINING: if checkEmpty(f, "name") { // TODO improve result return errorCode(nil, http.StatusBadRequest, nil) } name := f.Get("name") if name != model.Name { LoadError(w, "/models/edit.html", "delete-model-card", c.addMap(AnyMap{ "NameDoesNotMatch": true, "Model": model, })) return nil } deleteModel(handle, id, w, c, model) return nil default: panic("Do not know how to handle model in status:" + strconv.Itoa(model.Status)) } }) }