feat: implement bot handler with remake functionality

bot/handler.go

@@ -1,91 +1,47 @@
-package command
+package bot
 
 import (
 	"database/sql"
-	"encoding/json"
 	"fmt"
 	"github.com/purofle/remake_bot/quotely"
+	"go.uber.org/zap"
 	tele "gopkg.in/telebot.v3"
-	"log"
 	"math/rand"
-	crand "math/rand"
-	"os"
 	"strconv"
-	"strings"
 	"sync"
 	"time"
 )
 
-type Country struct {
+type Handler struct {
-	CountryName string `json:"country"`
+	bot      *tele.Bot
-	Population  int64  `json:"population"`
+	logger   *zap.Logger
-}
+	database *sql.DB
 
-type RemakeData struct {
-	count   int64
-	country string
-	gender  string
-}
-
-var (
-	countryList     []Country
-	userList        []string
-	totalPopulation int64
 	mutex  sync.Mutex
-	remakeCount     map[int64]*RemakeData
+	remake *Remake
-)
-
-var database *sql.DB
-
-func InitHandler() {
-
-	connStr := "postgresql://postgres:114514@localhost:5432/postgres?sslmode=disable"
-	db, err := sql.Open("postgres", connStr)
-	if err != nil {
-		log.Fatal(err)
-	}
-	database = db
-
-	err = initList()
-	if err != nil {
-		log.Fatal(err)
 }
 
-	remakeCount = make(map[int64]*RemakeData)
+func NewHandler(bot *tele.Bot, logger *zap.Logger) *Handler {
+	return &Handler{
+		bot:    bot,
+		logger: logger,
+	}
 }
 
-func initList() error {
+func (h *Handler) RegisterAll() {
-	rawJson, err := os.ReadFile("countries.json")
+	h.bot.Handle(tele.OnQuery, h.InlineQuery)
-	if err != nil {
+	h.bot.Handle("/remake", h.CommandRemake)
-		return err
+	h.bot.Handle("/remake_data", h.CommandRemakeData)
+	h.bot.Handle(tele.OnText, h.CommandOnText)
 }
 
-	if err = json.Unmarshal(rawJson, &countryList); err != nil {
+func (h *Handler) getRandomCountry() Country {
-		return err
-	}
-
-	totalPopulation = int64(0)
-	for _, country := range countryList {
-		totalPopulation += country.Population
-	}
-
-	rawJson, err = os.ReadFile("user_list.json")
-	if err != nil {
-		return err
-	}
-	if err = json.Unmarshal(rawJson, &userList); err != nil {
-		return err
-	}
-	return nil
-}
-
-func getRandomCountry() Country {
 	// 生成随机数
-	randomNum := rand.Int63n(totalPopulation)
+	randomNum := rand.Int63n(h.remake.TotalPopulation)
 
 	// 根据随机数获取对应的国家
 	index := 0
-	for i, country := range countryList {
+	for i, country := range h.remake.CountryList {
 		if randomNum < country.Population {
 			index = i
 			break
@@ -93,30 +49,31 @@ func getRandomCountry() Country {
 		randomNum -= country.Population
 	}
 
-	return countryList[index]
+	return h.remake.CountryList[index]
 }
 
-func CommandRemake(c tele.Context) error {
+func (h *Handler) CommandRemake(c tele.Context) error {
-
 	msg := c.Message()
 
 	remakeData := []string{"男孩子", "女孩子", "MtF", "FtM", "MtC", "萝莉", "正太", "武装直升机", "沃尔玛购物袋", "星巴克", "无性别", "扶她", "死胎"}
 
 	remakeResult := rand.Intn(len(remakeData))
-	randomCountry := getRandomCountry()
+	randomCountry := h.getRandomCountry()
 
-	mutex.Lock()
+	func() {
-	_, hasKey := remakeCount[c.Sender().ID]
+		h.mutex.Lock()
+		defer h.mutex.Unlock()
+		_, hasKey := h.remake.RemakeCount[c.Sender().ID]
 		if !hasKey {
-		remakeCount[c.Sender().ID] = new(RemakeData)
+			h.remake.RemakeCount[c.Sender().ID] = new(RemakeData)
 		}
-	oldGender := remakeCount[c.Sender().ID].count
+		oldGender := h.remake.RemakeCount[c.Sender().ID].count
-	remakeCount[c.Sender().ID] = &RemakeData{
+		h.remake.RemakeCount[c.Sender().ID] = &RemakeData{
 			country: randomCountry.CountryName,
 			gender:  remakeData[remakeResult],
 			count:   oldGender + 1,
 		}
-	mutex.Unlock()
+	}()
 
 	text := fmt.Sprintf("转生成功！您现在是 %s 的 %s 了。", randomCountry.CountryName, remakeData[remakeResult])
 
@@ -139,12 +96,12 @@ func CommandRemake(c tele.Context) error {
 	return nil
 }
 
-func CommandRemakeData(c tele.Context) error {
+func (h *Handler) CommandRemakeData(c tele.Context) error {
 
 	msg := c.Message()
 
 	var text string
-	userData, hasKey := remakeCount[c.Sender().ID]
+	userData, hasKey := h.remake.RemakeCount[c.Sender().ID]
 	if hasKey {
 		text = fmt.Sprintf("您现在是 %s 的 %s，共 remake 了 %d 次", userData.country, userData.gender, userData.count)
 	} else {
@@ -170,42 +127,7 @@ func CommandRemakeData(c tele.Context) error {
 	return nil
 }
 
-func CommandEat(c tele.Context) error {
+func (*Handler) CommandOnText(c tele.Context) error {
-	if !(c.Chat().Type == tele.ChatPrivate || c.Chat().ID == -1001965344356) {
-		fmt.Println(c.Chat().ID)
-		return nil
-	}
-
-	method := []string{"炒", "蒸", "煮", "红烧", "爆炒", "烤", "炸", "煎", "炖", "焖", "炖", "卤"}
-
-	// 获取时间段
-	hour := time.Now().Hour()
-	var hourText string
-	switch {
-	case hour > 6 && hour <= 10:
-		hourText = "早上"
-	case hour > 10 && hour <= 14:
-		hourText = "中午"
-	case hour > 14 && hour <= 17:
-		hourText = "下午"
-	case hour > 18 && hour <= 21:
-		hourText = "晚上"
-	default:
-		hourText = "宵夜"
-	}
-
-	var name string
-	if strings.Contains(c.Sender().FirstName, " | ") {
-		name = strings.Split(c.Sender().FirstName, " | ")[0]
-	} else {
-		name = c.Sender().FirstName
-	}
-
-	result := fmt.Sprintf("今天%s吃 %s %s %s", hourText, name, method[rand.Intn(len(method))], userList[crand.Intn(len(userList))])
-	return c.Reply(result)
-}
-
-func CommandOnText(c tele.Context) error {
 	if c.Chat().ID != -1001965344356 {
 		return nil
 	}
@@ -219,15 +141,15 @@ func CommandOnText(c tele.Context) error {
 	return nil
 }
 
-func getQuote(text string) (error, []string, []string) {
+func (h *Handler) getQuote(text string) (error, []string, []string) {
 	var rows *sql.Rows
 	var err error
 	if text == "" {
 		query := "select text, \"from\" from result_new where from_id not like 'channel%' order by random() limit 50"
-		rows, err = database.Query(query)
+		rows, err = h.database.Query(query)
 	} else {
 		query := "select text, \"from\" from result_new where from_id not like 'channel%' AND text like '%' || $1 || '%' order by random() limit 50"
-		rows, err = database.Query(query, text)
+		rows, err = h.database.Query(query, text)
 	}
 	if err != nil {
 		return err, nil, nil
@@ -254,7 +176,7 @@ func getQuote(text string) (error, []string, []string) {
 	return nil, resultText, from
 }
 
-func InlineQuery(c tele.Context) error {
+func (h *Handler) InlineQuery(c tele.Context) error {
 	member, err := c.Bot().ChatMemberOf(
 		&tele.Chat{ID: -1001965344356},
 		c.Sender(),
@@ -272,9 +194,9 @@ func InlineQuery(c tele.Context) error {
 	var from []string
 
 	if c.Query().Text == "" {
-		err, resultText, from = getQuote("")
+		err, resultText, from = h.getQuote("")
 	} else {
-		err, resultText, from = getQuote(c.Query().Text)
+		err, resultText, from = h.getQuote(c.Query().Text)
 	}
 	results := make(tele.Results, len(resultText))
 
@@ -296,3 +218,38 @@ func InlineQuery(c tele.Context) error {
 		CacheTime: 0,
 	})
 }
+
+//func CommandEat(c tele.Context) error {
+//	if !(c.Chat().Type == tele.ChatPrivate || c.Chat().ID == -1001965344356) {
+//		fmt.Println(c.Chat().ID)
+//		return nil
+//	}
+//
+//	method := []string{"炒", "蒸", "煮", "红烧", "爆炒", "烤", "炸", "煎", "炖", "焖", "炖", "卤"}
+//
+//	// 获取时间段
+//	hour := time.Now().Hour()
+//	var hourText string
+//	switch {
+//	case hour > 6 && hour <= 10:
+//		hourText = "早上"
+//	case hour > 10 && hour <= 14:
+//		hourText = "中午"
+//	case hour > 14 && hour <= 17:
+//		hourText = "下午"
+//	case hour > 18 && hour <= 21:
+//		hourText = "晚上"
+//	default:
+//		hourText = "宵夜"
+//	}
+//
+//	var name string
+//	if strings.Contains(c.Sender().FirstName, " | ") {
+//		name = strings.Split(c.Sender().FirstName, " | ")[0]
+//	} else {
+//		name = c.Sender().FirstName
+//	}
+//
+//	result := fmt.Sprintf("今天%s吃 %s %s %s", hourText, name, method[rand.Intn(len(method))], userList[crand.Intn(len(userList))])
+//	return c.Reply(result)
+//}

bot/module.go

@@ -0,0 +1,11 @@
+package bot
+
+import "go.uber.org/fx"
+
+var Module = fx.Options(
+	fx.Provide(NewRemake),
+	fx.Provide(NewHandler),
+	fx.Invoke(func(handler *Handler) {
+		handler.RegisterAll()
+	}),
+)

bot/remake.go

@@ -0,0 +1,44 @@
+package bot
+
+import (
+	"encoding/json"
+	"log"
+	"os"
+)
+
+type Country struct {
+	CountryName string `json:"country"`
+	Population  int64  `json:"population"`
+}
+
+type RemakeData struct {
+	count   int64
+	country string
+	gender  string
+}
+
+type Remake struct {
+	CountryList     []Country
+	TotalPopulation int64
+	RemakeCount     map[int64]*RemakeData
+}
+
+func NewRemake() *Remake {
+
+	var remake Remake
+
+	rawJson, err := os.ReadFile("countries.json")
+	if err != nil {
+		log.Fatal("Error reading countries.json:", err)
+	}
+
+	if err = json.Unmarshal(rawJson, &remake.CountryList); err != nil {
+		log.Fatal("Error unmarshalling countries.json:", err)
+	}
+
+	remake.TotalPopulation = int64(0)
+	for _, country := range remake.CountryList {
+		remake.TotalPopulation += country.Population
+	}
+	return &remake
+}

bot/tele.go

@@ -2,7 +2,6 @@ package bot
 
 import (
 	"context"
-	"github.com/purofle/remake_bot/command"
 	"go.uber.org/fx"
 	"go.uber.org/zap"
 	tele "gopkg.in/telebot.v3"
@@ -25,17 +24,15 @@ func NewRemakeBot(lc fx.Lifecycle, logger *zap.Logger) *tele.Bot {
 	}
 
 	lc.Append(fx.Hook{OnStart: func(ctx context.Context) error {
-
-		command.InitHandler()
-
-		b.Handle(tele.OnQuery, command.InlineQuery)
-		b.Handle("/remake", command.CommandRemake)
-		b.Handle("/remake_data", command.CommandRemakeData)
-		b.Handle("/eat", command.CommandEat)
-		b.Handle(tele.OnText, command.CommandOnText)
 		go b.Start()
 
-		logger.Info("Remake Bot is now running...")
+		logger.Info("remake Bot is now running...")
+
+		return nil
+	}, OnStop: func(ctx context.Context) error {
+		b.Stop()
+
+		logger.Info("remake Bot is now stopped")
 
 		return nil
 	}})

main.go

@@ -15,7 +15,9 @@ func main() {
 		fx.WithLogger(func(log *zap.Logger) fxevent.Logger {
 			return &fxevent.ZapLogger{Logger: log}
 		}),
+
 		fx.Provide(bot.NewRemakeBot),
+		bot.Module,
 		fx.Invoke(func(bot *telebot.Bot) {}),
 	).Run()
 }