package main import ( "bytes" "context" "encoding/base64" "encoding/json" "fmt" "io" "log" "net/http" "os" "path/filepath" "screenmark/screenmark/.gen/haystack/haystack/model" "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" "github.com/lib/pq" ) type TestAiClient struct { ImageInfo ImageInfo } func (client TestAiClient) GetImageInfo(imageName string, imageData []byte) (ImageInfo, error) { return client.ImageInfo, nil } func GetAiClient() (AiClient, error) { mode := os.Getenv("MODE") if mode == "TESTING" { return TestAiClient{ ImageInfo: ImageInfo{ Tags: []string{"tag"}, Links: []string{"links"}, Text: []string{"text"}, }, }, nil } return CreateOpenAiClient() } func main() { err := godotenv.Load() if err != nil { panic(err) } mode := os.Getenv("MODE") log.Printf("Mode: %s\n", mode) db, err := models.InitDatabase() if err != nil { panic(err) } imageModel := models.NewImageModel(db) linkModel := models.NewLinkModel(db) tagModel := models.NewTagModel(db) textModel := models.NewTextModel(db) 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) } for { select { case parameters := <-listener.Notify: imageId := uuid.MustParse(parameters.Extra) ctx := context.Background() go func() { openAiClient, err := GetAiClient() if err != nil { panic(err) } image, err := imageModel.GetToProcessWithData(ctx, imageId) if err != nil { return } imageInfo, err := openAiClient.GetImageInfo(image.Image.ImageName, image.Image.Image) if err != nil { return } savedImage, err := imageModel.FinishProcessing(ctx, image.ID) if err != nil { return } err = tagModel.SaveToImage(ctx, savedImage.ID, imageInfo.Tags) if err != nil { return } err = linkModel.Save(ctx, savedImage.ID, imageInfo.Links) if err != nil { return } err = textModel.Save(ctx, savedImage.ID, imageInfo.Text) if err != nil { return } }() } } }() r := chi.NewRouter() r.Use(middleware.Logger) 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.Header().Add("Access-Control-Allow-Origin", "*") w.Header().Add("Access-Control-Allow-Credentials", "*") w.Header().Add("Access-Control-Allow-Headers", "*") }) r.Get("/image", func(w http.ResponseWriter, r *http.Request) { 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", "*") images, err := imageModel.ListWithProperties(r.Context(), uuid.MustParse(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) }) r.Get("/image/{id}", func(w http.ResponseWriter, r *http.Request) { imageId := r.PathValue("id") w.Header().Add("Access-Control-Allow-Origin", "*") w.Header().Add("Access-Control-Allow-Credentials", "*") w.Header().Add("Access-Control-Allow-Headers", "*") // 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") }) log.Println("Listening and serving on port 3040.") http.ListenAndServe(":3040", r) }