From 26f599768731bfebdc03676ee51f36f16ef52fa4 Mon Sep 17 00:00:00 2001
From: Migan178 <migan9857@gmail.com>
Date: Wed, 25 Jan 2023 20:23:01 +0900
Subject: [PATCH] feat: turn on/off user train

---
 src/ChatBot.ts | 31 ++++++++++++++++++++++++-------
 src/Client.ts  | 31 +++++++++++++++++++++----------
 2 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/src/ChatBot.ts b/src/ChatBot.ts
index 89fdcad..e1ef25d 100644
--- a/src/ChatBot.ts
+++ b/src/ChatBot.ts
@@ -1,11 +1,16 @@
 import sqlite3 from 'sqlite3'
 import type { Client, Message } from 'discord.js'
+import { dirname, join } from 'node:path'
+import { fileURLToPath } from 'node:url'
+
+type TrainType = 'muffinOnly' | 'All'
 
 export default class ChatBot {
-  private db: sqlite3.Database
-  public constructor(dbPath: string) {
-    this.db = new sqlite3.Database(dbPath)
-  }
+  private db = new sqlite3.Database(
+    join(dirname(fileURLToPath(import.meta.url)), '..', 'db', 'db.sqlite3')
+  )
+  private trainType: TrainType = 'All'
+  public constructor() {}
 
   public getResponse(msg: Message, sendMsg?: boolean): ChatBot {
     this.db.all(
@@ -26,7 +31,7 @@ export default class ChatBot {
     return this
   }
 
-  public train(client: Client, user?: boolean): ChatBot {
+  public train(client: Client): ChatBot {
     client.on('messageCreate', msg => {
       if (msg.author.bot) return
       if (msg.author.id === '1026185545837191238') {
@@ -38,7 +43,7 @@ export default class ChatBot {
           }
         )
       } else {
-        if (!user) return
+        if (this.trainType !== 'All') return
         if (!msg.content.startsWith('머핀아 ')) return
         const sql = `INSERT INTO statement(text, persona) VALUES('${msg.content
           .replace('머핀아 ', '')
@@ -51,7 +56,19 @@ export default class ChatBot {
     return this
   }
 
-  public destroy() {
+  public changeTrainType(): TrainType {
+    switch (this.trainType) {
+      case 'muffinOnly':
+        this.trainType = 'All'
+        break
+      case 'All':
+        this.trainType = 'muffinOnly'
+        break
+    }
+    return this.trainType
+  }
+
+  public destroy(): void {
     this.db.close()
   }
 }
diff --git a/src/Client.ts b/src/Client.ts
index e2f38bf..597cf2e 100644
--- a/src/Client.ts
+++ b/src/Client.ts
@@ -1,7 +1,5 @@
 import { ActivityType, Client, GatewayIntentBits, Message } from 'discord.js'
 import ChatBot from './ChatBot.js'
-import { join, dirname } from 'node:path'
-import { fileURLToPath } from 'node:url'
 import Dokdo from 'dokdo'
 import 'dotenv/config'
 
@@ -17,10 +15,15 @@ function noPerm(msg: Message) {
   })
 }
 
+function isNotOwner(msg: Message): boolean {
+  if (msg.author.id !== '415135882006495242') {
+    noPerm(msg)
+    return false
+  } else return true
+}
+
 export default class MuffinAI extends Client {
-  private chatBot = new ChatBot(
-    join(dirname(fileURLToPath(import.meta.url)), '..', 'db', 'db.sqlite3')
-  )
+  private chatBot = new ChatBot()
   public constructor() {
     super({
       intents: [
@@ -32,7 +35,7 @@ export default class MuffinAI extends Client {
   }
 
   public override login(): Promise<string> {
-    this.chatBot.train(this, true)
+    this.chatBot.train(this)
     this.once('ready', client => {
       client.user!.setActivity({
         type: ActivityType.Playing,
@@ -49,11 +52,19 @@ export default class MuffinAI extends Client {
       }).run(msg)
       if (msg.content.startsWith('머핀아 ')) this.chatBot.getResponse(msg, true)
       else if (msg.content.startsWith('멒힌아 봇꺼')) {
-        if (msg.author.id !== '415135882006495242') {
-          noPerm(msg)
-          return
-        }
+        if (!isNotOwner(msg)) return
         this.destroy()
+      } else if (msg.content.startsWith('멒힌아 모드변경')) {
+        if (!isNotOwner(msg)) return
+        const a = this.chatBot.changeTrainType()
+        switch (a) {
+          case 'muffinOnly':
+            msg.channel.send('현재 모드: 머핀만 학습')
+            break
+          case 'All':
+            msg.channel.send('현재 모드: 전체 학습')
+            break
+        }
       } else return
     })
     return super.login()