Haystack/backend/main.go

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"screenmark/screenmark/models"
	"time"

	"github.com/joho/godotenv"
	"github.com/lib/pq"
)

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

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

	listener := pq.NewListener(os.Getenv("DB_CONNECTION"), time.Second, time.Second, func(event pq.ListenerEventType, err error) {
		if err != nil {
			panic(err)
		}
	})
	defer listener.Close()

	go func() {
		err := listener.Listen("new_image")
		if err != nil {
			panic(err)
		}

		select {
		case parameters := <-listener.Notify:
			log.Println("received notification, new image available: " + parameters.Extra)

			go func() {
				openAiClient, err := CreateOpenAiClient()
				if err != nil {
					panic(err)
				}

				image, err := models.GetImage(parameters.Extra)
				if err != nil {
					log.Println(err)
					return
				}

				imageInfo, err := openAiClient.GetImageInfo(image.ImageName, image.Image)
				if err != nil {
					log.Println(err)
					return
				}

				log.Printf("Info: %+v\n", imageInfo)
			}()
		}
	}()

	mux := http.NewServeMux()

	mux.HandleFunc("OPTIONS /image/{name}", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Add("Access-Control-Allow-Origin", "*")
		w.Header().Add("Access-Control-Allow-Credentials", "*")
		w.Header().Add("Access-Control-Allow-Headers", "*")
	})

	mux.HandleFunc("GET /image", func(w http.ResponseWriter, r *http.Request) {
		userId := r.Header.Get("userId")

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

		jsonImages, err := json.Marshal(images)
		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)
	})

	mux.HandleFunc("GET /image/{id}", func(w http.ResponseWriter, r *http.Request) {
		imageId := r.PathValue("id")

		// TODO: really need authorization here!
		image, err := models.GetImage(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)
	})

	mux.HandleFunc("POST /image/{name}", func(w http.ResponseWriter, r *http.Request) {
		imageName := r.PathValue("name")

		userId := r.Header.Get("userId")

		w.Header().Add("Access-Control-Allow-Origin", "*")
		w.Header().Add("Access-Control-Allow-Credentials", "*")
		w.Header().Add("Access-Control-Allow-Headers", "*")

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

		image, err := io.ReadAll(r.Body)

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

		userImage, err := models.SaveImage(userId, imageName, image)
		if err != nil {
			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(err)
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "Could not create JSON response for this image")
			return
		}

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

		w.WriteHeader(http.StatusCreated)
	})

	log.Println("Listening and serving on port 3040.")

	http.ListenAndServe(":3040", mux)
}