package handler import ( "context" "fmt" "log/slog" "os" "path/filepath" "strconv" "strings" "sync" "time" tgbot "github.com/go-telegram/bot" "github.com/go-telegram/bot/models" "github.com/google/uuid" "uptodownBot/internal/dl" "uptodownBot/internal/domain" "uptodownBot/internal/repo" ) type PendingKind string const ( PendingURL PendingKind = "url" ) type PendingInput struct { Kind PendingKind UserID int64 CreatedAt time.Time MsgID int Data map[string]string } type Handler struct { Bot *tgbot.Bot Repo repo.Repository YtDlp *dl.Client AllowedUsers map[int64]bool activeJobs map[int64]*domain.DownloadJob stepStates map[int64]*stepState lastMsgTime map[int64]time.Time pendingInput map[int64]*PendingInput mu sync.Mutex downloadDir string maxFileSize int64 } func NewHandler(bot *tgbot.Bot, store repo.Repository, ytdlp *dl.Client, allowed map[int64]bool, downloadDir string, maxFileSize int64) *Handler { h := &Handler{ Bot: bot, Repo: store, YtDlp: ytdlp, AllowedUsers: allowed, activeJobs: make(map[int64]*domain.DownloadJob), stepStates: make(map[int64]*stepState), lastMsgTime: make(map[int64]time.Time), pendingInput: make(map[int64]*PendingInput), downloadDir: downloadDir, maxFileSize: maxFileSize, } go h.periodicCleanup() return h } func (h *Handler) periodicCleanup() { for { time.Sleep(domain.CleanupInterval) h.mu.Lock() // Clean rate limit entries cutoff := time.Now().Add(-domain.RateLimitInterval * 2) for uid, t := range h.lastMsgTime { if t.Before(cutoff) { delete(h.lastMsgTime, uid) } } // Clean stale pending inputs pendingCutoff := time.Now().Add(-5 * time.Minute) for chatID, pi := range h.pendingInput { if pi.CreatedAt.Before(pendingCutoff) { delete(h.pendingInput, chatID) } } h.mu.Unlock() // Clean old temp files (>1h) h.cleanOldTempFiles() } } func (h *Handler) cleanOldTempFiles() { entries, err := os.ReadDir(h.downloadDir) if err != nil { return } cutoff := time.Now().Add(-domain.FileRetention) for _, entry := range entries { info, err := entry.Info() if err != nil { continue } if info.ModTime().Before(cutoff) { path := filepath.Join(h.downloadDir, entry.Name()) if err := os.Remove(path); err != nil { slog.Warn("failed to remove old temp file", "path", path, "error", err) } } } } 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) < domain.RateLimitInterval { return true } h.lastMsgTime[userID] = now return false } func (h *Handler) getOrCreateUser(chatID int64) (int64, error) { return h.Repo.GetOrCreateUser(chatID) } 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 } userID, err := h.getOrCreateUser(msg.Chat.ID) if err != nil { return } if h.isRateLimited(userID) { return } slog.Info("message", "chat_id", msg.Chat.ID, "text", msg.Text) if msg.Text[0] == '/' { h.mu.Lock() delete(h.pendingInput, msg.Chat.ID) h.mu.Unlock() h.routeCommand(ctx, msg, userID) return } h.mu.Lock() pi, hasPI := h.pendingInput[msg.Chat.ID] if hasPI { delete(h.pendingInput, msg.Chat.ID) } h.mu.Unlock() if hasPI { h.processPendingInput(ctx, msg, userID, pi) return } // Treat plain text as a potential URL h.handleURLInput(ctx, msg.Chat.ID, msg.Text, userID) } func (h *Handler) routeCommand(ctx context.Context, msg *models.Message, userID int64) { cmd := msg.Text if i := strings.IndexByte(msg.Text, ' '); i >= 0 { cmd = msg.Text[:i] } switch cmd { case "/start": h.handleStart(ctx, msg, userID) case "/settings": h.handleSettings(ctx, msg, userID) case "/help": h.handleHelp(ctx, msg) case "/history": h.handleHistory(ctx, msg, userID) case "/download": h.handleDownloadCommand(ctx, msg, userID) default: h.sendText(ctx, msg.Chat.ID, "Unknown command") } } func (h *Handler) HandleCallback(ctx context.Context, cb *models.CallbackQuery) { chatID := cb.Message.Message.Chat.ID if len(h.AllowedUsers) > 0 && !h.AllowedUsers[chatID] { h.answerCbWithText(ctx, cb.ID, "Unauthorized") return } userID, err := h.getOrCreateUser(chatID) if err != nil { h.answerCb(ctx, cb.ID) return } if h.isRateLimited(userID) { h.answerCb(ctx, cb.ID) return } slog.Info("callback", "chat_id", chatID, "data", cb.Data) h.mu.Lock() delete(h.pendingInput, chatID) h.mu.Unlock() msgID := cb.Message.Message.ID data := cb.Data switch { case data == "noop": h.answerCb(ctx, cb.ID) case data == "cancel_dl": h.handleCancelDownload(ctx, chatID, msgID, cb.ID, userID) case data == "back_step": h.handleBackStep(ctx, chatID, msgID, cb.ID, userID) case data == "back_menu": h.backToMenu(ctx, chatID, msgID, cb.ID, userID) case data == "download": h.handleDownloadPrompt(ctx, chatID, msgID, cb.ID, userID) h.answerCb(ctx, cb.ID) case data == "settings": h.settingsCallback(ctx, chatID, msgID, cb.ID, userID) case data == "help": h.handleHelp(ctx, cb.Message.Message) h.answerCb(ctx, cb.ID) case data == "history": h.historyCallback(ctx, chatID, msgID, cb.ID, userID) case data == "deleteaccount": h.deleteAccountPrompt(ctx, chatID, msgID, userID) h.answerCb(ctx, cb.ID) case data == "pending_cancel": h.backToMenu(ctx, chatID, msgID, cb.ID, userID) default: h.routePrefixedCallback(ctx, chatID, msgID, cb.ID, userID, data) } } func (h *Handler) routePrefixedCallback(ctx context.Context, chatID int64, msgID int, cbID string, userID int64, data string) { defer h.answerCb(ctx, cbID) switch { case strings.HasPrefix(data, "format_"): h.handleFormatSelection(ctx, chatID, msgID, userID, data[7:]) case strings.HasPrefix(data, "secondary_"): h.handleSecondarySelection(ctx, chatID, msgID, userID, data[10:]) case strings.HasPrefix(data, "lang_"): h.handleLanguageSelection(ctx, chatID, msgID, userID, data[5:]) case strings.HasPrefix(data, "container_"): h.handleContainerSelection(ctx, chatID, msgID, userID, data[10:]) case strings.HasPrefix(data, "pl_page_"): h.handlePlaylistPage(ctx, chatID, msgID, userID, data[8:]) case strings.HasPrefix(data, "pl_entry_"): h.handlePlaylistEntry(ctx, chatID, msgID, userID, data[9:]) case data == "pl_select_all": h.handlePlaylistSelectAll(ctx, chatID, msgID, userID) case data == "pl_confirm": h.handlePlaylistConfirm(ctx, chatID, msgID, userID) case strings.HasPrefix(data, "delaccount_"): if uid, err := strconv.ParseInt(data[11:], 10, 64); err == nil && uid > 0 { h.deleteAccountConfirm(ctx, chatID, msgID, uid) } case data == "back_quality": h.handleSettingsSelection(ctx, chatID, msgID, userID, "quality") case data == "back_container": h.handleSettingsSelection(ctx, chatID, msgID, userID, "container") case data == "back_settings": h.backToSettings(ctx, chatID, msgID, userID) case strings.HasPrefix(data, "settings_set_"): h.handleSettingsSet(ctx, chatID, msgID, userID, data[13:]) case strings.HasPrefix(data, "settings_"): h.handleSettingsSelection(ctx, chatID, msgID, userID, data[9:]) } } func (h *Handler) setPending(ctx context.Context, chatID int64, msgID int, userID int64, kind PendingKind, data map[string]string) { h.mu.Lock() h.pendingInput[chatID] = &PendingInput{ Kind: kind, UserID: userID, CreatedAt: time.Now(), MsgID: msgID, Data: data, } h.mu.Unlock() } func (h *Handler) processPendingInput(ctx context.Context, msg *models.Message, userID int64, pi *PendingInput) { text := strings.TrimSpace(msg.Text) switch pi.Kind { case PendingURL: h.handleURLInput(ctx, msg.Chat.ID, text, userID) default: h.sendText(ctx, msg.Chat.ID, "Unknown input type") } } func (h *Handler) promptForInput(ctx context.Context, chatID int64, msgID int, userID int64, kind PendingKind, prompt string, data map[string]string) { kb := models.InlineKeyboardMarkup{ InlineKeyboard: [][]models.InlineKeyboardButton{ {{Text: "Cancel", CallbackData: "pending_cancel"}}, }, } if msgID == 0 { msg, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{ ChatID: chatID, Text: prompt, ReplyMarkup: kb, }) if err != nil { slog.Warn("send prompt failed", "chat_id", chatID, "error", err) return } h.setPending(ctx, chatID, msg.ID, userID, kind, data) } else { h.editText(ctx, chatID, msgID, prompt, &kb) h.setPending(ctx, chatID, msgID, userID, kind, data) } } 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) } } // -- helpers -- func (h *Handler) sendText(ctx context.Context, chatID int64, text string) { if _, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{ChatID: chatID, Text: text}); err != nil { slog.Warn("send text failed", "chat_id", chatID, "error", err) } } func (h *Handler) editText(ctx context.Context, chatID int64, msgID int, text string, kb *models.InlineKeyboardMarkup) { p := &tgbot.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) } } func (h *Handler) sendWithKB(ctx context.Context, chatID int64, text string, kb models.InlineKeyboardMarkup) { if _, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{ChatID: chatID, Text: text, ReplyMarkup: kb}); err != nil { slog.Warn("send with keyboard failed", "chat_id", chatID, "error", err) } } func (h *Handler) answerCb(ctx context.Context, callbackID string) { if _, err := h.Bot.AnswerCallbackQuery(ctx, &tgbot.AnswerCallbackQueryParams{CallbackQueryID: callbackID}); err != nil { slog.Debug("answer callback failed", "error", err) } } func (h *Handler) answerCbWithText(ctx context.Context, callbackID, text string) { if _, err := h.Bot.AnswerCallbackQuery(ctx, &tgbot.AnswerCallbackQueryParams{CallbackQueryID: callbackID, Text: text}); err != nil { slog.Debug("answer callback with text failed", "error", err) } } func (h *Handler) backToMenu(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) { defer h.answerCb(ctx, cbID) h.clearActiveJob(userID) h.editText(ctx, chatID, msgID, "Main Menu:", &models.InlineKeyboardMarkup{ InlineKeyboard: mainKeyboard().InlineKeyboard, }) h.clearStepState(chatID) } func (h *Handler) clearActiveJob(userID int64) { h.mu.Lock() defer h.mu.Unlock() if job, ok := h.activeJobs[userID]; ok { dl.CancelDownload(job) delete(h.activeJobs, userID) } } func (h *Handler) getActiveJob(userID int64) *domain.DownloadJob { h.mu.Lock() defer h.mu.Unlock() return h.activeJobs[userID] } func (h *Handler) setActiveJob(userID int64, job *domain.DownloadJob) { h.mu.Lock() defer h.mu.Unlock() h.activeJobs[userID] = job } func formatBytes(bytes int64) string { const unit = 1024 if bytes < unit { return strconv.FormatInt(bytes, 10) + "B" } div, exp := int64(unit), 0 for n := bytes / unit; n >= unit; n /= unit { div *= unit exp++ } switch exp { case 0: return fmt.Sprintf("%.1fKB", float64(bytes)/float64(div)) case 1: return fmt.Sprintf("%.1fMB", float64(bytes)/float64(div)) case 2: return fmt.Sprintf("%.1fGB", float64(bytes)/float64(div)) default: return fmt.Sprintf("%.1fTB", float64(bytes)/float64(div)) } } func (h *Handler) generateJobID() string { return uuid.NewString()[:8] } // handleStart shows the main menu. func (h *Handler) handleStart(ctx context.Context, msg *models.Message, userID int64) { slog.Info("start", "user_id", userID, "chat_id", msg.Chat.ID) h.sendWithKB(ctx, msg.Chat.ID, "Main Menu:", mainKeyboard()) } // handleHelp shows the help text. func (h *Handler) handleHelp(ctx context.Context, msg *models.Message) { text := `Commands: /start - Show main menu /settings - Open settings /help - Show this help /history - View download history Just send me a URL to start downloading.` h.sendText(ctx, msg.Chat.ID, text) } // handleSettings opens the settings menu. func (h *Handler) handleSettings(ctx context.Context, msg *models.Message, userID int64) { prefs, err := h.Repo.GetOrCreatePreferences(userID) if err != nil { h.sendText(ctx, msg.Chat.ID, "Error loading settings") return } text, kb := h.buildSettingsKeyboard(userID, prefs) h.sendWithKB(ctx, msg.Chat.ID, text, kb) } // handleHistory opens the download history. func (h *Handler) handleHistory(ctx context.Context, msg *models.Message, userID int64) { records, err := h.Repo.GetHistory(userID, 10, 0) if err != nil { h.sendText(ctx, msg.Chat.ID, "Error loading history.") return } if len(records) == 0 { h.sendText(ctx, msg.Chat.ID, "No download history yet.") return } text := "Recent downloads:\n" for i, r := range records { text += fmt.Sprintf("%d. %s (%s)\n", i+1, r.Title, formatBytes(r.FileSize)) } h.sendText(ctx, msg.Chat.ID, text) } // historyCallback handles the history button from the main menu. func (h *Handler) historyCallback(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) { defer h.answerCb(ctx, cbID) records, err := h.Repo.GetHistory(userID, 10, 0) if err != nil { h.editText(ctx, chatID, msgID, "Error loading history.", &models.InlineKeyboardMarkup{ InlineKeyboard: [][]models.InlineKeyboardButton{ {{Text: "Back to Menu", CallbackData: "back_menu"}}, }, }) return } if len(records) == 0 { h.editText(ctx, chatID, msgID, "No download history yet.", &models.InlineKeyboardMarkup{ InlineKeyboard: [][]models.InlineKeyboardButton{ {{Text: "Back to Menu", CallbackData: "back_menu"}}, }, }) return } text := "Recent downloads:\n" for i, r := range records { text += fmt.Sprintf("%d. %s (%s)\n", i+1, r.Title, formatBytes(r.FileSize)) } h.editText(ctx, chatID, msgID, text, &models.InlineKeyboardMarkup{ InlineKeyboard: [][]models.InlineKeyboardButton{ {{Text: "Back to Menu", CallbackData: "back_menu"}}, }, }) } // applyQuota adds file size to the user's quota counters. func (h *Handler) applyQuota(userID int64, fileSize int64) { q, err := h.Repo.GetOrCreateQuotas(userID) if err != nil { return } q = h.resetQuotasIfNeeded(q) sizeMB := fileSize / (1024 * 1024) if sizeMB < 1 { sizeMB = 1 } q.DailyUsedMB += sizeMB q.WeeklyUsedMB += sizeMB q.MonthlyUsedMB += sizeMB _ = h.Repo.UpdateQuotas(userID, q) } // resetQuotasIfNeeded resets quota counters if the period has changed. func (h *Handler) resetQuotasIfNeeded(q *domain.UserQuotas) *domain.UserQuotas { today := time.Now().UTC().Format(domain.DateLayout) _, weekNum := time.Now().UTC().ISOWeek() weekStart := fmt.Sprintf("%d-W%02d", time.Now().UTC().Year(), weekNum) monthStart := time.Now().UTC().Format("2006-01") if q.DailyDate != today { q.DailyUsedMB = 0 q.DailyDate = today } if q.WeekStart != weekStart { q.WeeklyUsedMB = 0 q.WeekStart = weekStart } if q.MonthStart != monthStart { q.MonthlyUsedMB = 0 q.MonthStart = monthStart } return q }