package bot import ( "context" "fmt" "log/slog" "strings" "sync" "time" bot "github.com/go-telegram/bot" "github.com/go-telegram/bot/models" "worktimeBot/internal/db" ) const ( rateLimitInterval = 1 * time.Second DateLayout = "2006-01-02" TimeLayout = "15:04" secondsPerHour = 3600 secondsPerDay = 86400 defaultBreakThreshold = 300 maxBreakThresholdMins = 480 streakBonusCap = 500 xpCurveMultiplier = 50 salaryRoundFactor = 100 maxHourlyRate = 999999 ) type pendingNote struct { eventID int64 createdAt time.Time } // Handler holds the bot instance, database store, user allowlist, and rate-limit state. type Handler struct { Bot *bot.Bot DB *db.Store AllowedUsers map[int64]bool lastMsgTime map[int64]time.Time pendingNotes map[int64]*pendingNote mu sync.Mutex } // NewHandler creates a Handler and starts the rate-limit cleanup goroutine. func NewHandler(bot *bot.Bot, store *db.Store, allowed map[int64]bool) *Handler { h := &Handler{ Bot: bot, DB: store, AllowedUsers: allowed, lastMsgTime: make(map[int64]time.Time), pendingNotes: make(map[int64]*pendingNote), } go h.periodicCleanup() return h } // periodicCleanup runs in the background, purging stale rate-limit entries and pending notes 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) } } noteCutoff := time.Now().Add(-5 * time.Minute) for chatID, pn := range h.pendingNotes { if pn.createdAt.Before(noteCutoff) { delete(h.pendingNotes, chatID) } } 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(ctx context.Context, msg *models.Message, fn actionFunc) { text, err := fn(msg.Chat.ID) if err != nil { h.sendText(ctx, msg.Chat.ID, err.Error()) return } kb := mainKeyboard() h.sendWithKB(ctx, msg.Chat.ID, text, kb) } // handleActionCallback runs an action function and edits the callback message with the result. func (h *Handler) handleActionCallback(ctx context.Context, chatID int64, msgID int, callbackID string, fn actionFunc) { defer h.answerCb(ctx, callbackID) text, err := fn(chatID) if err != nil { text = err.Error() } kb := backKeyboard() h.editText(ctx, chatID, msgID, text, &kb) } // HandleMessage routes incoming text messages to the appropriate handler. func (h *Handler) HandleMessage(ctx context.Context, msg *models.Message) { if len(h.AllowedUsers) > 0 && !h.AllowedUsers[msg.Chat.ID] { return } if msg.Text == "" { 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) // Check for pending note — intercepts plain text only, commands always pass through. h.mu.Lock() pn, hasPN := h.pendingNotes[msg.Chat.ID] if hasPN { delete(h.pendingNotes, msg.Chat.ID) } h.mu.Unlock() if hasPN && msg.Text[0] != '/' { loc := loadLocation(user.Timezone) event, err := h.getEventByID(user.ID, pn.eventID) if err == nil { if err := h.DB.SetEventNote(pn.eventID, SanitizeNote(msg.Text)); err != nil { slog.Error("failed to save event note", "event_id", pn.eventID, "error", err) } date := time.Unix(event.OccurredAt, 0).In(loc).Format(DateLayout) h.sendDayView(ctx, msg.Chat.ID, 0, user, date, loc, true) } return } if msg.Text[0] != '/' { h.sendText(ctx, msg.Chat.ID, "Unknown command") return } cmd := msg.Text if i := strings.IndexByte(msg.Text, ' '); i >= 0 { cmd = msg.Text[:i] } switch cmd { case "/start": h.handleStart(ctx, msg) case "/help": h.handleHelp(ctx, msg) case "/clockin": h.handleActionMsg(ctx, msg, h.clockIn) case "/clockout": h.handleActionMsg(ctx, msg, h.clockOut) case "/worktype": h.handleWorkTypeMsg(ctx, msg) case "/report": h.handleActionMsg(ctx, msg, h.report) case "/export": h.handleExport(ctx, msg) case "/dayoff": h.handleActionMsg(ctx, msg, h.dayOff) case "/history": h.handleHistoryMsg(ctx, msg) case "/edit": h.handleEditMsg(ctx, msg) case "/reporttoggle": h.handleActionMsg(ctx, msg, h.toggleReport) case "/setreporttime": h.handleSetReportTime(ctx, msg) case "/setaccent": h.handleAccentMsg(ctx, msg) case "/settimezone": h.handleSetTimezone(ctx, msg) case "/note": h.handleNoteMsg(ctx, msg) case "/summary": h.handleSummaryMsg(ctx, msg) case "/setbreak": h.handleSetBreakMsg(ctx, msg) case "/rpg": h.handleActionMsg(ctx, msg, h.rpgStatus) case "/achievements": h.handleActionMsg(ctx, msg, h.achievementsList) case "/salary": h.handleActionMsg(ctx, msg, h.salaryEstimate) case "/burnout": h.handleActionMsg(ctx, msg, h.burnoutAssessment) case "/league": h.handleLeagueMsg(ctx, msg) case "/setrate": h.handleSetRateMsg(ctx, msg) case "/setcurrency": h.handleSetCurrencyMsg(ctx, msg) case "/rpgtoggle": h.handleRPGSettingsMsg(ctx, msg) case "/leaguetoggle": h.handleLeagueSettingsMsg(ctx, msg) case "/setusername": h.handleSetUsername(ctx, msg) default: h.sendText(ctx, msg.Chat.ID, "Unknown command") } } // HandleCallback routes inline callback queries to the appropriate handler. func (h *Handler) HandleCallback(ctx context.Context, cb *models.CallbackQuery) { if len(h.AllowedUsers) > 0 && !h.AllowedUsers[cb.Message.Message.Chat.ID] { h.answerCbWithText(ctx, cb.ID, "Unauthorized") return } user, err := h.getOrCreateUser(cb.Message.Message.Chat.ID) if err != nil { h.answerCb(ctx, cb.ID) return } if h.isRateLimited(user.ID) { h.answerCb(ctx, cb.ID) return } slog.Info("callback", "chat_id", cb.Message.Message.Chat.ID, "data", cb.Data) chatID := cb.Message.Message.Chat.ID msgID := cb.Message.Message.ID switch cb.Data { case "clockin": h.handleActionCallback(ctx, chatID, msgID, cb.ID, h.clockIn) case "clockout": h.handleActionCallback(ctx, chatID, msgID, cb.ID, h.clockOut) case "worktype": h.workTypeCallback(ctx, chatID, msgID, cb.ID) case "report": h.handleActionCallback(ctx, chatID, msgID, cb.ID, h.report) case "export": h.exportCallback(ctx, chatID, msgID, cb.ID) case "noop": h.answerCb(ctx, cb.ID) case "dayoff": h.handleActionCallback(ctx, chatID, msgID, cb.ID, h.dayOff) case "settings": h.settingsCallback(ctx, chatID, msgID, cb.ID) case "caltype": h.calTypeCallback(ctx, chatID, msgID, cb.ID) case "accent": h.accentCallback(ctx, chatID, msgID, cb.ID) case "timezone": h.timezoneCallback(ctx, chatID, msgID, cb.ID) case "history": h.historyCallback(ctx, chatID, msgID, cb.ID) case "reporttime": h.reportTimeCallback(ctx, chatID, msgID, cb.ID) case "breakthreshold": h.breakThresholdCallback(ctx, chatID, msgID, cb.ID) case "back_menu": h.backToMenu(ctx, chatID, msgID, cb.ID) case "back_settings": h.backToSettings(ctx, chatID, msgID, cb.ID) case "rpg": h.handleActionCallback(ctx, chatID, msgID, cb.ID, h.rpgStatus) case "achievements": h.handleActionCallback(ctx, chatID, msgID, cb.ID, h.achievementsList) case "salary": h.handleActionCallback(ctx, chatID, msgID, cb.ID, h.salaryEstimate) case "burnout": h.handleActionCallback(ctx, chatID, msgID, cb.ID, h.burnoutAssessment) case "league": h.handleLeagueCallback(ctx, chatID, msgID, cb.ID, "xp") case "league_xp": h.handleLeagueCallback(ctx, chatID, msgID, cb.ID, "xp") case "league_level": h.handleLeagueCallback(ctx, chatID, msgID, cb.ID, "level") case "league_streak": h.handleLeagueCallback(ctx, chatID, msgID, cb.ID, "streak") case "league_hours": h.handleLeagueCallback(ctx, chatID, msgID, cb.ID, "hours") case "rpgtoggle": h.rpgToggleCallback(ctx, chatID, msgID, cb.ID) case "leaguetoggle": h.leagueToggleCallback(ctx, chatID, msgID, cb.ID) case "reporttoggle": h.reportToggleCallback(ctx, chatID, msgID, cb.ID) case "currency": h.currencyCallback(ctx, chatID, msgID, cb.ID) case "setrate": h.rateCallback(ctx, chatID, msgID, cb.ID) default: h.handlePrefixedCallback(ctx, chatID, msgID, cb.ID, cb.Data) } } func (h *Handler) handlePrefixedCallback(ctx context.Context, chatID int64, msgID int, callbackID, data string) { if len(data) > 9 && data[:9] == "worktype_" { h.selectWorkType(ctx, chatID, msgID, callbackID, data[9:]) return } if len(data) > 7 && data[:7] == "accent_" { h.selectAccent(ctx, chatID, msgID, callbackID, data[7:]) return } if len(data) > 8 && data[:8] == "caltype_" { h.selectCalendar(ctx, chatID, msgID, callbackID, data[8:]) return } if len(data) > 4 && data[:4] == "exp_" { h.handleExportCallback(ctx, chatID, msgID, callbackID, data) return } if len(data) > 15 && data[:15] == "breakthreshold_" { h.selectBreakThreshold(ctx, chatID, msgID, callbackID, data[15:]) return } if len(data) > 9 && data[:9] == "currency_" { h.selectCurrency(ctx, chatID, msgID, callbackID, data[9:]) return } switch data { case "rpg_enable": h.handleSelectCallback(ctx, chatID, msgID, callbackID, "rpg_enable", true) case "rpg_disable": h.handleSelectCallback(ctx, chatID, msgID, callbackID, "rpg_enable", false) case "league_enable": h.handleSelectCallback(ctx, chatID, msgID, callbackID, "league_enable", true) case "league_disable": h.handleSelectCallback(ctx, chatID, msgID, callbackID, "league_enable", false) case "report_enable": h.handleSelectCallback(ctx, chatID, msgID, callbackID, "report_enable", true) case "report_disable": h.handleSelectCallback(ctx, chatID, msgID, callbackID, "report_enable", false) default: h.handleHistoryCallback(ctx, chatID, msgID, callbackID, 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(DateLayout) 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(ctx context.Context, msg *models.Message) { user, err := h.getOrCreateUser(msg.Chat.ID) if err != nil { slog.Error("start: create user", "error", err) h.sendText(ctx, msg.Chat.ID, "Error setting up your profile") return } slog.Info("start", "user_id", user.ID, "chat_id", msg.Chat.ID) kb := mainKeyboard() h.sendWithKB(ctx, msg.Chat.ID, h.buildStatusText(msg.Chat.ID), kb) } // handleHelp shows the list of available commands. func (h *Handler) handleHelp(ctx context.Context, msg *models.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 /history - View calendar history /edit YYYY-MM-DD - Edit a specific day /note [YYYY-MM-DD] HH:MM - Set note for an event /summary [YYYY-MM] - Show monthly summary /reporttoggle - Enable/disable auto report /setreporttime HH:MM - Set daily report time /setaccent - Select accent color /settimezone - Set timezone (e.g. Asia/Tehran) /setbreak - Set break threshold in minutes /setusername - Set your display name /rpg - Show RPG stats and XP /achievements - View unlocked achievements /salary - Show salary estimate /burnout - View burnout assessment /league - View WorkTime League rankings /setrate - Set hourly rate (e.g. 25.50) /setcurrency - Set currency (e.g. $ USD) /rpgtoggle - Enable/disable RPG system /leaguetoggle - Enable/disable league participation Dates use your calendar type (Gregorian/Jalali/Hijri). Buttons on the main menu also work.` h.sendText(ctx, msg.Chat.ID, text) } // handleSetUsername handles /setusername . func (h *Handler) handleSetUsername(ctx context.Context, msg *models.Message) { name := strings.TrimSpace(strings.TrimPrefix(msg.Text, "/setusername")) if name == "" { user, err := h.getOrCreateUser(msg.Chat.ID) if err != nil { h.sendText(ctx, msg.Chat.ID, "Error loading profile") return } h.sendText(ctx, msg.Chat.ID, fmt.Sprintf("Current username: %s\nUsage: /setusername ", user.Username)) return } name = SanitizeDisplayName(name) if name == "" { h.sendText(ctx, msg.Chat.ID, "Invalid username.") return } user, err := h.getOrCreateUser(msg.Chat.ID) if err != nil { h.sendText(ctx, msg.Chat.ID, "Error loading profile") return } user.Username = name if err := h.DB.UpdateUser(user); err != nil { h.sendText(ctx, msg.Chat.ID, "Error saving username") return } h.sendText(ctx, msg.Chat.ID, fmt.Sprintf("Username set to: %s", name)) } // handleNoteMsg processes the /note command: /note [YYYY-MM-DD] HH:MM func (h *Handler) handleNoteMsg(ctx context.Context, msg *models.Message) { user, err := h.getOrCreateUser(msg.Chat.ID) if err != nil { h.sendText(ctx, msg.Chat.ID, "Error loading profile") return } loc := loadLocation(user.Timezone) parts := strings.Fields(msg.Text) if len(parts) < 3 { h.sendText(ctx, msg.Chat.ID, "Usage: /note [YYYY-MM-DD] HH:MM ") return } var dateStr string var timeStr string var noteParts []string // Check if first argument looks like a date if len(parts[1]) == 10 && parts[1][4] == '-' && parts[1][7] == '-' { dateStr = parts[1] timeStr = parts[2] noteParts = parts[3:] } else { dateStr = time.Now().In(loc).Format(DateLayout) timeStr = parts[1] noteParts = parts[2:] } events, err := h.DB.EventsForDayByDate(user.ID, dateStr) if err != nil || len(events) == 0 { h.sendText(ctx, msg.Chat.ID, "No events found for that day") return } var targetEvent *db.Event for i := range events { t := time.Unix(events[i].OccurredAt, 0).In(loc).Format(TimeLayout) if t == timeStr { targetEvent = &events[i] break } } if targetEvent == nil { h.sendText(ctx, msg.Chat.ID, "No event found at that time") return } note := SanitizeNote(strings.Join(noteParts, " ")) if err := h.DB.SetEventNote(targetEvent.ID, note); err != nil { h.sendText(ctx, msg.Chat.ID, "Error saving note") return } h.sendText(ctx, msg.Chat.ID, "Note saved") } // handleWorkTypeMsg shows the work type picker for today. func (h *Handler) handleWorkTypeMsg(ctx context.Context, msg *models.Message) { user, day, err := h.getUserToday(msg.Chat.ID) if err != nil { h.sendText(ctx, msg.Chat.ID, "Error loading work types") return } text, kb := h.buildWorkTypeKeyboard(user, day) h.sendWithKB(ctx, msg.Chat.ID, text, kb) } // workTypeCallback shows the work type picker (inline edit). func (h *Handler) workTypeCallback(ctx context.Context, chatID int64, msgID int, callbackID string) { defer h.answerCb(ctx, callbackID) user, day, err := h.getUserToday(chatID) if err != nil { kb := backKeyboard() h.editText(ctx, chatID, msgID, "Error loading work types", &kb) return } text, kb := h.buildWorkTypeKeyboard(user, day) h.editText(ctx, chatID, msgID, text, &kb) } // buildWorkTypeKeyboard returns the work type selection keyboard. func (h *Handler) buildWorkTypeKeyboard(user *db.User, day *db.Day) (string, models.InlineKeyboardMarkup) { wts, err := h.DB.GetWorkTypes() if err != nil || len(wts) == 0 { return "No work types available.", backKeyboard() } rows := [][]models.InlineKeyboardButton{} for _, wt := range wts { label := wt.Name if wt.ID == day.CurrentWorkTypeID { label = "> " + wt.Name } rows = append(rows, []models.InlineKeyboardButton{ {Text: label, CallbackData: fmt.Sprintf("worktype_%d", wt.ID)}, }) } rows = append(rows, []models.InlineKeyboardButton{ {Text: "Back to Menu", CallbackData: "back_menu"}, }) return "Select work type:", models.InlineKeyboardMarkup{InlineKeyboard: rows} } // selectWorkType updates the day's work type from a callback. func (h *Handler) selectWorkType(ctx context.Context, chatID int64, msgID int, callbackID, idStr string) { defer h.answerCb(ctx, 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 } kb := backKeyboard() h.editText(ctx, chatID, msgID, fmt.Sprintf("Work type set to: %s", wt.Name), &kb) } // mainKeyboard returns the persistent inline menu. func mainKeyboard() models.InlineKeyboardMarkup { return models.InlineKeyboardMarkup{ InlineKeyboard: [][]models.InlineKeyboardButton{ { {Text: "Clock In", CallbackData: "clockin"}, {Text: "Clock Out", CallbackData: "clockout"}, }, { {Text: "Work Type", CallbackData: "worktype"}, {Text: "Day Off", CallbackData: "dayoff"}, }, { {Text: "Report", CallbackData: "report"}, {Text: "Export", CallbackData: "export"}, {Text: "History", CallbackData: "history"}, }, { {Text: "RPG", CallbackData: "rpg"}, {Text: "League", CallbackData: "league"}, {Text: "Burnout", CallbackData: "burnout"}, }, { {Text: "Salary", CallbackData: "salary"}, {Text: "Settings", CallbackData: "settings"}, }, }, } } // backKeyboard returns a simple "Back to Menu" button. func backKeyboard() models.InlineKeyboardMarkup { return models.InlineKeyboardMarkup{ InlineKeyboard: [][]models.InlineKeyboardButton{ {{Text: "Back to Menu", CallbackData: "back_menu"}}, }, } } // formatDuration formats seconds as a human-readable duration (e.g. "7h 30m"). func formatDuration(seconds int64) string { hours := seconds / secondsPerHour mins := (seconds % secondsPerHour) / 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(ctx context.Context, chatID int64, text string) { if _, err := h.Bot.SendMessage(ctx, &bot.SendMessageParams{ChatID: chatID, Text: text}); err != nil { slog.Warn("send text failed", "chat_id", chatID, "error", err) } } // editText edits an existing message, logging errors. func (h *Handler) editText(ctx context.Context, chatID int64, msgID int, text string, kb *models.InlineKeyboardMarkup) { p := &bot.EditMessageTextParams{ChatID: chatID, MessageID: msgID, Text: text} if kb != nil { p.ReplyMarkup = kb } if _, err := h.Bot.EditMessageText(ctx, p); err != nil { slog.Warn("edit text failed", "chat_id", chatID, "msg_id", msgID, "error", err) } } // sendWithKB sends a message with an inline keyboard, logging errors. func (h *Handler) sendWithKB(ctx context.Context, chatID int64, text string, kb models.InlineKeyboardMarkup) { if _, err := h.Bot.SendMessage(ctx, &bot.SendMessageParams{ChatID: chatID, Text: text, ReplyMarkup: kb}); err != nil { slog.Warn("send with keyboard failed", "chat_id", chatID, "error", err) } } // answerCb answers a callback query, logging errors. func (h *Handler) answerCb(ctx context.Context, callbackID string) { if _, err := h.Bot.AnswerCallbackQuery(ctx, &bot.AnswerCallbackQueryParams{CallbackQueryID: callbackID}); err != nil { slog.Debug("answer callback failed", "error", err) } } // answerCbWithText answers a callback query with a toast message, logging errors. func (h *Handler) answerCbWithText(ctx context.Context, callbackID, text string) { if _, err := h.Bot.AnswerCallbackQuery(ctx, &bot.AnswerCallbackQueryParams{CallbackQueryID: callbackID, Text: text}); err != nil { slog.Debug("answer callback with text failed", "error", err) } } // 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(ctx context.Context, chatID int64, msgID int, text string, kb *models.InlineKeyboardMarkup) { if msgID == 0 { if kb != nil { h.sendWithKB(ctx, chatID, text, *kb) } else { h.sendText(ctx, chatID, text) } } else { h.editText(ctx, chatID, msgID, text, kb) } }