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

	eventManager := NewEventManager()

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

	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.Group(func(r chi.Router) {
		r.Use(ProtectedRoute)

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

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

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

			dataTypes := make([]DataType, 0)

			// lord
			// forgive me
			idMap := make(map[uuid.UUID]bool)

			for _, image := range images {
				for _, location := range image.Locations {
					_, exists := idMap[location.ID]
					if exists {
						continue
					}
					dataTypes = append(dataTypes, DataType{
						Type: "location",
						Data: location,
					})

					idMap[location.ID] = true
				}

				for _, event := range image.Events {
					_, exists := idMap[event.ID]
					if exists {
						continue
					}
					dataTypes = append(dataTypes, DataType{
						Type: "event",
						Data: event,
					})

					idMap[event.ID] = true
				}

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

				for _, contact := range image.Contacts {
					_, exists := idMap[contact.ID]
					if exists {
						continue
					}

					dataTypes = append(dataTypes, DataType{
						Type: "contact",
						Data: contact,
					})
					idMap[contact.ID] = true
				}
			}

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

			// TODO: really need authorization here!
			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.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.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")

		imageNotifier, exists := eventManager.listeners[id]
		if !exists {
			fmt.Println("Not found!")
			w.WriteHeader(http.StatusNotFound)
			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()

		ctx, cancel := context.WithCancel(r.Context())

		for {
			select {
			case <-ctx.Done():
				fmt.Fprint(w, "event: close\ndata: Connection closed\n\n")
				w.(http.Flusher).Flush()
				cancel()
				return
			case data := <-imageNotifier:
				fmt.Printf("Status received: %s\n", data)
				fmt.Fprintf(w, "data: %s-%s\n", data, time.Now().String())
				w.(http.Flusher).Flush()
				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,
		}

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