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" "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.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 } w.WriteHeader(http.StatusOK) w.Header().Add("Content-Type", "application/json") fmt.Fprint(w, string(json)) }) log.Println("Listening and serving on port 3040.") if err := http.ListenAndServe(":3040", r); err != nil { log.Println(err) return } }