package models_train

import (
	"database/sql"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path"
	"strconv"
	"text/template"

	model_classes "git.andr3h3nriqu3s.com/andr3/fyp/logic/models/classes"
	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/models/utils"
	. "git.andr3h3nriqu3s.com/andr3/fyp/logic/utils"
)

func MakeDefenition(db *sql.DB, model_id string, target_accuracy int) (id string, err error) {
	id = ""
	rows, err := db.Query("insert into model_definition (model_id, target_accuracy) values ($1, $2) returning id;", model_id, target_accuracy)
	if err != nil {
		return
	}
	defer rows.Close()
	if !rows.Next() {
		return id, errors.New("Something wrong!")
	}
	err = rows.Scan(&id)
	return
}

type ModelDefinitionStatus int

const (
	MODEL_DEFINITION_STATUS_FAILED_TRAINING                       = -3
	MODEL_DEFINITION_STATUS_PRE_INIT        ModelDefinitionStatus = 1
	MODEL_DEFINITION_STATUS_INIT                                  = 2
	MODEL_DEFINITION_STATUS_TRAINING                              = 3
	MODEL_DEFINITION_STATUS_TRANIED                               = 4
	MODEL_DEFINITION_STATUS_READY                                 = 5
)

type LayerType int

const (
	LAYER_INPUT   LayerType = 1
	LAYER_DENSE             = 2
	LAYER_FLATTEN           = 3
)

func ModelDefinitionUpdateStatus(c *Context, id string, status ModelDefinitionStatus) (err error) {
	_, err = c.Db.Exec("update model_definition set status = $1 where id = $2", status, id)
	return
}

func MakeLayer(db *sql.DB, def_id string, layer_order int, layer_type LayerType, shape string) (err error) {
	_, err = db.Exec("insert into model_definition_layer (def_id, layer_order, layer_type, shape) values ($1, $2, $3, $4)", def_id, layer_order, layer_type, shape)
	return
}

func generateCvs(c *Context, run_path string, model_id string) (count int, err error) {

	classes, err := c.Db.Query("select count(*) from model_classes where model_id=$1;", model_id)
	if err != nil {
		return
	}
	defer classes.Close()
	if !classes.Next() {
		return
	}
	if err = classes.Scan(&count); err != nil {
		return
	}

	data, err := c.Db.Query("select mdp.id, mc.class_order, mdp.file_path from model_data_point as mdp inner join model_classes as mc on mc.id = mdp.class_id where mc.model_id = $1;", model_id)
	if err != nil {
		return
	}
	defer data.Close()

	f, err := os.Create(path.Join(run_path, "train.csv"))
	if err != nil {
		return
	}
	defer f.Close()
	f.Write([]byte("Id,Index\n"))

	for data.Next() {
		var id string
		var class_order int
		var file_path string
		if err = data.Scan(&id, &class_order, &file_path); err != nil {
			return
		}
		if file_path == "id://" {
			f.Write([]byte(id + "," + strconv.Itoa(class_order) + "\n"))
		} else {
			return count, errors.New("TODO generateCvs to file_path " + file_path)
		}
	}

	return
}

func trainDefinition(c *Context, model *BaseModel, definition_id string) (accuracy float64, err error) {
	accuracy = 0
	layers, err := c.Db.Query("select layer_type, shape from model_definition_layer where def_id=$1 order by layer_order asc;", definition_id)
	if err != nil {
		return
	}
	defer layers.Close()

	type layerrow struct {
		LayerType int
		Shape     string
	}

	got := []layerrow{}

	for layers.Next() {
		var row = layerrow{}
		if err = layers.Scan(&row.LayerType, &row.Shape); err != nil {
			return
		}
		row.Shape = shapeToSize(row.Shape)
		got = append(got, row)
	}

	// Generate run folder
	run_path := path.Join("/tmp", model.Id, "defs", definition_id)

	err = os.MkdirAll(run_path, os.ModePerm)
	if err != nil {
		return
	}

	_, err = generateCvs(c, run_path, model.Id)
	if err != nil {
		return
	}

	// Create python script
	f, err := os.Create(path.Join(run_path, "run.py"))
	if err != nil {
		return
	}
	defer f.Close()

	tmpl, err := template.New("python_model_template.py").ParseFiles("views/py/python_model_template.py")
	if err != nil {
		return
	}

	if err = tmpl.Execute(f, AnyMap{
		"Layers":    got,
		"Size":      got[0].Shape,
		"DataDir":   path.Join(getDir(), "savedData", model.Id, "data"),
		"RunPath":   run_path,
		"ColorMode": model.ImageMode,
        "Model": model,
	}); err != nil {
		return
	}

	// Run the command
	out, err := exec.Command("bash", "-c", fmt.Sprintf("cd %s && python run.py", run_path)).Output()
	if err != nil {
        c.Logger.Debug(string(out))
		return
	}

	// Copy result around
	result_path := path.Join("savedData", model.Id, "defs", definition_id)

	if err = os.MkdirAll(result_path, os.ModePerm); err != nil {
		return
	}

	if err = exec.Command("cp", "-r", path.Join(run_path, "model"), path.Join(result_path, "model")).Run(); err != nil {
		return
	}

	accuracy_file, err := os.Open(path.Join(run_path, "accuracy.val"))
	if err != nil {
		return
	}
	defer accuracy_file.Close()

	accuracy_file_bytes, err := io.ReadAll(accuracy_file)
	if err != nil {
		return
	}

	fmt.Println(string(accuracy_file_bytes))

	accuracy, err = strconv.ParseFloat(string(accuracy_file_bytes), 64)
	if err != nil {
		return
	}

	os.RemoveAll(run_path)
	return
}

