package handler import ( "context" "fmt" "log/slog" "os" "path/filepath" "strings" tgbot "github.com/go-telegram/bot" "github.com/go-telegram/bot/models" "uptodownBot/internal/dl" "uptodownBot/internal/domain" ) // Step state stored per user during the download flow type stepState struct { URL string MediaInfo *domain.MediaInfo Step int // 0=type, 1=quality, 2=language, 3=container MediaType domain.MediaType Quality string Language string Container string FormatIDs []string PlaylistState *domain.PlaylistState URLMessageID int // original URL message ID for editing } // getUserStepState retrieves the current step state for a user. func (h *Handler) getUserStepState(chatID int64) *stepState { h.mu.Lock() defer h.mu.Unlock() return h.stepStates[chatID] } func (h *Handler) setUserStepState(chatID int64, ss *stepState) { h.mu.Lock() defer h.mu.Unlock() h.stepStates[chatID] = ss } func (h *Handler) clearStepState(chatID int64) { h.mu.Lock() defer h.mu.Unlock() delete(h.stepStates, chatID) } // handleURLInput is the entry point when a user sends a URL or types in Download mode. func (h *Handler) handleURLInput(ctx context.Context, chatID int64, text string, userID int64) { url, ok := domain.ValidateURL(text) if !ok { h.sendText(ctx, chatID, "Please send a valid URL (http/https).") return } // Check for active job if job := h.getActiveJob(userID); job != nil { h.sendText(ctx, chatID, "You already have an active download. Cancel it first.") return } // Check quota info, err := h.YtDlp.GetMediaInfo(url) if err != nil { errMsg := err.Error() if strings.Contains(errMsg, "Unsupported URL") || strings.Contains(errMsg, "is not a valid URL") { h.sendText(ctx, chatID, "This URL is not supported.") } else if strings.Contains(errMsg, "HTTP Error 403") || strings.Contains(errMsg, "cookies") { h.sendText(ctx, chatID, "This content requires authentication (cookies).") } else if strings.Contains(errMsg, "live") { h.sendText(ctx, chatID, "Cannot download live streams.") } else { h.sendText(ctx, chatID, fmt.Sprintf("Error: %s", errMsg)) } return } if info.IsPlaylist { h.handlePlaylist(ctx, chatID, userID, url, info) return } if len(info.Formats) == 0 { h.sendText(ctx, chatID, "No formats available for this URL.") return } // Check quota before starting flow q, err := h.Repo.GetOrCreateQuotas(userID) if err == nil { q = h.resetQuotasIfNeeded(q) estimatedMB := info.EstimatedSize / (1024 * 1024) if estimatedMB > 0 { if q.DailyUsedMB+estimatedMB > q.DailyLimitMB || q.WeeklyUsedMB+estimatedMB > q.WeeklyLimitMB || q.MonthlyUsedMB+estimatedMB > q.MonthlyLimitMB { h.sendText(ctx, chatID, fmt.Sprintf( "Quota exceeded. Daily: %d/%d MB, Weekly: %d/%d MB, Monthly: %d/%d MB.", q.DailyUsedMB, q.DailyLimitMB, q.WeeklyUsedMB, q.WeeklyLimitMB, q.MonthlyUsedMB, q.MonthlyLimitMB)) return } } } ss := &stepState{ URL: url, MediaInfo: info, Step: 0, } h.setUserStepState(chatID, ss) h.editText(ctx, chatID, 0, fmt.Sprintf("URL: %s\nTitle: %s\n\nSelect media type:", url, info.Title), &models.InlineKeyboardMarkup{ InlineKeyboard: mediaTypeKeyboard().InlineKeyboard, }) } // handleTypeSelection processes the media type choice. func (h *Handler) handleTypeSelection(ctx context.Context, chatID int64, msgID int, userID int64, choice string) { ss := h.getUserStepState(chatID) if ss == nil { return } switch choice { case "audio": ss.MediaType = domain.MediaTypeAudio case "video": ss.MediaType = domain.MediaTypeVideo case "both": ss.MediaType = domain.MediaTypeVideoAudio default: return } ss.Step = 1 // Filter and deduplicate formats by type filtered := dl.FilterFormatsByType(ss.MediaInfo.Formats, ss.MediaType) if len(filtered) == 0 { h.editText(ctx, chatID, msgID, fmt.Sprintf("No %s formats available for this URL.", choice), &models.InlineKeyboardMarkup{ InlineKeyboard: [][]models.InlineKeyboardButton{ {{Text: "Back", CallbackData: "back_step"}, {Text: "Cancel", CallbackData: "cancel_dl"}}, }, }) return } deduped := dl.DeduplicateResolutions(filtered) ss.FormatIDs = make([]string, len(deduped)) labels := make([]string, len(deduped)) for i, f := range deduped { ss.FormatIDs[i] = f.ID labels[i] = formatLabelForDisplay(f) } // Add "Best" option first ss.FormatIDs = append([]string{"best"}, ss.FormatIDs...) labels = append([]string{"Best"}, labels...) h.setUserStepState(chatID, ss) kb := formatKeyboard("quality_", labels) h.editText(ctx, chatID, msgID, fmt.Sprintf("Title: %s\n\nSelect quality:", ss.MediaInfo.Title), &kb) } func formatLabelForDisplay(f domain.Format) string { if f.Resolution != "" { label := f.Resolution if f.Filesize > 0 { label += fmt.Sprintf(" (%s)", formatBytes(f.Filesize)) } return label } if f.Bitrate != "" { label := f.Bitrate if f.Filesize > 0 { label += fmt.Sprintf(" (%s)", formatBytes(f.Filesize)) } return label } return f.ID } // handleQualitySelection processes the quality choice. func (h *Handler) handleQualitySelection(ctx context.Context, chatID int64, msgID int, userID int64, quality string) { ss := h.getUserStepState(chatID) if ss == nil { return } ss.Quality = quality ss.Step = 2 // Check if we need language selection if len(ss.MediaInfo.Languages) > 1 { kb := languageKeyboard(ss.MediaInfo.Languages) h.editText(ctx, chatID, msgID, "Select language:", &kb) return } // Skip to container selection ss.Step = 3 containers := dl.ContainerOptions(ss.MediaType) kb := containerKeyboard(containers) h.editText(ctx, chatID, msgID, "Select container format:", &kb) } // handleLanguageSelection processes the language choice. func (h *Handler) handleLanguageSelection(ctx context.Context, chatID int64, msgID int, userID int64, language string) { ss := h.getUserStepState(chatID) if ss == nil { return } ss.Language = language ss.Step = 3 h.setUserStepState(chatID, ss) containers := dl.ContainerOptions(ss.MediaType) kb := containerKeyboard(containers) h.editText(ctx, chatID, msgID, "Select container format:", &kb) } // handleContainerSelection processes the container format choice and starts the download. func (h *Handler) handleContainerSelection(ctx context.Context, chatID int64, msgID int, userID int64, container string) { ss := h.getUserStepState(chatID) if ss == nil { return } ss.Container = container // Build the format object format := &domain.Format{ID: ss.Quality} for _, f := range ss.MediaInfo.Formats { if f.ID == ss.Quality { format = &f break } } // Create the download job job := &domain.DownloadJob{ ID: h.generateJobID(), UserID: userID, ChatID: chatID, MessageID: msgID, URL: ss.URL, MediaType: ss.MediaType, SelectedFormat: format, Container: container, Language: ss.Language, Status: domain.StatusDownloading, Title: ss.MediaInfo.Title, FileSize: format.Filesize, } h.setActiveJob(userID, job) h.clearStepState(chatID) // Start download in background go h.runDownload(ctx, job, msgID) } // runDownload manages the download lifecycle, progress updates, and completion. func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob, msgID int) { // Show initial progress message kb := progressKeyboard(0, "", "") msg, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{ ChatID: job.ChatID, Text: fmt.Sprintf("Downloading: %s", job.Title), ReplyMarkup: &kb, }) if err != nil { slog.Error("send progress msg failed", "error", err) return } job.MessageID = msg.ID updates, err := h.YtDlp.StartDownload(job, h.downloadDir) if err != nil { h.editText(ctx, job.ChatID, job.MessageID, fmt.Sprintf("Download failed: %s", err.Error()), nil) job.Status = domain.StatusFailed h.mu.Lock() delete(h.activeJobs, job.UserID) h.mu.Unlock() return } lastProgress := 0.0 for upd := range updates { h.mu.Lock() currentJob := h.activeJobs[job.UserID] h.mu.Unlock() if currentJob == nil { // Job was cancelled return } if upd.Status == domain.StatusCancelled { h.editText(ctx, job.ChatID, job.MessageID, "Download cancelled.", nil) h.mu.Lock() delete(h.activeJobs, job.UserID) h.mu.Unlock() return } if upd.Status == domain.StatusFailed { h.editText(ctx, job.ChatID, job.MessageID, "Download failed.", nil) h.mu.Lock() delete(h.activeJobs, job.UserID) h.mu.Unlock() return } // Update progress every ~5% if upd.Progress-lastProgress >= domain.ProgressThreshold || upd.Status == domain.StatusCompleted { lastProgress = upd.Progress // Apply quota at >50% if upd.Progress > domain.QuotaProgressThreshold && !job.QuotaApplied && job.FileSize > 0 { job.QuotaApplied = true h.applyQuota(job.UserID, job.FileSize) } kb := progressKeyboard(upd.Progress, upd.Speed, upd.ETA) text := fmt.Sprintf("Downloading: %s", job.Title) h.editText(ctx, job.ChatID, job.MessageID, text, &kb) } } // Download completed h.mu.Lock() delete(h.activeJobs, job.UserID) h.mu.Unlock() if job.Status == domain.StatusCompleted { // Find the downloaded file files, err := filepath.Glob(filepath.Join(h.downloadDir, "*")) if err == nil { for _, f := range files { fileInfo, err := os.Stat(f) if err != nil { continue } if fileInfo.Size() > h.maxFileSize { h.editText(ctx, job.ChatID, job.MessageID, fmt.Sprintf("File too large (%s). Max: %s", formatBytes(fileInfo.Size()), formatBytes(h.maxFileSize)), nil) os.Remove(f) return } // Check file size quota after download if !job.QuotaApplied && job.FileSize > 0 { job.QuotaApplied = true h.applyQuota(job.UserID, job.FileSize) } formatID := job.SelectedFormat.ID h.sendDocument(ctx, job.ChatID, f, job.Title) os.Remove(f) // Save to history _ = h.Repo.AddHistory(job.UserID, &domain.DownloadRecord{ URL: job.URL, Title: job.Title, MediaType: job.MediaType.String(), FormatID: formatID, Container: job.Container, FileSize: job.FileSize, Status: "completed", }) h.editText(ctx, job.ChatID, job.MessageID, "Download completed!", nil) return } } h.editText(ctx, job.ChatID, job.MessageID, "Download completed but file not found.", nil) } } func (h *Handler) sendDocument(ctx context.Context, chatID int64, filePath, title string) { f, err := os.Open(filePath) if err != nil { slog.Error("open file for send", "path", filePath, "error", err) return } defer f.Close() doc := &models.InputFileUpload{ Filename: title + filepath.Ext(filePath), Data: f, } if _, err := h.Bot.SendDocument(ctx, &tgbot.SendDocumentParams{ ChatID: chatID, Document: doc, }); err != nil { slog.Error("send document failed", "error", err) } } // handleCancelDownload cancels the current download for a user. func (h *Handler) handleCancelDownload(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) { defer h.answerCb(ctx, cbID) job := h.getActiveJob(userID) if job != nil { dl.CancelDownload(job) h.mu.Lock() delete(h.activeJobs, userID) h.mu.Unlock() h.editText(ctx, chatID, msgID, "Download cancelled.", nil) return } h.clearStepState(chatID) h.editText(ctx, chatID, msgID, "Cancelled.", nil) } // handleBackStep goes back one step in the download flow. func (h *Handler) handleBackStep(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) { defer h.answerCb(ctx, cbID) ss := h.getUserStepState(chatID) if ss == nil { return } ss.Step-- switch ss.Step { case -1: // Go back to URL input h.clearStepState(chatID) h.handleDownloadPrompt(ctx, chatID, msgID, cbID, userID) case 0: h.editText(ctx, chatID, msgID, "Select media type:", &models.InlineKeyboardMarkup{ InlineKeyboard: mediaTypeKeyboard().InlineKeyboard, }) case 1: // Re-show quality selector filtered := dl.FilterFormatsByType(ss.MediaInfo.Formats, ss.MediaType) deduped := dl.DeduplicateResolutions(filtered) labels := make([]string, len(deduped)) for i, f := range deduped { labels[i] = formatLabelForDisplay(f) } labels = append([]string{"Best"}, labels...) kb := formatKeyboard("quality_", labels) h.editText(ctx, chatID, msgID, "Select quality:", &kb) case 2: if len(ss.MediaInfo.Languages) > 1 { kb := languageKeyboard(ss.MediaInfo.Languages) h.editText(ctx, chatID, msgID, "Select language:", &kb) } else { ss.Step = 1 h.handleBackStep(ctx, chatID, msgID, cbID, userID) } } } // handleDownloadPrompt shows the download prompt (asks for URL). func (h *Handler) handleDownloadPrompt(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) { defer h.answerCb(ctx, cbID) h.promptForInput(ctx, chatID, msgID, userID, PendingURL, "Send me the URL to download:", nil) }