From 6ea5198fd250730cf4234520a935049235a09eee Mon Sep 17 00:00:00 2001
From: Siwoo Jeon <me@migan.co.kr>
Date: Mon, 24 Mar 2025 21:07:40 +0900
Subject: [PATCH] feat: Add save

---
 configs/config.go        | 22 +++++++++++++++++-----
 databases/Text.go        | 10 +++++++++-
 handler/messageCreate.go | 25 +++++++++++++++++++++++--
 main.go                  |  2 +-
 4 files changed, 50 insertions(+), 9 deletions(-)

diff --git a/configs/config.go b/configs/config.go
index 940aa2d..49516d3 100644
--- a/configs/config.go
+++ b/configs/config.go
@@ -10,10 +10,19 @@ import (
 
 var MUFFIN_VERSION = "0.0.0-gopher_canary.250323b"
 
+type botConfig struct {
+	Token  string
+	Prefix string
+}
+
+type trainConfig struct {
+	UserID string
+}
+
 // MuffinConfig for Muffin bot
 type MuffinConfig struct {
-	Token       string
-	Prefix      string
+	Bot         botConfig
+	Train       trainConfig
 	DatabaseURL string
 }
 
@@ -23,15 +32,18 @@ func loadConfig() *MuffinConfig {
 		fmt.Println("[goMuffin] 봇의 설절파일을 불러올 수가 없어요.")
 		log.Fatalln(err)
 	}
-	config := MuffinConfig{}
+	config := MuffinConfig{Bot: botConfig{}, Train: trainConfig{}}
 	setConfig(&config)
 
 	return &config
 }
 
 func setConfig(config *MuffinConfig) {
-	config.Prefix = os.Getenv("BOT_PREFIX")
-	config.Token = os.Getenv("BOT_TOKEN")
+	config.Bot.Prefix = os.Getenv("BOT_PREFIX")
+	config.Bot.Token = os.Getenv("BOT_TOKEN")
+
+	config.Train.UserID = os.Getenv("TRAIN_USER_ID")
+
 	config.DatabaseURL = os.Getenv("DATABASE_URL")
 }
 
diff --git a/databases/Text.go b/databases/Text.go
index 93a3bef..295d688 100644
--- a/databases/Text.go
+++ b/databases/Text.go
@@ -1,15 +1,23 @@
 package databases
 
 import (
+	"time"
+
 	"go.mongodb.org/mongo-driver/v2/bson"
 	"go.mongodb.org/mongo-driver/v2/mongo"
 )
 
+type InsertText struct {
+	Text      string
+	Persona   string
+	CreatedAt time.Time `bson:"created_at"`
+}
+
 type Text struct {
 	Id        bson.ObjectID `bson:"_id"`
 	Text      string
 	Persona   string
-	CreatedAt bson.DateTime `bson:"created_at"`
+	CreatedAt time.Time `bson:"created_at"`
 }
 
 var Texts *mongo.Collection = Client.Database("muffin_ai_test").Collection("text")
diff --git a/handler/messageCreate.go b/handler/messageCreate.go
index e72bcd1..4d303df 100644
--- a/handler/messageCreate.go
+++ b/handler/messageCreate.go
@@ -5,6 +5,7 @@ import (
 	"log"
 	"math/rand"
 	"strings"
+	"time"
 
 	"git.wh64.net/muffin/goMuffin/commands"
 	"git.wh64.net/muffin/goMuffin/configs"
@@ -20,10 +21,20 @@ func MessageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
 		return
 	}
 
-	if strings.HasPrefix(m.Content, config.Prefix) {
-		content := strings.TrimPrefix(m.Content, config.Prefix)
+	if strings.HasPrefix(m.Content, config.Bot.Prefix) {
+		content := strings.TrimPrefix(m.Content, config.Bot.Prefix)
 		command := commands.Discommand.Aliases[content]
 
+		if m.Author.ID == config.Train.UserID {
+			if _, err := databases.Texts.InsertOne(context.TODO(), databases.InsertText{
+				Text:      content,
+				Persona:   "muffin",
+				CreatedAt: time.Now(),
+			}); err != nil {
+				log.Fatalln(err)
+			}
+		}
+
 		if command == "" {
 			s.ChannelTyping(m.ChannelID)
 
@@ -33,6 +44,14 @@ func MessageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
 			channel, _ := s.Channel(m.ChannelID)
 			if channel.NSFW {
 				filter = bson.D{{}}
+
+				if _, err := databases.Texts.InsertOne(context.TODO(), databases.InsertText{
+					Text:      content,
+					Persona:   "user:" + m.Author.Username,
+					CreatedAt: time.Now(),
+				}); err != nil {
+					log.Fatalln(err)
+				}
 			} else {
 				filter = bson.D{{Key: "persona", Value: "muffin"}}
 			}
@@ -46,9 +65,11 @@ func MessageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
 			cur.All(context.TODO(), &datas)
 
 			s.ChannelMessageSendReply(m.ChannelID, datas[rand.Intn(len(datas))].Text, m.Reference())
+			return
 		}
 
 		commands.Discommand.MessageRun(command, s, m)
+		return
 	} else {
 		return
 	}
diff --git a/main.go b/main.go
index f74d453..d4d6294 100644
--- a/main.go
+++ b/main.go
@@ -17,7 +17,7 @@ import (
 func main() {
 	config := configs.Config
 
-	dg, err := discordgo.New("Bot " + config.Token)
+	dg, err := discordgo.New("Bot " + config.Bot.Token)
 	if err != nil {
 		fmt.Println("[goMuffin] 봇의 세션을 만들수가 없어요.")
 		log.Fatalln(err)