Haystack/backend/main.go

package main

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

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

	go ListenNewImageEvents(db)

	r := chi.NewRouter()

	r.Use(middleware.Logger)
	r.Use(CorsMiddleware)
	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.Options("/*", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	r.Get("/image", func(w http.ResponseWriter, r *http.Request) {
		token := r.Header.Get("Authorization")[7:]

		fmt.Println(token)

		userId, err := GetUserIdFromAccess(token)
		if err != nil {
			log.Println(err)
			w.WriteHeader(http.StatusForbidden)
			fmt.Fprintf(w, "Get out of here!")
			return
		}

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

		log.Println(images)

		type DataType struct {
			Type string `json:"type"`
			Data any    `json:"data"`
		}

		dataTypes := make([]DataType, 0)
		for _, image := range images {
			for _, location := range image.Locations {
				dataTypes = append(dataTypes, DataType{
					Type: "location",
					Data: location,
				})
			}

			for _, event := range image.Events {
				dataTypes = append(dataTypes, DataType{
					Type: "event",
					Data: event,
				})
			}

			for _, note := range image.Notes {
				dataTypes = append(dataTypes, DataType{
					Type: "note",
					Data: note,
				})
			}
		}

		jsonImages, err := json.Marshal(dataTypes)
		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/{id}", func(w http.ResponseWriter, r *http.Request) {
		imageId := r.PathValue("id")

		// TODO: really need authorization here!
		image, err := imageModel.Get(r.Context(), uuid.MustParse(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.Image.ImageName)
		extension = extension[1:]

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

	r.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
		}

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

		log.Println(contentType)

		// 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(), uuid.MustParse(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.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
		}

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

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

		accessCookie := http.Cookie{
			Name:     "access",
			Value:    access,
			Expires:  time.Now().Add(1 * time.Hour),
			Secure:   false,
			HttpOnly: true,
		}

		refreshCookie := http.Cookie{
			Name:     "refresh",
			Value:    access,
			Expires:  time.Now().Add(24 * 30 * time.Hour),
			Secure:   false,
			HttpOnly: true,
			Path:     "/refresh",
		}

		http.SetCookie(w, &accessCookie)  // TODO
		http.SetCookie(w, &refreshCookie) // TODO

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

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