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 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, 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("2006-01-02") 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) 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 "reporttoggle": h.handleActionCallback(ctx, chatID, msgID, cb.ID, h.toggleReport) 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) default: // Delegate prefixed callbacks to the appropriate sub-handler if len(cb.Data) > 9 && cb.Data[:9] == "worktype_" { h.selectWorkType(ctx, chatID, msgID, cb.ID, cb.Data[9:]) return } if len(cb.Data) > 7 && cb.Data[:7] == "accent_" { h.selectAccent(ctx, chatID, msgID, cb.ID, cb.Data[7:]) return } if len(cb.Data) > 8 && cb.Data[:8] == "caltype_" { h.selectCalendar(ctx, chatID, msgID, cb.ID, cb.Data[8:]) return } if len(cb.Data) > 4 && cb.Data[:4] == "exp_" { h.handleExportCallback(ctx, chatID, msgID, cb.ID, cb.Data) return } h.handleHistoryCallback(ctx, 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(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 (current month if omitted) /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 (current month if omitted) /reporttoggle - Enable/disable auto report /setreporttime HH:MM - Set daily report time /setaccent - Select accent color /settimezone - Set your timezone (e.g. Asia/Tehran) /setbreak - Set minimum break threshold in minutes Dates use your calendar type (Gregorian/Jalali/Hijri). Buttons on the main menu also work.` h.sendText(ctx, msg.Chat.ID, text) } // 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("2006-01-02") 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("15:04") if t == timeStr { targetEvent = &events[i] break } } if targetEvent == nil { h.sendText(ctx, msg.Chat.ID, "No event found at that time") return } note := 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: "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 / 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(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) } }