func trainModel(c *Context, model *BaseModel) {
	definitionsRows, err := c.Db.Query("select id, target_accuracy from model_definition where status=$1 and model_id=$2", MODEL_DEFINITION_STATUS_INIT, model.Id)
	if err != nil {
		c.Logger.Error("Failed to trainModel!Err:")
		c.Logger.Error(err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}
	defer definitionsRows.Close()

	type row struct {
		id              string
		target_accuracy int
	}

	definitions := []row{}

	for definitionsRows.Next() {
		var rowv row
		if err = definitionsRows.Scan(&rowv.id, &rowv.target_accuracy); err != nil {
			c.Logger.Error("Failed to train Model Could not read definition from db!Err:")
			c.Logger.Error(err)
			ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
			return
		}
		definitions = append(definitions, rowv)
	}

	if len(definitions) == 0 {
		c.Logger.Error("No Definitions defined!")
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}

	for _, def := range definitions {
		accuracy, err := trainDefinition(c, model, def.id)
		if err != nil {
			c.Logger.Error("Failed to train definition!Err:")
			c.Logger.Error(err)
			ModelDefinitionUpdateStatus(c, def.id, MODEL_DEFINITION_STATUS_FAILED_TRAINING)
			continue
		}

		int_accuracy := int(accuracy * 100)

		if int_accuracy < def.target_accuracy {
			fmt.Printf("Failed to train definition! Accuracy less %d < %d\n", int_accuracy, def.target_accuracy)
			ModelDefinitionUpdateStatus(c, def.id, MODEL_DEFINITION_STATUS_FAILED_TRAINING)
			continue
		}

		_, err = c.Db.Exec("update model_definition set accuracy=$1, status=$2 where id=$3", int_accuracy, MODEL_DEFINITION_STATUS_TRANIED, def.id)
		if err != nil {
			fmt.Printf("Failed to train definition!Err:\n")
			fmt.Println(err)
			ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
			return
		}
	}

	rows, err := c.Db.Query("select id from model_definition where model_id=$1 and status=$2 order by accuracy desc limit 1;", model.Id, MODEL_DEFINITION_STATUS_TRANIED)
	if err != nil {
		c.Logger.Error("DB: failed to read definition")
		c.Logger.Error(err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}
	defer rows.Close()

	if !rows.Next() {
		// TODO Make the Model status have a message
		c.Logger.Error("All definitions failed to train!")
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}

	var id string
	if err = rows.Scan(&id); err != nil {
		c.Logger.Error("Failed to read id:")
		c.Logger.Error(err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}

	if _, err = c.Db.Exec("update model_definition set status=$1 where id=$2;", MODEL_DEFINITION_STATUS_READY, id); err != nil {
		c.Logger.Error("Failed to update model definition")
		c.Logger.Error(err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}

	to_delete, err := c.Db.Query("select id from model_definition where status != $1 and model_id=$2", MODEL_DEFINITION_STATUS_READY, model.Id)
	if err != nil {
		c.Logger.Error("Failed to select model_definition to delete")
		c.Logger.Error(err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}
	defer to_delete.Close()

	for to_delete.Next() {
		var id string
		if to_delete.Scan(&id); err != nil {
			c.Logger.Error("Failed to scan the id of a model_definition to delete")
			c.Logger.Error(err)
			ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
			return
		}
		os.RemoveAll(path.Join("savedData", model.Id, "defs", id))
	}

	// TODO Check if returning also works here
	if _, err = c.Db.Exec("delete from model_definition where status!=$1 and model_id=$2;", MODEL_DEFINITION_STATUS_READY, model.Id); err != nil {
		c.Logger.Error("Failed to delete model_definition")
		c.Logger.Error(err)
		ModelUpdateStatus(c, model.Id, FAILED_TRAINING)
		return
	}

	ModelUpdateStatus(c, model.Id, READY)
}

func removeFailedDataPoints(db *sql.DB, model *BaseModel) (err error) {
	rows, err := db.Query("select mdp.id from model_data_point as mdp join model_classes as mc on mc.id=mdp.class_id where mc.model_id=$1 and mdp.status=-1;", model.Id)
	if err != nil {
		return
	}
	defer rows.Close()

    base_path := path.Join("savedData", model.Id, "data")

	for rows.Next() {
		var dataPointId string
		err = rows.Scan(&dataPointId)
		if err != nil {
			return
		}
        err = os.RemoveAll(path.Join(base_path, dataPointId + model.Format))
        if err != nil {
            return
        }
	}

    _, err = db.Exec("delete from model_data_point as mdp using model_classes as mc where mdp.class_id = mc.id and mc.model_id=$1 and mdp.status=-1;", model.Id)
	return
}

func handleTrain(handle *Handle) {
	handle.Post("/models/train", func(w http.ResponseWriter, r *http.Request, c *Context) *Error {
		if !CheckAuthLevel(1, w, r, c) {
			return nil
		}
		if c.Mode == JSON {
			panic("TODO /models/train JSON")
		}

		r.ParseForm()
		f := r.Form

		number_of_models := 0
		accuracy := 0

		if !CheckId(f, "id") || CheckEmpty(f, "model_type") || !CheckNumber(f, "number_of_models", &number_of_models) || !CheckNumber(f, "accuracy", &accuracy) {
			fmt.Println(
				!CheckId(f, "id"), CheckEmpty(f, "model_type"), !CheckNumber(f, "number_of_models", &number_of_models), !CheckNumber(f, "accuracy", &accuracy),
			)
			// TODO improve this response
			return ErrorCode(nil, 400, c.AddMap(nil))
		}

		id := f.Get("id")
		model_type := f.Get("model_type")
		// Its not used rn
		_ = model_type

		// TODO check if the model has data
		/*rows, err := handle.Db.Query("select mc.name, mdp.file_path from model_classes as mc join model_data_point as mdp on mdp.class_id = mc.id where mdp.model_mode = 1 and mc.model_id = $1 limit 1;", id)
		if err != nil {
			return Error500(err)
		}
		defer rows.Close()

		if !rows.Next() {
			return Error500(err)
		}

		var name string
		var file_path string
		err = rows.Scan(&name, &file_path)
		if err != nil {
			return Error500(err)
		}*/

		model, err := GetBaseModel(handle.Db, id)
		if err == ModelNotFoundError {
			return ErrorCode(nil, http.StatusNotFound, c.AddMap(AnyMap{
				"NotFoundMessage": "Model not found",
				"GoBackLink":      "/models",
			}))
		} else if err != nil {
			// TODO improve this response
			return Error500(err)
		}

		if model.Status != CONFIRM_PRE_TRAINING {
			// TODO improve this response
			return ErrorCode(nil, 400, c.AddMap(nil))
		}

		cls, err := model_classes.ListClasses(handle.Db, model.Id)
		if err != nil {
			ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
			// TODO improve this response
			return c.Error500(err)
		}

		err = removeFailedDataPoints(c.Db, model)
		if err != nil {
			return c.Error500(err)
		}

		var fid string
		for i := 0; i < number_of_models; i++ {
			def_id, err := MakeDefenition(handle.Db, model.Id, accuracy)
			if err != nil {
				ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
				// TODO improve this response
				return c.Error500(err)
			}

			if fid == "" {
				fid = def_id
			}

			// TODO change shape of it depends on the type of the image
			err = MakeLayer(handle.Db, def_id, 1, LAYER_INPUT, fmt.Sprintf("%d,%d,1", model.Width, model.Height))
			if err != nil {
				ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
				// TODO improve this response
				return c.Error500(err)
			}
			err = MakeLayer(handle.Db, def_id, 4, LAYER_FLATTEN, fmt.Sprintf("%d,1", len(cls)))
			if err != nil {
				ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
				// TODO improve this response
				return c.Error500(err)
			}
			err = MakeLayer(handle.Db, def_id, 5, LAYER_DENSE, fmt.Sprintf("%d,1", len(cls)*3))
			if err != nil {
				ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
				// TODO improve this response
				return c.Error500(err)
			}
			// Using sparce
			err = MakeLayer(handle.Db, def_id, 5, LAYER_DENSE, fmt.Sprintf("%d, 1", len(cls)))
			if err != nil {
				ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
				// TODO improve this response
				return c.Error500(err)
			}

			err = ModelDefinitionUpdateStatus(c, def_id, MODEL_DEFINITION_STATUS_INIT)
			if err != nil {
				ModelUpdateStatus(c, model.Id, FAILED_PREPARING_TRAINING)
				// TODO improve this response
				return c.Error500(err)
			}
		}

		// TODO start training with id fid
		go trainModel(c, model)

		ModelUpdateStatus(c, model.Id, TRAINING)
		Redirect("/models/edit?id="+model.Id, c.Mode, w, r)
		return nil
	})
}