Haystack/backend/main.go

package main

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"path/filepath"
	"screenmark/screenmark/.gen/haystack/haystack/model"
	"screenmark/screenmark/agents/client"
	"screenmark/screenmark/models"

	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
	"github.com/google/uuid"
	"github.com/joho/godotenv"
)

type TestAiClient struct {
	ImageInfo client.ImageMessageContent
}

func (client TestAiClient) GetImageInfo(imageName string, imageData []byte) (client.ImageMessageContent, error) {
	return client.ImageInfo, nil
}

func main() {
	err := godotenv.Load()
	if err != nil {
		panic(err)
	}

	db, err := models.InitDatabase()
	if err != nil {
		panic(err)
	}

	imageModel := models.NewImageModel(db)
	userModel := models.NewUserModel(db)

	mail, err := CreateMailClient()
	if err != nil {
		panic(err)
	}

	auth := CreateAuth(mail)

	eventManager := NewEventManager()

	go ListenNewImageEvents(db, &eventManager)
	go ListenProcessingImageStatus(db, &eventManager)

	r := chi.NewRouter()

	r.Use(middleware.Logger)
	r.Use(CorsMiddleware)
	r.Options("/*", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	// Temporarily not in protect route because we aren't using cookies.
	// Therefore they don't get automatically attached to the request.
	// So <img src=""> cannot send the tokensend the token
	r.Get("/image/{id}", func(w http.ResponseWriter, r *http.Request) {
		stringImageId := r.PathValue("id")
		// userId := r.Context().Value(USER_ID).(uuid.UUID)

		imageId, err := uuid.Parse(stringImageId)
		if err != nil {
			w.WriteHeader(http.StatusForbidden)
			fmt.Fprintf(w, "You cannot read this")
			return
		}

		// if authorized := imageModel.IsUserAuthorized(r.Context(), imageId, userId); !authorized {
		// 	w.WriteHeader(http.StatusForbidden)
		// 	fmt.Fprintf(w, "You cannot read this")
		// 	return
		// }

		image, err := imageModel.Get(r.Context(), imageId)
		if err != nil {
			log.Println(err)
			w.WriteHeader(http.StatusNotFound)
			fmt.Fprintf(w, "Could not get image")
			return
		}

		// TODO: this could be part of the db table
		extension := filepath.Ext(image.ImageName)
		extension = extension[1:]

		w.Header().Add("Content-Type", "image/"+extension)
		w.Write(image.Image)
	})

	r.Group(func(r chi.Router) {
		r.Use(ProtectedRoute)
		r.Use(func(next http.Handler) http.Handler {
			return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Add("Content-Type", "application/json")

				next.ServeHTTP(w, r)
			})
		})

		r.Get("/image", func(w http.ResponseWriter, r *http.Request) {
			userId := r.Context().Value(USER_ID).(uuid.UUID)
			if err != nil {
				w.WriteHeader(http.StatusForbidden)
				fmt.Fprintf(w, "You cannot read this")
				return
			}

			imageProperties, err := userModel.ListWithProperties(r.Context(), userId)
			if err != nil {
				log.Println(err)
				w.WriteHeader(http.StatusNotFound)
				fmt.Fprintf(w, "Something went wrong")
				return
			}

			images, err := userModel.GetUserImages(r.Context(), userId)
			if err != nil {
				log.Println(err)
				w.WriteHeader(http.StatusNotFound)
				fmt.Fprintf(w, "Something went wrong")
				return
			}

			type ImagesReturn struct {
				UserImages      []models.UserImageWithImage
				ImageProperties []models.TypedProperties
			}

			imagesReturn := ImagesReturn{
				UserImages:      images,
				ImageProperties: models.GetTypedImageProperties(imageProperties),
			}

			jsonImages, err := json.Marshal(imagesReturn)
			if err != nil {
				log.Println(err)
				w.WriteHeader(http.StatusBadRequest)
				fmt.Fprintf(w, "Could not create JSON response for this image")
				return
			}

			w.Write(jsonImages)
		})

		r.Get("/image-properties/{id}", func(w http.ResponseWriter, r *http.Request) {
			userId := r.Context().Value(USER_ID).(uuid.UUID)
			stringImageId := r.PathValue("id")

			imageId, err := uuid.Parse(stringImageId)
			if err != nil {
				w.WriteHeader(http.StatusForbidden)
				fmt.Fprintf(w, "You cannot read this")
				return
			}

			image, err := userModel.ListImageWithProperties(r.Context(), userId, imageId)
			if err != nil {
				log.Println(err)
				w.WriteHeader(http.StatusNotFound)
				fmt.Fprintf(w, "Something went wrong")
				return
			}

			jsonImages, err := json.Marshal(models.GetTypedImageProperties([]models.ImageWithProperties{image})[0])
			if err != nil {
				log.Println(err)
				w.WriteHeader(http.StatusBadRequest)
				fmt.Fprintf(w, "Could not create JSON response for this image")
				return
			}

			w.Write(jsonImages)
		})

		r.Post("/image/{name}", func(w http.ResponseWriter, r *http.Request) {
			imageName := r.PathValue("name")
			userId := r.Context().Value(USER_ID).(uuid.UUID)

			if len(imageName) == 0 {
				w.WriteHeader(http.StatusBadRequest)
				fmt.Fprintf(w, "You need to provide a name in the path")
				return
			}

			contentType := r.Header.Get("Content-Type")

			// TODO: length checks on body
			// TODO: extract this shit out
			image := make([]byte, 0)
			if contentType == "application/base64" {
				decoder := base64.NewDecoder(base64.StdEncoding, r.Body)
				buf := &bytes.Buffer{}

				decodedIamge, err := io.Copy(buf, decoder)
				if err != nil {
					log.Println(err)
					w.WriteHeader(http.StatusBadRequest)
					fmt.Fprintf(w, "bruh, base64 aint decoding")
					return
				}

				fmt.Println(string(image))
				fmt.Println(decodedIamge)

				image = buf.Bytes()
			} else if contentType == "application/oclet-stream" {
				bodyData, err := io.ReadAll(r.Body)
				if err != nil {
					log.Println(err)
					w.WriteHeader(http.StatusBadRequest)
					fmt.Fprintf(w, "bruh, binary aint binaring")
					return
				}
				// TODO: check headers

				image = bodyData
			} else {
				log.Println("bad stuff?")
				w.WriteHeader(http.StatusBadRequest)
				fmt.Fprintf(w, "Bruh, you need oclet stream or base64")
				return
			}

			if err != nil {
				log.Println("First case")
				w.WriteHeader(http.StatusBadRequest)
				fmt.Fprintf(w, "Couldnt read the image from the request body")
				return
			}

			userImage, err := imageModel.Process(r.Context(), userId, model.Image{
				Image:     image,
				ImageName: imageName,
			})
			if err != nil {
				log.Println("Second case")
				log.Println(err)
				w.WriteHeader(http.StatusBadRequest)
				fmt.Fprintf(w, "Could not save image to DB")
				return
			}

			jsonUserImage, err := json.Marshal(userImage)
			if err != nil {
				log.Println("Third case")
				log.Println(err)
				w.WriteHeader(http.StatusBadRequest)
				fmt.Fprintf(w, "Could not create JSON response for this image")
				return
			}

			w.WriteHeader(http.StatusCreated)

			fmt.Fprint(w, string(jsonUserImage))
			w.Header().Add("Content-Type", "application/json")
		})

	})

	r.Get("/image-events/{id}", func(w http.ResponseWriter, r *http.Request) {
		// TODO: authentication :)

		id := r.PathValue("id")

		// TODO: get the current status of the image and send it across.
		ctx, cancel := context.WithCancel(r.Context())

		imageNotifier, exists := eventManager.listeners[id]
		if !exists {
			fmt.Println("Not found!")
			w.WriteHeader(http.StatusNotFound)
			w.(http.Flusher).Flush()
			cancel()
			return
		}

		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.Header().Set("Connection", "keep-alive")
		w.(http.Flusher).Flush()

		for {
			select {
			case <-ctx.Done():
				fmt.Fprint(w, "event: close\ndata: Connection closed\n\n")
				w.(http.Flusher).Flush()
				cancel()
				return
			case data := <-imageNotifier:
				if data == "" {
					cancel()
					continue
				}

				fmt.Printf("Status received: %s\n", data)
				fmt.Fprintf(w, "event: data\ndata: %s\n\n", data)
				w.(http.Flusher).Flush()

				if data == "complete" {
					cancel()
				}
			}
		}
	})

	r.Post("/login", func(w http.ResponseWriter, r *http.Request) {
		type LoginBody struct {
			Email string `json:"email"`
		}

		loginBody := LoginBody{}
		err := json.NewDecoder(r.Body).Decode(&loginBody)
		if err != nil {
			log.Println(err)
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "Request body was not correct")
			return
		}

		// TODO: validate it's an email

		auth.CreateCode(loginBody.Email)

		w.WriteHeader(http.StatusOK)
	})

	r.Post("/code", func(w http.ResponseWriter, r *http.Request) {
		type CodeBody struct {
			Email string `json:"email"`
			Code  string `json:"code"`
		}

		type CodeReturn struct {
			Access  string `json:"access"`
			Refresh string `json:"refresh"`
		}

		codeBody := CodeBody{}
		if err := json.NewDecoder(r.Body).Decode(&codeBody); err != nil {
			log.Println(err)
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "Request body was not correct")
			return
		}

		if err := auth.UseCode(codeBody.Email, codeBody.Code); err != nil {
			log.Println(err)
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "email or code are incorrect")
			return
		}

		if exists := userModel.DoesUserExist(r.Context(), codeBody.Email); !exists {
			userModel.Save(r.Context(), model.Users{
				Email: codeBody.Email,
			})
		}

		uuid, err := userModel.GetUserIdFromEmail(r.Context(), codeBody.Email)
		if err != nil {
			log.Println(err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "Something went wrong.")
			return
		}

		refresh := CreateRefreshToken(uuid)
		access := CreateAccessToken(uuid)

		codeReturn := CodeReturn{
			Access:  access,
			Refresh: refresh,
		}

		fmt.Println(codeReturn)

		json, err := json.Marshal(codeReturn)
		if err != nil {
			log.Println(err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "Something went wrong.")
			return
		}

		w.WriteHeader(http.StatusOK)
		w.Header().Add("Content-Type", "application/json")

		fmt.Fprint(w, string(json))
	})

	logWriter := DatabaseWriter{
		dbPool: db,
	}

	r.Route("/logs", createLogHandler(&logWriter))

	log.Println("Listening and serving on port 3040.")
	if err := http.ListenAndServe(":3040", r); err != nil {
		log.Println(err)
		return
	}
}