Jade/Chat.go

package main

import (
	"context"
	"fmt"
	"net/http"
	"sort"
	"time"

	"github.com/flosch/pongo2"
	"github.com/gofiber/fiber/v2"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
)

type Message struct {
	ID               primitive.ObjectID `bson:"_id"`
	Content          string             `bson:"message"`
	Role             string             `bson:"role"`
	Date             time.Time          `bson:"date"`
	Model            string             `bson:"model"`
	Selected         bool               `bson:"selected"`
	LinkedMessageIds []string           `bson:"linked_message_ids"`
}

type Conversation struct {
	ID       string
	Messages []Message
}

func ChatPageHandler(c *fiber.Ctx) error {
	return c.Render("chat", fiber.Map{}, "layouts/main")
}

func LoadModelSelectionHandler(c *fiber.Ctx) error {
	CheckedModels := []string{"gpt-3.5-turbo"} // Default model
	out, err := popoverTmpl.Execute(pongo2.Context{
		"CompanyInfos":  CompanyInfos,
		"CheckedModels": CheckedModels,
	})
	if err != nil {
		c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
			"error": "Error rendering template",
		})
	}
	return c.SendString(out)
}

func DeleteMessageHandler(c *fiber.Ctx) error {
	messageId := c.FormValue("id")

	// messageId is a string that look like ObjectID("662d6556e5328bbc8357cfb1")
	// Get the message from the db and delete it and all messages after it
	collection := mongoClient.Database("chat").Collection("messages")

	// Convert the messageId string to a bson.ObjectId
	objectID, err := primitive.ObjectIDFromHex(messageId)
	if err != nil {
		return c.Status(http.StatusBadRequest).JSON(fiber.Map{
			"error": "Invalid message ID",
		})
	}

	// Find the message with the given ID
	message := &Message{}
	err = collection.FindOne(c.Context(), bson.M{"_id": objectID}).Decode(message)
	if err != nil {
		if err == mongo.ErrNoDocuments {
			return c.Status(http.StatusNotFound).JSON(fiber.Map{
				"error": "Message not found",
			})
		}
		return c.Status(http.StatusInternalServerError).JSON(fiber.Map{
			"error": "Error finding message",
		})
	}

	// Delete the message and all messages after it
	_, err = collection.DeleteMany(c.Context(), bson.M{
		"date": bson.M{"$gte": message.Date},
	})
	if err != nil {
		return c.Status(http.StatusInternalServerError).JSON(fiber.Map{
			"error": "Error deleting messages",
		})
	}

	return c.SendString(generateChatHTML())
}

func LoadChatHandler(c *fiber.Ctx) error {
	return c.SendString(generateChatHTML())
}

func ChangeSelectedMessageHandler(c *fiber.Ctx) error {
	messageId := c.FormValue("id")
	collection := mongoClient.Database("chat").Collection("messages")
	objectID, err := primitive.ObjectIDFromHex(messageId)
	if err != nil {
		return c.Status(http.StatusBadRequest).JSON(fiber.Map{
			"error": "Invalid message ID",
		})
	}

	var selectedMessage Message
	err = collection.FindOne(c.Context(), bson.M{"_id": objectID}).Decode(&selectedMessage)
	if err != nil {
		return c.Status(http.StatusInternalServerError).JSON(fiber.Map{
			"error": "Error finding message",
		})
	}

	for i := range selectedMessage.LinkedMessageIds {
		objectID, _ = primitive.ObjectIDFromHex(selectedMessage.LinkedMessageIds[i])
		_, err = collection.UpdateOne(c.Context(), bson.M{"_id": objectID}, bson.M{"$set": bson.M{"selected": false}})
		if err != nil {
			return c.Status(http.StatusInternalServerError).JSON(fiber.Map{
				"error": "Error updating message",
			})
		}
	}

	_, err = collection.UpdateOne(c.Context(), bson.M{"_id": objectID}, bson.M{"$set": bson.M{"selected": true}})
	if err != nil {
		return c.Status(http.StatusInternalServerError).JSON(fiber.Map{
			"error": "Error updating message",
		})
	}

	return c.SendString(generateChatHTML())
}

type NextMessage struct {
	Icon    string
	Content string
	Hidden  bool
	Id      string
}

func generateChatHTML() string {
	// Get the messages from the database
	collection := mongoClient.Database("chat").Collection("messages")

	// Get all messages
	cursor, err := collection.Find(context.TODO(), bson.D{})
	if err != nil {
		fmt.Println(err)
	}

	// Convert the cursor to an array of messages
	var Messages []Message
	if err = cursor.All(context.TODO(), &Messages); err != nil {
		fmt.Println(err)
	}

	htmlString := "<div class='columns is-centered' id='chat-container'><div class='column is-12-mobile is-8-tablet is-6-desktop' id='chat-messages'>"

	var NextMessages []NextMessage // Use the custom NextMessage struct

	for i, message := range Messages {
		if message.Role == "user" {
			htmlContent := markdownToHTML(message.Content)
			userOut, err := userTmpl.Execute(pongo2.Context{"Content": htmlContent, "ID": message.ID.Hex()})
			if err != nil {
				panic(err)
			}
			htmlString += userOut
			// Reset NextMessages when a user message is encountered
			NextMessages = []NextMessage{}
		} else {
			fmt.Println(i)
			// For bot messages, add them to NextMessages with only the needed fields
			nextMsg := NextMessage{
				Icon:    model2Icon(message.Model), // Assuming Icon is a field you want to include from Message
				Content: markdownToHTML(message.Content),
				Hidden:  !message.Selected, // Assuming Hidden is a field you want to include from Message
				Id:      message.ID.Hex(),
			}
			NextMessages = append(NextMessages, nextMsg)

			// Check if the next message is not a bot or if it's the last message
			if i+1 == len(Messages) || Messages[i+1].Role != "bot" {
				sort.Slice(NextMessages, func(i, j int) bool {
					if !NextMessages[i].Hidden && NextMessages[j].Hidden {
						return true
					}
					if NextMessages[i].Hidden && !NextMessages[j].Hidden {
						return false
					}
					return true
				})

				botOut, err := botTmpl.Execute(pongo2.Context{"Messages": NextMessages, "ConversationAreaId": i})
				if err != nil {
					panic(err)
				}
				htmlString += botOut
				htmlString += "<div style='height: 10px;'></div>"
			}
		}
	}

	htmlString += "</div></div>"

	// Render the HTML template with the messages
	return htmlString
}

func GetMessageContentHandler(c *fiber.Ctx) error {
	messageId := c.FormValue("id")
	fmt.Println(messageId)
	collection := mongoClient.Database("chat").Collection("messages")
	objectID, err := primitive.ObjectIDFromHex(messageId)
	if err != nil {
		return c.Status(http.StatusBadRequest).JSON(fiber.Map{
			"error": "Invalid message ID",
		})
	}

	var selectedMessage Message
	err = collection.FindOne(c.Context(), bson.M{"_id": objectID}).Decode(&selectedMessage)
	if err != nil {
		return c.Status(http.StatusInternalServerError).JSON(fiber.Map{
			"error": "Error finding message",
		})
	}

	fmt.Println(selectedMessage)

	return c.SendString(selectedMessage.Content)
}