package main import ( "context" "fmt" "log" "os" "os/signal" "strconv" "strings" "syscall" "time" "dcdev.ro/CfrTrainInfoTelegramBot/pkg/database" "dcdev.ro/CfrTrainInfoTelegramBot/pkg/handlers" "dcdev.ro/CfrTrainInfoTelegramBot/pkg/subscriptions" "dcdev.ro/CfrTrainInfoTelegramBot/pkg/utils" tgBot "github.com/go-telegram/bot" "github.com/go-telegram/bot/models" "gorm.io/driver/sqlite" "gorm.io/gorm" ) const ( trainInfoCommand = "/train_info" stationInfoCommand = "/station_info" routeCommand = "/route" cancelCommand = "/cancel" initialMessage = `Hello. 😄 You can send the following commands: ` + trainInfoCommand + ` - Find information about a certain train. ` + stationInfoCommand + ` - Find departures or arrivals at a certain station. ` + routeCommand + ` - Find trains for a certain route. You may use ` + cancelCommand + ` to cancel any ongoing command.` waitingForTrainNumberMessage = "Please send the number of the train you want information for." pleaseWaitMessage = "Please wait..." cancelResponseMessage = "Command cancelled." chooseDateMessage = `Please choose the date of departure from the first station for this train. You may also send the date as a message in the following formats: dd.mm.yyyy, m/d/yyyy, yyyy-mm-dd, UNIX timestamp. Keep in mind that, for night trains, this date might be yesterday.` invalidDateMessage = "Invalid date. Please try again or us " + cancelCommand + " to cancel." ) func main() { ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer cancel() log.SetOutput(os.Stderr) botToken := os.Getenv("CFR_BOT.TOKEN") botToken = strings.TrimSpace(botToken) if len(botToken) == 0 { log.Fatal("ERROR: No bot token supplied; supply with CFR_BOT.TOKEN") } db, err := gorm.Open(sqlite.Open("bot_db.sqlite"), &gorm.Config{}) if err != nil { panic(err) } if err := db.AutoMigrate(&handlers.ChatFlow{}); err != nil { panic(err) } if err := db.AutoMigrate(&subscriptions.SubData{}); err != nil { panic(err) } database.SetDatabase(db) subBot, err := tgBot.New(botToken) if err != nil { panic(err) } subs, err := subscriptions.LoadSubscriptions(subBot) if err != nil { subs = nil fmt.Printf("WARN : Could not load subscriptions: %s\n", err.Error()) } go subs.CheckSubscriptions(ctx) bot, err := tgBot.New(botToken, tgBot.WithDefaultHandler(handlerBuilder(subs))) if err != nil { panic(err) } log.Print("INFO : Starting...") bot.Start(ctx) } func handlerBuilder(subs *subscriptions.Subscriptions) func(context.Context, *tgBot.Bot, *models.Update) { return func(ctx context.Context, b *tgBot.Bot, update *models.Update) { handler(ctx, b, update, subs) } } func handler(ctx context.Context, b *tgBot.Bot, update *models.Update, subs *subscriptions.Subscriptions) { var response *handlers.HandlerResponse var toEditId int defer func() { if response == nil { return } if response.ProgressMessageToEditId != 0 { toEditId = response.ProgressMessageToEditId } if response.Message != nil { response.Message.ChatID = response.Injected.ChatId if toEditId != 0 { b.EditMessageText(ctx, &tgBot.EditMessageTextParams{ ChatID: response.Message.ChatID, MessageID: toEditId, Text: response.Message.Text, ParseMode: response.Message.ParseMode, Entities: response.Message.Entities, DisableWebPagePreview: response.Message.DisableWebPagePreview, ReplyMarkup: response.Message.ReplyMarkup, }) } else { b.SendMessage(ctx, response.Message) } } if response.CallbackAnswer != nil { b.AnswerCallbackQuery(ctx, response.CallbackAnswer) } for _, edit := range response.MessageEdits { if (edit.ChatID == nil || edit.MessageID == 0) && edit.InlineMessageID == "" { edit.ChatID = response.Injected.ChatId edit.MessageID = response.Injected.MessageId } b.EditMessageText(ctx, edit) } for _, edit := range response.MessageMarkupEdits { if (edit.ChatID == nil || edit.MessageID == 0) && edit.InlineMessageID == "" { edit.ChatID = response.Injected.ChatId edit.MessageID = response.Injected.MessageId } b.EditMessageReplyMarkup(ctx, edit) } }() if update.Message != nil { defer func() { if response == nil { response = &handlers.HandlerResponse{} } response.Injected.ChatId = update.Message.Chat.ID response.Injected.MessageId = update.Message.ID }() log.Printf("DEBUG: Got message: %s\n", update.Message.Text) chatFlow := handlers.GetChatFlow(update.Message.Chat.ID) switch { case strings.HasPrefix(update.Message.Text, trainInfoCommand): response = handleFindTrainStages(ctx, b, update) case strings.HasPrefix(update.Message.Text, cancelCommand): handlers.SetChatFlow(chatFlow, handlers.InitialFlowType, handlers.InitialFlowType, "") response = &handlers.HandlerResponse{ Message: &tgBot.SendMessageParams{ Text: cancelResponseMessage, }, } default: switch chatFlow.Type { case handlers.InitialFlowType: b.SendMessage(ctx, &tgBot.SendMessageParams{ ChatID: update.Message.Chat.ID, Text: initialMessage, }) case handlers.TrainInfoFlowType: log.Printf("DEBUG: trainInfoFlowType with stage %s\n", chatFlow.Stage) response = handleFindTrainStages(ctx, b, update) } } } if update.CallbackQuery != nil { defer func() { if response == nil { response = &handlers.HandlerResponse{ CallbackAnswer: &tgBot.AnswerCallbackQueryParams{ CallbackQueryID: update.CallbackQuery.ID, }, } } response.Injected.ChatId = update.CallbackQuery.Message.Chat.ID response.Injected.MessageId = update.CallbackQuery.Message.ID if response.CallbackAnswer == nil { response.CallbackAnswer = &tgBot.AnswerCallbackQueryParams{ CallbackQueryID: update.CallbackQuery.ID, } } if response.CallbackAnswer.CallbackQueryID == "" { response.CallbackAnswer.CallbackQueryID = update.CallbackQuery.ID } }() chatFlow := handlers.GetChatFlow(update.CallbackQuery.Message.Chat.ID) if len(update.CallbackQuery.Data) != 0 { splitted := strings.Split(update.CallbackQuery.Data, "\x1b") switch splitted[0] { case handlers.TrainInfoChooseDateCallbackQuery: trainNumber := splitted[1] dateInt, _ := strconv.ParseInt(splitted[2], 10, 64) date := time.Unix(dateInt, 0) message, err := b.SendMessage(ctx, &tgBot.SendMessageParams{ ChatID: update.CallbackQuery.Message.Chat.ID, Text: pleaseWaitMessage, }) response, _ = handlers.HandleTrainNumberCommand(ctx, trainNumber, date, -1, false) if err == nil { response.ProgressMessageToEditId = message.ID } handlers.SetChatFlow(chatFlow, handlers.InitialFlowType, handlers.InitialFlowType, "") case handlers.TrainInfoChooseGroupCallbackQuery: trainNumber := splitted[1] dateInt, _ := strconv.ParseInt(splitted[2], 10, 64) date := time.Unix(dateInt, 0) groupIndex, _ := strconv.ParseInt(splitted[3], 10, 31) originalResponse, _ := handlers.HandleTrainNumberCommand(ctx, trainNumber, date, int(groupIndex), false) response = &handlers.HandlerResponse{ MessageEdits: []*tgBot.EditMessageTextParams{ { Text: originalResponse.Message.Text, ParseMode: originalResponse.Message.ParseMode, Entities: originalResponse.Message.Entities, DisableWebPagePreview: originalResponse.Message.DisableWebPagePreview, ReplyMarkup: originalResponse.Message.ReplyMarkup, }, }, } case handlers.TrainInfoSubscribeCallbackQuery: trainNumber := splitted[1] dateInt, _ := strconv.ParseInt(splitted[2], 10, 64) date := time.Unix(dateInt, 0) groupIndex, _ := strconv.ParseInt(splitted[3], 10, 31) err := subs.InsertSubscription(subscriptions.SubData{ ChatId: update.CallbackQuery.Message.Chat.ID, MessageId: update.CallbackQuery.Message.ID, TrainNumber: trainNumber, Date: date, GroupIndex: int(groupIndex), }) if err != nil { log.Printf("ERROR: Subscribe error: %s", err.Error()) response = &handlers.HandlerResponse{ CallbackAnswer: &tgBot.AnswerCallbackQueryParams{ Text: fmt.Sprintf("Error when subscribing."), ShowAlert: true, }, } } else { // TODO: Update message to contain unsubscribe button response = &handlers.HandlerResponse{ CallbackAnswer: &tgBot.AnswerCallbackQueryParams{ Text: fmt.Sprintf("Subscribed successfully!"), }, MessageMarkupEdits: []*tgBot.EditMessageReplyMarkupParams{ { ChatID: update.CallbackQuery.Message.Chat.ID, MessageID: update.CallbackQuery.Message.ID, ReplyMarkup: handlers.GetTrainNumberCommandResponseButtons(trainNumber, date, int(groupIndex), handlers.TrainInfoResponseButtonIncludeUnsub), }, }, } } case handlers.TrainInfoUnsubscribeCallbackQuery: trainNumber := splitted[1] dateInt, _ := strconv.ParseInt(splitted[2], 10, 64) date := time.Unix(dateInt, 0) groupIndex, _ := strconv.ParseInt(splitted[3], 10, 31) _, err := subs.DeleteSubscription(update.CallbackQuery.Message.Chat.ID, update.CallbackQuery.Message.ID) if err != nil { log.Printf("ERROR: Unsubscribe error: %s", err.Error()) response = &handlers.HandlerResponse{ CallbackAnswer: &tgBot.AnswerCallbackQueryParams{ Text: fmt.Sprintf("Error when unsubscribing."), ShowAlert: true, }, } } else { // TODO: Update message to contain unsubscribe button response = &handlers.HandlerResponse{ CallbackAnswer: &tgBot.AnswerCallbackQueryParams{ Text: fmt.Sprintf("Unsubscribed successfully!"), }, MessageMarkupEdits: []*tgBot.EditMessageReplyMarkupParams{ { ChatID: update.CallbackQuery.Message.Chat.ID, MessageID: update.CallbackQuery.Message.ID, ReplyMarkup: handlers.GetTrainNumberCommandResponseButtons(trainNumber, date, int(groupIndex), handlers.TrainInfoResponseButtonIncludeSub), }, }, } } } } } } func handleFindTrainStages(ctx context.Context, b *tgBot.Bot, update *models.Update) *handlers.HandlerResponse { log.Println("DEBUG: handleFindTrainStages") var response *handlers.HandlerResponse var chatId int64 if update.Message != nil { chatId = update.Message.Chat.ID } if update.CallbackQuery != nil { chatId = update.CallbackQuery.Message.Chat.ID } chatFlow := handlers.GetChatFlow(chatId) switch chatFlow.Type { case handlers.InitialFlowType: // Only command is possible here commandParamsString := strings.TrimPrefix(update.Message.Text, trainInfoCommand) commandParamsString = strings.TrimSpace(commandParamsString) commandParams := strings.Split(commandParamsString, " ") if len(commandParams) > 1 { message, err := b.SendMessage(ctx, &tgBot.SendMessageParams{ ChatID: update.Message.Chat.ID, Text: pleaseWaitMessage, }) trainNumber := commandParams[0] date := time.Now() groupIndex := -1 if len(commandParams) > 1 { date, _ = time.Parse(time.RFC3339, commandParams[1]) } if len(commandParams) > 2 { groupIndex, _ = strconv.Atoi(commandParams[2]) } response, _ = handlers.HandleTrainNumberCommand(ctx, trainNumber, date, groupIndex, false) if err == nil { response.ProgressMessageToEditId = message.ID } } else if len(commandParams) > 0 && len(commandParams[0]) != 0 { // Got only train number trainNumber := commandParams[0] response = getTrainInfoChooseDateResponse(trainNumber) handlers.SetChatFlow(chatFlow, handlers.TrainInfoFlowType, handlers.WaitingForDateStage, trainNumber) } else { response = &handlers.HandlerResponse{ Message: &tgBot.SendMessageParams{ Text: waitingForTrainNumberMessage, }, } handlers.SetChatFlow(chatFlow, handlers.TrainInfoFlowType, handlers.WaitingForTrainNumberStage, "") } case handlers.TrainInfoFlowType: switch chatFlow.Stage { case handlers.WaitingForTrainNumberStage: trainNumber := update.Message.Text response = getTrainInfoChooseDateResponse(trainNumber) handlers.SetChatFlow(chatFlow, handlers.TrainInfoFlowType, handlers.WaitingForDateStage, trainNumber) case handlers.WaitingForDateStage: date, err := utils.ParseDate(update.Message.Text) if err != nil { response = &handlers.HandlerResponse{ Message: &tgBot.SendMessageParams{ Text: invalidDateMessage, }, } } else { message, err := b.SendMessage(ctx, &tgBot.SendMessageParams{ ChatID: update.Message.Chat.ID, Text: pleaseWaitMessage, }) response, _ = handlers.HandleTrainNumberCommand(ctx, chatFlow.Extra, date, -1, false) if err == nil { response.ProgressMessageToEditId = message.ID } handlers.SetChatFlow(chatFlow, handlers.InitialFlowType, handlers.InitialFlowType, "") } } } return response } func getTrainInfoChooseDateResponse(trainNumber string) *handlers.HandlerResponse { replyButtons := make([][]models.InlineKeyboardButton, 0, 4) replyButtons = append(replyButtons, []models.InlineKeyboardButton{ { Text: fmt.Sprintf("Yesterday (%s)", time.Now().Add(time.Hour*-24).In(utils.Location).Format("02.01.2006")), CallbackData: fmt.Sprintf(handlers.TrainInfoChooseDateCallbackQuery+"\x1b%s\x1b%d", trainNumber, time.Now().Add(time.Hour*-24).Unix()), }, { Text: fmt.Sprintf("Today (%s)", time.Now().In(utils.Location).Format("02.01.2006")), CallbackData: fmt.Sprintf(handlers.TrainInfoChooseDateCallbackQuery+"\x1b%s\x1b%d", trainNumber, time.Now().Unix()), }, }) for i := 1; i < 4; i++ { arr := make([]models.InlineKeyboardButton, 0, 7) for j := 0; j < 7; j++ { ts := time.Now().Add(time.Hour * time.Duration(24*(j+(i-1)*7+1))).In(utils.Location) arr = append(arr, models.InlineKeyboardButton{ Text: ts.Format("02.01"), CallbackData: fmt.Sprintf(handlers.TrainInfoChooseDateCallbackQuery+"\x1b%s\x1b%d", trainNumber, ts.Unix()), }) } replyButtons = append(replyButtons, arr) } return &handlers.HandlerResponse{ Message: &tgBot.SendMessageParams{ Text: chooseDateMessage, ReplyMarkup: models.InlineKeyboardMarkup{ InlineKeyboard: replyButtons, }, }, } }