Haystack/backend/events.go

package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"net/http"
	"os"
	"screenmark/screenmark/agents"
	"screenmark/screenmark/middleware"
	"screenmark/screenmark/models"
	"strconv"
	"sync"
	"time"

	"github.com/charmbracelet/log"
	"github.com/google/uuid"
	"github.com/lib/pq"
)

const (
	IMAGE_TYPE = "image"
	LIST_TYPE  = "list"
)

type imageNotification struct {
	Type string

	ImageID   uuid.UUID
	ImageName string

	Status string
}

type listNotification struct {
	Type string

	ListID uuid.UUID
	Name   string

	Status string
}

type Notification struct {
	image *imageNotification
	list  *listNotification
}

func getImageNotification(image imageNotification) Notification {
	return Notification{
		image: &image,
	}
}

func getListNotification(list listNotification) Notification {
	return Notification{
		list: &list,
	}
}

func (n Notification) MarshalJSON() ([]byte, error) {
	if n.image != nil {
		return json.Marshal(n.image)
	}

	if n.list != nil {
		return json.Marshal(n.list)
	}

	return nil, fmt.Errorf("no image or list present")
}

func (n *Notification) UnmarshalJSON(data []byte) error {
	return fmt.Errorf("unimplemented")
}

func ListenNewImageEvents(db *sql.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()

	imageModel := models.NewImageModel(db)
	listModel := models.NewListModel(db)

	databaseEventLog := createLogger("Database Events 🤖", os.Stdout)
	databaseEventLog.SetLevel(log.DebugLevel)

	err := listener.Listen("new_image")
	if err != nil {
		panic(err)
	}

	for parameters := range listener.Notify {
		imageId := uuid.MustParse(parameters.Extra)

		databaseEventLog.Debug("Starting processing image", "ImageID", imageId)

		ctx := context.Background()

		go func() {
			image, err := imageModel.GetToProcessWithData(ctx, imageId)
			if err != nil {
				databaseEventLog.Error("Failed to GetToProcessWithData", "error", err)
				return
			}

			splitWriter := createDbStdoutWriter(db, image.ImageID)

			if err := imageModel.StartProcessing(ctx, image.ID); err != nil {
				databaseEventLog.Error("Failed to FinishProcessing", "error", err)
				return
			}

			descriptionAgent := agents.NewDescriptionAgent(createLogger("Description 📝", splitWriter), imageModel)
			listAgent := agents.NewListAgent(createLogger("Lists 🖋️", splitWriter), listModel)

			var wg sync.WaitGroup
			wg.Add(2)

			go func() {
				defer wg.Done()

				descriptionAgent.Describe(createLogger("Description 📓", splitWriter), image.Image.ID, image.Image.ImageName, image.Image.Image)
			}()

			go func() {
				defer wg.Done()

				listAgent.RunAgent(image.UserID, image.ImageID, image.Image.ImageName, image.Image.Image)
			}()

			wg.Wait()

			_, err = imageModel.FinishProcessing(ctx, image.ID)
			if err != nil {
				databaseEventLog.Error("Failed to finish processing", "ImageID", imageId, "error", err)
				return
			}

			databaseEventLog.Debug("Finished processing image", "ImageID", imageId)
		}()
	}
}

func ListenProcessingImageStatus(db *sql.DB, images models.ImageModel, notifier *Notifier[Notification]) {
	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()

	logger := createLogger("Image Status 📊", os.Stdout)

	if err := listener.Listen("new_processing_image_status"); err != nil {
		panic(err)
	}

	for data := range listener.Notify {
		imageStringUuid := data.Extra[0:36]
		status := data.Extra[36:]

		imageUuid, err := uuid.Parse(imageStringUuid)
		if err != nil {
			logger.Error(err)
			continue
		}

		processingImage, err := images.GetToProcess(context.Background(), imageUuid)
		if err != nil {
			logger.Error("GetToProcess failed", "err", err)
			continue
		}

		logger.Info("Update", "id", imageStringUuid, "status", status)

		notification := getImageNotification(imageNotification{
			Type:      IMAGE_TYPE,
			ImageID:   processingImage.ImageID,
			ImageName: processingImage.Image.ImageName,
			Status:    status,
		})

		if err := notifier.SendAndCreate(processingImage.UserID.String(), notification); err != nil {
			logger.Error(err)
		}
	}
}

