fyp/models.go

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))
		}
	})
}