Jade/main.go

package main

import (
	"bufio"
	"context"
	"fmt"
	"time"

	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/middleware/logger"
	"github.com/gofiber/template/django/v3"
	"github.com/valyala/fasthttp"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

var mongoClient *mongo.Client

type Message struct {
	ID      primitive.ObjectID `bson:"_id"`
	Content string             `bson:"message"`
	Role    string             `bson:"role"`
	Date    time.Time          `bson:"date"`
}

type Conversation struct {
	ID       string
	Messages []Message
}

type User struct {
	ID         string `bson:"_id"`
	Username   string `bson:"username"`
	OAth2Token string `bson:"oauth2token"`
	IsSub      bool   `bson:"isSub"`
}

type CurrentSession struct {
	ID                    string
	CurrentConversationID string
	CurrentUserID         string
}

func connectToMongoDB(uri string) {
	serverAPI := options.ServerAPI(options.ServerAPIVersion1)
	opts := options.Client().ApplyURI(uri).SetServerAPIOptions(serverAPI)
	client, err := mongo.Connect(context.TODO(), opts)
	if err != nil {
		panic(err)
	}
	mongoClient = client

	if err := mongoClient.Ping(context.TODO(), nil); err != nil {
		panic(err)
	}
}

func main() {
	// Import HTML using django engine/template
	engine := django.New("./views", ".html")

	if engine == nil {
		panic("Failed to create django engine")
	}

	// Create new Fiber instance. Can use any framework. I use fiber for speed and simplicity
	app := fiber.New(fiber.Config{
		Views:             engine,
		AppName:           "JADE 2.0",
		EnablePrintRoutes: true,
	})

	// Initialize
	go connectToMongoDB("mongodb://localhost:27017")

	// Add default logger
	app.Use(logger.New())

	// Add static files
	app.Static("/", "./static")

	// Add routes
	app.Get("/", indexHandler)
	app.Get("/isMongoDBConnected", isMongoDBConnectedHandler)

	app.Get("/chat", chatPageHandler)         // Complete chat page
	app.Put("/chat", addMessageHandler)       // Add message
	app.Delete("/chat", deleteMessageHandler) // Delete message
	app.Get("/loadChat", generateChatHTML)    // Load chat

	app.Get("/generateOpenai", addOpenaiMessage)

	app.Get("/sse", sseHandler) // SSE handler

	// Add test button
	app.Get("/test-button", testButtonHandler)

	// Start server
	app.Listen(":3000")
}

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

func chatPageHandler(c *fiber.Ctx) error {
	return c.Render("chat", fiber.Map{
		"Messages": []Message{},
	}, "layouts/main")
}

func testButtonHandler(c *fiber.Ctx) error {
	return c.Render("partials/test-button", fiber.Map{})
}

func addMessageHandler(c *fiber.Ctx) error {
	message := c.FormValue("message")
	lastMessageAsked = message

	return c.Render("partials/user-message", fiber.Map{
		"Message": Message{
			Content: message,
			Role:    "user",
			Date:    time.Now(),
		},
		"IncludePlaceholder": true,
	})
}

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

	// Convert the string ID to a MongoDB ObjectID
	objID, err := primitive.ObjectIDFromHex(messageId)
	if err != nil {
		// If the conversion fails, return an error response
		return c.Status(fiber.StatusBadRequest).JSON(fiber.Map{
			"error": "Invalid message ID format",
		})
	}

	// Delete messages in the database. This one and all the following ones
	collection := mongoClient.Database("chat").Collection("messages")

	// Get the date of the message
	var message Message
	err = collection.FindOne(context.TODO(), bson.M{"_id": objID}).Decode(&message)
	if err != nil {
		if err == mongo.ErrNoDocuments {
			return c.Status(fiber.StatusNotFound).JSON(fiber.Map{
				"error": "Message not found",
			})
		}
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
			"error": "Failed to find message",
		})
	}

	// Delete the message and all messages with a date after the specified date
	filter := bson.M{
		"$or": []bson.M{
			{"_id": objID},
			{"date": bson.M{"$gt": message.Date}},
		},
	}
	_, err = collection.DeleteMany(context.TODO(), filter)
	if err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
			"error": "Failed to delete messages",
		})
	}

	return c.SendString("")
}

func generateChatHTML(c *fiber.Ctx) error {
	// 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 {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
			"error": "Failed to find messages",
		})
	}

	// Convert the cursor to an array of messages
	var Messages []Message
	if err = cursor.All(context.TODO(), &Messages); err != nil {
		return c.Status(fiber.StatusInternalServerError).JSON(fiber.Map{
			"error": "Failed to convert cursor to array",
		})
	}

	// Render the HTML template with the messages
	return c.Render("partials/chat-messages", fiber.Map{
		"Messages":           Messages,
		"IncludePlaceholder": false,
	})
}

func isMongoDBConnectedHandler(c *fiber.Ctx) error {
	if mongoClient != nil {
		return c.SendString("<h1>Connected</h1>")
	}
	return c.SendString("<h1 hx-get='/isMongoDBConnected' hx-trigger='every 1s' hx-swap='outerHTM'>Not connected</h1>")
}

// SSE Stuff
var (
	eventChannel = make(chan string)
)

func sseHandler(c *fiber.Ctx) error {
	c.Set("Content-Type", "text/event-stream")
	c.Set("Cache-Control", "no-cache")
	c.Set("Connection", "keep-alive")
	c.Set("Transfer-Encoding", "chunked")

	c.Context().SetBodyStreamWriter(fasthttp.StreamWriter(func(w *bufio.Writer) {
		fmt.Println("WRITER")
		for {
			select {
			case msg := <-eventChannel:
				fmt.Fprintf(w, "data: %s\n\n", msg)
				err := w.Flush()
				if err != nil {
					fmt.Printf("Error while flushing: %v. Closing http connection.\n", err)
					return
				}
			default:
				if c.Context() != nil && c.Context().Done() != nil {
					select {
					case <-c.Context().Done():
						fmt.Println("Client connection closed")
						return
					default:
					}
				}
				time.Sleep(100 * time.Millisecond)
			}
		}
	}))

	return nil
}