feat: implement bot handler with remake functionality

This commit is contained in:
purofle 2025-05-06 02:01:39 +08:00
parent 5b37ae6794
commit 2c57dd9439
No known key found for this signature in database
5 changed files with 147 additions and 136 deletions

View File

@ -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 {
CountryName string `json:"country"`
Population int64 `json:"population"`
}
type Handler struct {
bot *tele.Bot
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
)
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)
remake *Remake
}
func initList() error {
rawJson, err := os.ReadFile("countries.json")
if err != nil {
return err
func NewHandler(bot *tele.Bot, logger *zap.Logger) *Handler {
return &Handler{
bot: bot,
logger: logger,
}
if err = json.Unmarshal(rawJson, &countryList); err != nil {
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 {
func (h *Handler) RegisterAll() {
h.bot.Handle(tele.OnQuery, h.InlineQuery)
h.bot.Handle("/remake", h.CommandRemake)
h.bot.Handle("/remake_data", h.CommandRemakeData)
h.bot.Handle(tele.OnText, h.CommandOnText)
}
func (h *Handler) 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()
_, hasKey := remakeCount[c.Sender().ID]
func() {
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
remakeCount[c.Sender().ID] = &RemakeData{
oldGender := h.remake.RemakeCount[c.Sender().ID].count
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 {
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 {
func (*Handler) 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)
//}

11
bot/module.go Normal file
View File

@ -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()
}),
)

44
bot/remake.go Normal file
View File

@ -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
}

View File

@ -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
}})

View File

@ -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()
}