package bot import ( "fmt" "log/slog" "sync" "time" tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5" "worktimeBot/internal/db" ) const rateLimitInterval = 1 * time.Second // Handler holds the bot instance, database store, user allowlist, and rate-limit state. type Handler struct { Bot *tgbotapi.BotAPI DB *db.Store AllowedUsers map[int64]bool lastMsgTime map[int64]time.Time mu sync.Mutex } // NewHandler creates a Handler and starts the rate-limit cleanup goroutine. func NewHandler(bot *tgbotapi.BotAPI, store *db.Store, allowed map[int64]bool) *Handler { h := &Handler{ Bot: bot, DB: store, AllowedUsers: allowed, lastMsgTime: make(map[int64]time.Time), } go h.periodicCleanup() return h } // periodicCleanup runs in the background, purging stale rate-limit entries every 10 seconds. func (h *Handler) periodicCleanup() { for { time.Sleep(10 * time.Second) h.mu.Lock() cutoff := time.Now().Add(-rateLimitInterval * 2) for uid, t := range h.lastMsgTime { if t.Before(cutoff) { delete(h.lastMsgTime, uid) } } h.mu.Unlock() } } type actionFunc func(int64) (string, error) // isRateLimited returns true if the user has sent a message within the last rateLimitInterval. func (h *Handler) isRateLimited(userID int64) bool { h.mu.Lock() defer h.mu.Unlock() now := time.Now() if last, ok := h.lastMsgTime[userID]; ok && now.Sub(last) < rateLimitInterval { return true } h.lastMsgTime[userID] = now return false } // handleActionMsg runs an action function and sends the result as a new message. func (h *Handler) handleActionMsg(msg *tgbotapi.Message, fn actionFunc) { text, err := fn(msg.Chat.ID) if err != nil { h.sendText(msg.Chat.ID, err.Error()) return } reply := tgbotapi.NewMessage(msg.Chat.ID, text) reply.ReplyMarkup = mainKeyboard() h.Bot.Send(reply) } // handleActionCallback runs an action function and edits the callback message with the result. func (h *Handler) handleActionCallback(chatID int64, msgID int, callbackID string, fn actionFunc) { defer h.Bot.Request(tgbotapi.NewCallback(callbackID, "")) text, err := fn(chatID) if err != nil { text = err.Error() } edit := tgbotapi.NewEditMessageText(chatID, msgID, text) kb := backKeyboard() edit.ReplyMarkup = &kb h.Bot.Send(edit) } // HandleMessage routes incoming text messages to the appropriate handler. func (h *Handler) HandleMessage(update tgbotapi.Update) { msg := update.Message if len(h.AllowedUsers) > 0 && !h.AllowedUsers[msg.Chat.ID] { return } user, err := h.getOrCreateUser(msg.Chat.ID) if err != nil { return } if h.isRateLimited(user.ID) { return } slog.Info("message", "chat_id", msg.Chat.ID, "text", msg.Text) switch msg.Text { case "/start": h.handleStart(msg) case "/help": h.handleHelp(msg) case "/clockin": h.handleActionMsg(msg, h.clockIn) case "/clockout": h.handleActionMsg(msg, h.clockOut) case "/worktype": h.handleWorkTypeMsg(msg) case "/report": h.handleActionMsg(msg, h.report) case "/export": h.handleExport(msg) case "/dayoff": h.handleActionMsg(msg, h.dayOff) case "/history": h.handleHistoryMsg(msg) case "/edit": h.handleEditMsg(msg) case "/reporttoggle": h.handleActionMsg(msg, h.toggleReport) case "/setreporttime": h.handleSetReportTime(msg) case "/setaccent": h.handleAccentMsg(msg) case "/settimezone": h.handleSetTimezone(msg) default: h.sendText(msg.Chat.ID, "Unknown command") } } // HandleCallback routes inline callback queries to the appropriate handler. func (h *Handler) HandleCallback(update tgbotapi.Update) { cb := update.CallbackQuery if len(h.AllowedUsers) > 0 && !h.AllowedUsers[cb.Message.Chat.ID] { h.Bot.Request(tgbotapi.NewCallback(cb.ID, "Unauthorized")) return } user, err := h.getOrCreateUser(cb.Message.Chat.ID) if err != nil { h.Bot.Request(tgbotapi.NewCallback(cb.ID, "")) return } if h.isRateLimited(user.ID) { h.Bot.Request(tgbotapi.NewCallback(cb.ID, "")) return } slog.Info("callback", "chat_id", cb.Message.Chat.ID, "data", cb.Data) chatID := cb.Message.Chat.ID msgID := cb.Message.MessageID switch cb.Data { case "clockin": h.handleActionCallback(chatID, msgID, cb.ID, h.clockIn) case "clockout": h.handleActionCallback(chatID, msgID, cb.ID, h.clockOut) case "worktype": h.workTypeCallback(chatID, msgID, cb.ID) case "report": h.handleActionCallback(chatID, msgID, cb.ID, h.report) case "export": h.exportCallback(chatID, msgID, cb.ID) case "noop": h.Bot.Request(tgbotapi.NewCallback(cb.ID, "")) case "dayoff": h.handleActionCallback(chatID, msgID, cb.ID, h.dayOff) case "settings": h.settingsCallback(chatID, msgID, cb.ID) case "caltype": h.calTypeCallback(chatID, msgID, cb.ID) case "accent": h.accentCallback(chatID, msgID, cb.ID) case "timezone": h.timezoneCallback(chatID, msgID, cb.ID) case "history": h.historyCallback(chatID, msgID, cb.ID) case "reporttoggle": h.handleActionCallback(chatID, msgID, cb.ID, h.toggleReport) case "reporttime": h.reportTimeCallback(chatID, msgID, cb.ID) case "back_menu": h.backToMenu(chatID, msgID, cb.ID) case "back_settings": h.backToSettings(chatID, msgID, cb.ID) default: // Delegate prefixed callbacks to the appropriate sub-handler if len(cb.Data) > 9 && cb.Data[:9] == "worktype_" { h.selectWorkType(chatID, msgID, cb.ID, cb.Data[9:]) return } if len(cb.Data) > 7 && cb.Data[:7] == "accent_" { h.selectAccent(chatID, msgID, cb.ID, cb.Data[7:]) return } if len(cb.Data) > 8 && cb.Data[:8] == "caltype_" { h.selectCalendar(chatID, msgID, cb.ID, cb.Data[8:]) return } if len(cb.Data) > 4 && cb.Data[:4] == "exp_" { h.handleExportCallback(chatID, msgID, cb.ID, cb.Data) return } h.handleHistoryCallback(chatID, msgID, cb.ID, cb.Data) } } // getOrCreateUser is a shortcut that delegates to the store. func (h *Handler) getOrCreateUser(chatID int64) (*db.User, error) { return h.DB.GetOrCreateUser(chatID) } // getUserToday fetches the user and today's day record, using the user's timezone. func (h *Handler) getUserToday(chatID int64) (*db.User, *db.Day, error) { user, err := h.getOrCreateUser(chatID) if err != nil { return nil, nil, err } loc, err := time.LoadLocation(user.Timezone) if err != nil { loc = time.UTC } today := time.Now().In(loc).Format("2006-01-02") day, err := h.DB.GetOrCreateDay(user.ID, today) if err != nil { return nil, nil, err } return user, day, nil } // handleStart shows the main menu with the user's current status. func (h *Handler) handleStart(msg *tgbotapi.Message) { user, err := h.getOrCreateUser(msg.Chat.ID) if err != nil { slog.Error("start: create user", "error", err) h.sendText(msg.Chat.ID, "Error setting up your profile") return } slog.Info("start", "user_id", user.ID, "chat_id", msg.Chat.ID) menu := tgbotapi.NewMessage(msg.Chat.ID, h.buildStatusText(msg.Chat.ID)) menu.ReplyMarkup = mainKeyboard() h.Bot.Send(menu) } // handleHelp shows the list of available commands. func (h *Handler) handleHelp(msg *tgbotapi.Message) { text := `Commands: /start - Show main menu /clockin - Clock in /clockout - Clock out /worktype - Select work type /dayoff - Toggle day off /report - Show today report /export [YYYY-MM] - Export monthly report (current month if omitted) /history - View calendar history /edit YYYY-MM-DD - Edit a specific day /reporttoggle - Enable/disable auto report /setreporttime HH:MM - Set daily report time /setaccent - Select accent color /settimezone - Set your timezone (e.g. Asia/Tehran) Dates use your calendar type (Gregorian/Jalali/Hijri). Buttons on the main menu also work.` h.sendText(msg.Chat.ID, text) } // handleWorkTypeMsg shows the work type picker for today. func (h *Handler) handleWorkTypeMsg(msg *tgbotapi.Message) { user, day, err := h.getUserToday(msg.Chat.ID) if err != nil { h.sendText(msg.Chat.ID, "Error loading work types") return } text, kb := h.buildWorkTypeKeyboard(user, day) reply := tgbotapi.NewMessage(msg.Chat.ID, text) reply.ReplyMarkup = &kb h.Bot.Send(reply) } // workTypeCallback shows the work type picker (inline edit). func (h *Handler) workTypeCallback(chatID int64, msgID int, callbackID string) { defer h.Bot.Request(tgbotapi.NewCallback(callbackID, "")) user, day, err := h.getUserToday(chatID) if err != nil { edit := tgbotapi.NewEditMessageText(chatID, msgID, "Error loading work types") kb := backKeyboard() edit.ReplyMarkup = &kb h.Bot.Send(edit) return } text, kb := h.buildWorkTypeKeyboard(user, day) edit := tgbotapi.NewEditMessageText(chatID, msgID, text) edit.ReplyMarkup = &kb h.Bot.Send(edit) } // buildWorkTypeKeyboard returns the work type selection keyboard. func (h *Handler) buildWorkTypeKeyboard(user *db.User, day *db.Day) (string, tgbotapi.InlineKeyboardMarkup) { wts, err := h.DB.GetWorkTypes() if err != nil || len(wts) == 0 { return "No work types available.", backKeyboard() } rows := [][]tgbotapi.InlineKeyboardButton{} for _, wt := range wts { label := wt.Name if wt.ID == day.CurrentWorkTypeID { label = "> " + wt.Name } rows = append(rows, tgbotapi.NewInlineKeyboardRow( tgbotapi.NewInlineKeyboardButtonData(label, fmt.Sprintf("worktype_%d", wt.ID)), )) } rows = append(rows, tgbotapi.NewInlineKeyboardRow( tgbotapi.NewInlineKeyboardButtonData("Back to Menu", "back_menu"), )) return "Select work type:", tgbotapi.NewInlineKeyboardMarkup(rows...) } // selectWorkType updates the day's work type from a callback. func (h *Handler) selectWorkType(chatID int64, msgID int, callbackID, idStr string) { defer h.Bot.Request(tgbotapi.NewCallback(callbackID, "")) _, day, err := h.getUserToday(chatID) if err != nil { return } var wtID int64 if _, err := fmt.Sscanf(idStr, "%d", &wtID); err != nil { return } wt, err := h.DB.GetWorkType(wtID) if err != nil { return } day.CurrentWorkTypeID = wtID if err := h.DB.UpdateDay(day); err != nil { return } edit := tgbotapi.NewEditMessageText(chatID, msgID, fmt.Sprintf("Work type set to: %s", wt.Name)) kb := backKeyboard() edit.ReplyMarkup = &kb h.Bot.Send(edit) } // mainKeyboard returns the persistent inline menu. func mainKeyboard() tgbotapi.InlineKeyboardMarkup { return tgbotapi.NewInlineKeyboardMarkup( tgbotapi.NewInlineKeyboardRow( tgbotapi.NewInlineKeyboardButtonData("Clock In", "clockin"), tgbotapi.NewInlineKeyboardButtonData("Clock Out", "clockout"), ), tgbotapi.NewInlineKeyboardRow( tgbotapi.NewInlineKeyboardButtonData("Work Type", "worktype"), tgbotapi.NewInlineKeyboardButtonData("Day Off", "dayoff"), ), tgbotapi.NewInlineKeyboardRow( tgbotapi.NewInlineKeyboardButtonData("Report", "report"), tgbotapi.NewInlineKeyboardButtonData("Export", "export"), ), tgbotapi.NewInlineKeyboardRow( tgbotapi.NewInlineKeyboardButtonData("Settings", "settings"), ), ) } // backKeyboard returns a simple "Back to Menu" button. func backKeyboard() tgbotapi.InlineKeyboardMarkup { return tgbotapi.NewInlineKeyboardMarkup( tgbotapi.NewInlineKeyboardRow( tgbotapi.NewInlineKeyboardButtonData("Back to Menu", "back_menu"), ), ) } // formatDuration formats seconds as a human-readable duration (e.g. "7h 30m"). func formatDuration(seconds int64) string { hours := seconds / 3600 mins := (seconds % 3600) / 60 if hours > 0 { return fmt.Sprintf("%dh %dm", hours, mins) } return fmt.Sprintf("%dm", mins) } // sendText sends a plain text message to the given chat. func (h *Handler) sendText(chatID int64, text string) { h.Bot.Send(tgbotapi.NewMessage(chatID, text)) } // navMonth returns the (prevY, prevM, nextY, nextM) for month navigation. func navMonth(year, month int) (prevY, prevM, nextY, nextM int) { prevY, prevM = year, month-1 if prevM < 1 { prevM = 12 prevY-- } nextY, nextM = year, month+1 if nextM > 12 { nextM = 1 nextY++ } return } // sendOrEdit sends a new message or edits an existing one depending on msgID. func (h *Handler) sendOrEdit(chatID int64, msgID int, text string, kb *tgbotapi.InlineKeyboardMarkup) { if msgID == 0 { msg := tgbotapi.NewMessage(chatID, text) msg.ReplyMarkup = kb h.Bot.Send(msg) } else { edit := tgbotapi.NewEditMessageText(chatID, msgID, text) edit.ReplyMarkup = kb h.Bot.Send(edit) } }