From 28dd02a47d943b4c70ce87a2197b0c04c70e1f60 Mon Sep 17 00:00:00 2001
From: John Costa <johncosta027@gmail.com>
Date: Tue, 18 Mar 2025 18:18:01 +0000
Subject: [PATCH] feat(locations): now working

---
 backend/main.go             | 15 +++++++++
 backend/models/locations.go | 61 +++++++++++++++++++++++++++++++++----
 backend/models/tags.go      |  1 +
 backend/openai.go           |  3 ++
 backend/schema.sql          |  2 +-
 5 files changed, 75 insertions(+), 7 deletions(-)

diff --git a/backend/main.go b/backend/main.go
index 3e5ee87..7d173b9 100644
--- a/backend/main.go
+++ b/backend/main.go
@@ -33,11 +33,18 @@ func (client TestAiClient) GetImageInfo(imageName string, imageData []byte) (Ima
 func GetAiClient() (AiClient, error) {
 	mode := os.Getenv("MODE")
 	if mode == "TESTING" {
+		address := "10 Downing Street"
+
 		return TestAiClient{
 			ImageInfo: ImageInfo{
 				Tags:  []string{"tag"},
 				Links: []string{"links"},
 				Text:  []string{"text"},
+				Locations: []model.Locations{{
+					ID:      uuid.Nil,
+					Name:    "London",
+					Address: &address,
+				}},
 			},
 		}, nil
 	}
@@ -63,6 +70,7 @@ func main() {
 	linkModel := models.NewLinkModel(db)
 	tagModel := models.NewTagModel(db)
 	textModel := models.NewTextModel(db)
+	locationModel := models.NewLocationModel(db)
 	userModel := models.NewUserModel(db)
 
 	listener := pq.NewListener(os.Getenv("DB_CONNECTION"), time.Second, time.Second, func(event pq.ListenerEventType, err error) {
@@ -135,6 +143,13 @@ func main() {
 						log.Println(err)
 						return
 					}
+
+					err = locationModel.SaveToImage(ctx, userImage.ImageID, imageInfo.Locations)
+					if err != nil {
+						log.Println("Failed to save location")
+						log.Println(err)
+						return
+					}
 				}()
 			}
 		}
diff --git a/backend/models/locations.go b/backend/models/locations.go
index fbac5dc..2cfcde4 100644
--- a/backend/models/locations.go
+++ b/backend/models/locations.go
@@ -3,27 +3,76 @@ package models
 import (
 	"context"
 	"database/sql"
+	"log"
 	"screenmark/screenmark/.gen/haystack/haystack/model"
 	. "screenmark/screenmark/.gen/haystack/haystack/table"
+
+	"github.com/google/uuid"
 )
 
 type LocationModel struct {
 	dbPool *sql.DB
 }
 
-func (m LocationModel) Save(ctx context.Context, location model.Locations) (model.Locations, error) {
+// This looks stupid
+func getValues(location model.Locations) []any {
+	arr := make([]any, 0)
+
+	if location.Address != nil {
+		arr = append(arr, *location.Address)
+	} else {
+		arr = append(arr, nil)
+	}
+
+	if location.Coordinates != nil {
+		arr = append(arr, *location.Coordinates)
+	} else {
+		arr = append(arr, nil)
+	}
+
+	if location.Description != nil {
+		arr = append(arr, *location.Description)
+	} else {
+		arr = append(arr, nil)
+	}
+
+	return arr
+}
+
+func (m LocationModel) Save(ctx context.Context, locations []model.Locations) (model.Locations, error) {
 	insertLocationStmt := Locations.
-		INSERT(Locations.Name, Locations.Address, Locations.Coordinates, Locations.Description).
-		VALUES(location.Name, location.Address, location.Coordinates, location.Description).
-		RETURNING(Locations.AllColumns)
+		INSERT(Locations.Name, Locations.Address, Locations.Coordinates, Locations.Description)
+
+	for _, location := range locations {
+		insertLocationStmt = insertLocationStmt.VALUES(location.Name, getValues(location)...)
+	}
+
+	insertLocationStmt = insertLocationStmt.RETURNING(Locations.AllColumns)
+
+	log.Println(insertLocationStmt.DebugSql())
 
 	insertedLocation := model.Locations{}
-
-	err := insertLocationStmt.QueryContext(ctx, m.dbPool, &insertLocationStmt)
+	err := insertLocationStmt.QueryContext(ctx, m.dbPool, &insertedLocation)
 
 	return insertedLocation, err
 }
 
+func (m LocationModel) SaveToImage(ctx context.Context, imageId uuid.UUID, locations []model.Locations) error {
+	location, err := m.Save(ctx, locations)
+
+	if err != nil {
+		return err
+	}
+
+	insertImageLocationStmt := ImageLocations.
+		INSERT(ImageLocations.ImageID, ImageLocations.LocationID).
+		VALUES(imageId, location.ID)
+
+	_, err = insertImageLocationStmt.ExecContext(ctx, m.dbPool)
+
+	return err
+}
+
 func NewLocationModel(db *sql.DB) LocationModel {
 	return LocationModel{dbPool: db}
 }
diff --git a/backend/models/tags.go b/backend/models/tags.go
index a6d3a93..4d32f95 100644
--- a/backend/models/tags.go
+++ b/backend/models/tags.go
@@ -8,6 +8,7 @@ import (
 	. "screenmark/screenmark/.gen/haystack/haystack/table"
 
 	. "github.com/go-jet/jet/v2/postgres"
+
 	"github.com/google/uuid"
 )
 
diff --git a/backend/openai.go b/backend/openai.go
index 4a90a34..cca416f 100644
--- a/backend/openai.go
+++ b/backend/openai.go
@@ -11,12 +11,15 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"screenmark/screenmark/.gen/haystack/haystack/model"
 )
 
 type ImageInfo struct {
 	Tags  []string `json:"tags"`
 	Text  []string `json:"text"`
 	Links []string `json:"links"`
+
+	Locations []model.Locations `json:"locations"`
 }
 
 type ResponseFormat struct {
diff --git a/backend/schema.sql b/backend/schema.sql
index a232e5e..1bd7ef3 100644
--- a/backend/schema.sql
+++ b/backend/schema.sql
@@ -63,7 +63,7 @@ CREATE TABLE haystack.locations (
 
 CREATE TABLE haystack.image_locations (
 	id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
-    location_id UUID NOT NULL REFERENCES haystack.image_locations (id),
+    location_id UUID NOT NULL REFERENCES haystack.locations (id),
     image_id UUID NOT NULL REFERENCES haystack.image (id)
 );