diff --git a/commands/command.go b/commands/command.go index 44355b8..db6f594 100644 --- a/commands/command.go +++ b/commands/command.go @@ -9,4 +9,5 @@ const ( Go = Command("/go") Stat = Command("/stat") Eat = Command("/eat") + Pause = Command("/pause") ) diff --git a/db/chat.go b/db/chat.go index 73356f0..17deff1 100644 --- a/db/chat.go +++ b/db/chat.go @@ -1,9 +1,10 @@ package db const ( - UserStateNone = UserState("") - UserStateGo = UserState("Go") - UserStateEat = UserState("Eat") + UserStateNone = UserState("") + UserStateGo = UserState("Go") + UserStateEat = UserState("Eat") + UserStatePause = UserState("Pause") ) type UserState string diff --git a/db/db.go b/db/db.go index 9311814..dc8feb4 100644 --- a/db/db.go +++ b/db/db.go @@ -2,12 +2,13 @@ package db import ( "context" - "go.mongodb.org/mongo-driver/bson" - "go.mongodb.org/mongo-driver/mongo" - "go.mongodb.org/mongo-driver/mongo/options" "log" "time" "valera/config" + + "go.mongodb.org/mongo-driver/bson" + "go.mongodb.org/mongo-driver/mongo" + "go.mongodb.org/mongo-driver/mongo/options" ) type DB struct { @@ -112,6 +113,19 @@ func (db *DB) SetStatusToChat(chatID int64, status UserState) error { return err } +func (db *DB) SetPause(chatID int64, duration time.Duration) error { + ctx := context.Background() + + loc, _ := time.LoadLocation("Asia/Novosibirsk") + t := time.Now().In(loc).Add(+duration) + _, err := db.chatsColection.UpdateOne( + ctx, + bson.M{"chat_id": chatID}, + bson.M{"$set": bson.M{"pause_until": t}}, + ) + return err +} + func (db *DB) GetStat(chatID int64, cron bool) (map[string]int, error) { ctx := context.Background() @@ -178,7 +192,21 @@ func (db *DB) GetStat(chatID int64, cron bool) (map[string]int, error) { func (db *DB) GetAllChats() ([]int64, error) { ctx := context.Background() - cursor, err := db.chatsColection.Find(ctx, bson.M{}) + loc, _ := time.LoadLocation("Asia/Novosibirsk") + t := time.Now().In(loc) + cursor, err := db.chatsColection.Find( + ctx, + bson.M{ + "$or": bson.A{ + bson.M{ + "pause_until": bson.M{"$lt": t}, + }, + bson.M{ + "pause_until": bson.M{"$exists": false}, + }, + }, + }, + ) if err != nil { return nil, err } diff --git a/main.go b/main.go index e03dbd8..46dc96e 100644 --- a/main.go +++ b/main.go @@ -12,6 +12,7 @@ import ( "valera/calories" "valera/commands" "valera/db" + "valera/pause" tgbot "github.com/go-telegram-bot-api/telegram-bot-api" "github.com/umputun/go-flags" @@ -27,7 +28,7 @@ var ( ) const ( - version = "v1.5.0" + version = "v1.6.0" ) type Opts struct { @@ -105,7 +106,7 @@ func run() { } go func() { - http.HandleFunc("/go", func(w http.ResponseWriter, r *http.Request) { + http.HandleFunc("/go", func(w http.ResponseWriter, _ *http.Request) { chats, err := dataBase.GetAllChats() if err != nil { w.Header().Add("Content-Type", "application/json") @@ -121,7 +122,7 @@ func run() { _, _ = fmt.Fprintf(w, `{"result":"ok"}`) }) - http.HandleFunc("/stat", func(w http.ResponseWriter, r *http.Request) { + http.HandleFunc("/stat", func(w http.ResponseWriter, _ *http.Request) { chats, err := dataBase.GetAllChats() if err != nil { fmt.Println(err) @@ -161,6 +162,16 @@ func run() { chatID := update.Message.Chat.ID username := update.Message.From.UserName + if text == "/c" { + msg := tgbot.NewMessage(chatID, "Чистка") + msg.ReplyMarkup = tgbot.NewRemoveKeyboard(false) + _, _ = bot.Send(msg) + if err := dataBase.SetStatusToChat(chatID, db.UserStateNone); err != nil { + log.Println(err) + } + continue + } + userInfoDTO, err := dataBase.GetChatInfo(chatID) if err != nil { log.Println(err) @@ -223,6 +234,22 @@ func run() { log.Println(err) } continue + case db.UserStatePause: + duration, err := pause.GetDuration(text) + if err != nil { + log.Println(err) + continue + } + if err := dataBase.SetPause(chatID, duration); err != nil { + log.Println(err) + } + msg := tgbot.NewMessage(chatID, fmt.Sprintf("Поставил паузу %v, отдыхай", duration)) + msg.ReplyMarkup = tgbot.NewRemoveKeyboard(false) + _, _ = bot.Send(msg) + if err := dataBase.SetStatusToChat(chatID, db.UserStateNone); err != nil { + log.Println(err) + } + continue } command := commands.Command(strings.Replace(text, opts.Name, "", 1)) @@ -249,6 +276,24 @@ func run() { log.Println(err) } } + case commands.Pause: + msg := tgbot.NewMessage(chatID, "Хочешь отдохнуть? Сколько времени тебе нужно?") + msg.ReplyMarkup = tgbot.NewReplyKeyboard([][]tgbot.KeyboardButton{ + tgbot.NewKeyboardButtonRow( + tgbot.NewKeyboardButton("1ч"), + tgbot.NewKeyboardButton("4ч"), + ), + tgbot.NewKeyboardButtonRow( + tgbot.NewKeyboardButton("1д"), + tgbot.NewKeyboardButton("2д"), + tgbot.NewKeyboardButton("7д"), + ), + }...) + if _, err := bot.Send(msg); err == nil { + if err := dataBase.SetStatusToChat(chatID, db.UserStatePause); err != nil { + log.Println(err) + } + } } } } diff --git a/pause/pause.go b/pause/pause.go new file mode 100644 index 0000000..e05307a --- /dev/null +++ b/pause/pause.go @@ -0,0 +1,21 @@ +package pause + +import ( + "strconv" + "strings" + "time" +) + +func GetDuration(text string) (time.Duration, error) { + d := strings.TrimSuffix(text, "ч") + d = strings.TrimSuffix(d, "д") + count, err := strconv.Atoi(d) + if err != nil { + return 0, err + } + res := time.Duration(count) * time.Hour + if strings.HasSuffix(text, "д") { + res *= 24 + } + return res, nil +}