package main import ( "context" "database/sql" "encoding/json" "fmt" "net/http" "os" "screenmark/screenmark/agents" "screenmark/screenmark/imageprocessor" "screenmark/screenmark/limits" "screenmark/screenmark/middleware" "screenmark/screenmark/models" "screenmark/screenmark/notifier" "strconv" "sync" "time" "github.com/charmbracelet/log" "github.com/google/uuid" "github.com/lib/pq" ) 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) limits := limits.CreateLimitsManager(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, limits) 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 getProcessingImageStatusChannel(db *sql.DB) chan string { 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() if err := listener.Listen("new_processing_image_status"); err != nil { panic(err) } msgChan := make(chan string) go func() { for data := range listener.Notify { msgChan <- data.Extra } }() return msgChan } 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 } if err := stackModel.EndProcessing(ctx, stackID); err != nil { newStacksLogger.Error("failed to finish processing", "error", err) return } newStacksLogger.Debug("Finished processing stack", "StackID", stackID) }() } } func ListenProcessingStackStatus(db *sql.DB, stacks models.ListModel, notifier *notifier.Notifier[imageprocessor.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 := imageprocessor.GetListNotification(imageprocessor.ListNotification{ Type: imageprocessor.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(notifierr *notifier.Notifier[imageprocessor.Notification]) http.HandlerFunc { counter := 0 userSplitters := make(map[string]*notifier.ChannelSplitter[imageprocessor.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 := notifierr.Listeners[userId]; !exists { notifierr.Create(userId) } userNotifications := notifierr.Listeners[userId] if _, exists := userSplitters[userId]; !exists { splitter := notifier.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() } } } }