func ListenNewStackEvents(db *sql.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()

	stackModel := models.NewListModel(db)

	newStacksLogger := createLogger("New Stacks 🤖", os.Stdout)
	newStacksLogger.SetLevel(log.DebugLevel)

	err := listener.Listen("new_stack")
	if err != nil {
		panic(err)
	}

	for parameters := range listener.Notify {
		stackID := uuid.MustParse(parameters.Extra)

		newStacksLogger.Debug("Starting processing stack", "StackID", stackID)

		ctx := context.Background()

		go func() {
			stack, err := stackModel.GetProcessing(ctx, stackID)
			if err != nil {
				newStacksLogger.Error("failed to get processing", "error", err)
				return
			}

			if err := stackModel.StartProcessing(ctx, stackID); err != nil {
				newStacksLogger.Error("failed to start processing", "error", err)
				return
			}

			listAgent := agents.NewCreateListAgent(newStacksLogger, stackModel)
			userListRequest := fmt.Sprintf("title=%s,fields=%s", stack.Title, stack.Fields)

			err = listAgent.CreateList(newStacksLogger, stack.UserID, userListRequest)
			if err != nil {
				newStacksLogger.Error("running agent", "err", err)
				return
			}

			newStacksLogger.Debug("Finished processing stack", "StackID", stackID)
		}()
	}
}

func ListenProcessingStackStatus(db *sql.DB, stacks models.ListModel, notifier *Notifier[Notification]) {
	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()

	logger := createLogger("Stack Status 📊", os.Stdout)

	if err := listener.Listen("new_processing_stack_status"); err != nil {
		panic(err)
	}

	for data := range listener.Notify {
		stackStringUUID := data.Extra[0:36]
		status := data.Extra[36:]

		stackUUID, err := uuid.Parse(stackStringUUID)
		if err != nil {
			logger.Error(err)
			continue
		}

		processingStack, err := stacks.GetToProcess(context.Background(), stackUUID)
		if err != nil {
			logger.Error("GetToProcess failed", "err", err)
			continue
		}

		logger.Info("Update", "id", stackStringUUID, "status", status)

		notification := getListNotification(listNotification{
			Type:   LIST_TYPE,
			Name:   processingStack.Title,
			ListID: stackUUID,
			Status: status,
		})

		if err := notifier.SendAndCreate(processingStack.UserID.String(), notification); err != nil {
			logger.Error(err)
		}
	}
}

/*
 * TODO: We have channels open every a user sends an image.
 * We never close these channels.
 *
 * What is a reasonable default? Close the channel after 1 minute of inactivity?
 */
func CreateEventsHandler(notifier *Notifier[Notification]) http.HandlerFunc {
	counter := 0

	userSplitters := make(map[string]*ChannelSplitter[Notification])

	return func(w http.ResponseWriter, r *http.Request) {
		_userId := r.Context().Value(middleware.USER_ID).(uuid.UUID)
		if _userId == uuid.Nil {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}

		userId := _userId.String()

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

		if _, exists := notifier.Listeners[userId]; !exists {
			notifier.Create(userId)
		}

		userNotifications := notifier.Listeners[userId]

		if _, exists := userSplitters[userId]; !exists {
			splitter := NewChannelSplitter(userNotifications)

			userSplitters[userId] = &splitter
			splitter.Listen()
		}

		splitter := userSplitters[userId]

		id := strconv.Itoa(counter)
		counter += 1

		notifications := splitter.Add(id)
		defer splitter.Remove(id)

		for {
			select {
			case <-r.Context().Done():
				fmt.Fprint(w, "event: close\ndata: Connection closed\n\n")
				w.(http.Flusher).Flush()
				return
			case msg := <-notifications:

				msgString, err := json.Marshal(msg)
				if err != nil {
					w.WriteHeader(http.StatusInternalServerError)
					return
				}

				fmt.Printf("Sending msg %s\n", msgString)

				fmt.Fprintf(w, "event: data\ndata: %s\n\n", string(msgString))
				w.(http.Flusher).Flush()
			}
		}
	}
}