package handler import ( "context" "fmt" "io" "log/slog" "math" "net/http" "os" "path/filepath" "strings" "time" tgbot "github.com/go-telegram/bot" "github.com/go-telegram/bot/models" "uptodownBot/internal/dl" "uptodownBot/internal/domain" ) type stepState struct { URL string MediaInfo *domain.MediaInfo Step int MainFormatID string MainFormat *domain.Format SecondaryID string SecondaryFmt *domain.Format Language string FormatIDs []string PlaylistState *domain.PlaylistState SearchState *domain.SearchState CreatedAt time.Time } const stepStateTTL = 30 * time.Minute func (h *Handler) getUserStepState(chatID int64) *stepState { h.stepStatesMu.Lock() defer h.stepStatesMu.Unlock() ss := h.stepStates[chatID] if ss != nil && time.Since(ss.CreatedAt) > stepStateTTL { delete(h.stepStates, chatID) return nil } return ss } func (h *Handler) setUserStepState(chatID int64, ss *stepState) { h.stepStatesMu.Lock() defer h.stepStatesMu.Unlock() h.stepStates[chatID] = ss } func (h *Handler) clearStepState(chatID int64) { h.stepStatesMu.Lock() defer h.stepStatesMu.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 } if job := h.getActiveJob(userID); job != nil { h.sendText(ctx, chatID, "You already have an active download. Cancel it first.") return } info, err := h.YtDlp.GetMediaInfo(url) if err != nil { errMsg := err.Error() // DRM content (Spotify, etc.): fall back to YouTube search if strings.Contains(errMsg, "DRM") { youtubeURL, ytTitle, searchErr := h.YtDlp.SearchYoutube(text) if searchErr != nil { kb := models.InlineKeyboardMarkup{ InlineKeyboard: [][]models.InlineKeyboardButton{ {{Text: "Search on YouTube", CallbackData: "drm_search"}}, {{Text: "Cancel", CallbackData: "pending_cancel"}}, }, } h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{ ChatID: chatID, Text: "This content is DRM protected and no YouTube alternative was found automatically. Try searching by name:", ReplyMarkup: kb, }) return } h.sendText(ctx, chatID, fmt.Sprintf("This content is DRM protected. Found on YouTube: %s", ytTitle)) info, err = h.YtDlp.GetMediaInfo(youtubeURL) if err != nil { h.sendText(ctx, chatID, fmt.Sprintf("YouTube fallback failed: %s", err.Error())) return } url = youtubeURL } else { switch { case strings.Contains(errMsg, "Unsupported URL") || strings.Contains(errMsg, "is not a valid URL"): h.sendText(ctx, chatID, "This URL is not supported.") case strings.Contains(errMsg, "HTTP Error 403") || strings.Contains(errMsg, "cookies"): h.sendText(ctx, chatID, "This content requires authentication (cookies).") case strings.Contains(errMsg, "live"): h.sendText(ctx, chatID, "Cannot download live streams.") default: 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 } // Reject if user has already exhausted quota (limit of 0 = unlimited). q, err := h.Repo.GetOrCreateQuotas(userID) if err == nil { q = h.resetQuotasIfNeeded(q) if (q.DailyLimitMB > 0 && q.DailyUsedMB >= q.DailyLimitMB) || (q.WeeklyLimitMB > 0 && q.WeeklyUsedMB >= q.WeeklyLimitMB) || (q.MonthlyLimitMB > 0 && q.MonthlyUsedMB >= 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, CreatedAt: time.Now(), } h.sendMediaPreview(ctx, chatID, info) ss.FormatIDs, _ = buildFormatList(info.Formats) h.setUserStepState(chatID, ss) labels := make([]string, len(ss.FormatIDs)) for i, id := range ss.FormatIDs { if id == "best" { labels[i] = "Best" } else { f := findFormatByID(info.Formats, id) if f != nil { labels[i] = formatLabelForDisplay(*f) } else { labels[i] = id } } } kb := formatKeyboard("format_", labels, ss.FormatIDs) h.sendWithKB(ctx, chatID, "Select format:", kb) } // handleFormatSelection processes a format pick from the unified list. func (h *Handler) handleFormatSelection(ctx context.Context, chatID int64, msgID int, userID int64, formatID string) { ss := h.getUserStepState(chatID) if ss == nil { slog.Warn("nil stepState in handleFormatSelection", "chat_id", chatID, "user_id", userID) return } ss.MainFormatID = formatID if formatID == "best" { if ss.MediaInfo.IsAudioOnlyContent() { ss.MainFormat = &domain.Format{ID: "best", HasAudio: true, IsAudioOnly: true} } else { ss.MainFormat = &domain.Format{ID: "best", HasVideo: true, HasAudio: true} } h.advanceToLanguage(ctx, chatID, msgID, userID, ss) return } for i := range ss.MediaInfo.Formats { f := &ss.MediaInfo.Formats[i] if f.ID == formatID { ss.MainFormat = f break } } if ss.MainFormat == nil { return } // Video-only: always offer audio track if ss.MainFormat.HasVideo && !ss.MainFormat.HasAudio { h.showSecondaryFormats(ctx, chatID, msgID, userID, ss, "audio") return } // Audio-only: ask if they want to add video if !ss.MainFormat.HasVideo && ss.MainFormat.HasAudio { if ss.MediaInfo.IsAudioOnlyContent() { h.advanceToLanguage(ctx, chatID, msgID, userID, ss) return } text := "This format is audio-only. Add video?" kb := models.InlineKeyboardMarkup{ InlineKeyboard: [][]models.InlineKeyboardButton{ {{Text: "Add Video", CallbackData: "secondary_video"}}, {{Text: "Audio Only", CallbackData: "secondary_skip"}}, {{Text: "Back", CallbackData: "back_step"}, {Text: "Cancel", CallbackData: "cancel_dl"}}, }, } h.editText(ctx, chatID, msgID, text, &kb) return } // Has both: proceed h.advanceToLanguage(ctx, chatID, msgID, userID, ss) } // showSecondaryFormats shows formats of the missing type to combine. func (h *Handler) showSecondaryFormats(ctx context.Context, chatID int64, msgID int, userID int64, ss *stepState, kind string) { var mt domain.MediaType if kind == "audio" { mt = domain.MediaTypeAudio } else { mt = domain.MediaTypeVideo } filtered := dl.FilterFormatsByType(ss.MediaInfo.Formats, mt) if len(filtered) == 0 { ss.SecondaryFmt = nil h.advanceToLanguage(ctx, chatID, msgID, userID, ss) return } deduped := dl.DeduplicateResolutions(filtered) ids := make([]string, len(deduped)) labels := make([]string, len(deduped)) for i, f := range deduped { ids[i] = f.ID labels[i] = formatLabelForDisplay(f) } ss.FormatIDs = ids h.setUserStepState(chatID, ss) prompt := fmt.Sprintf("Select %s format to combine:", kind) kb := formatKeyboard("secondary_", labels, ids) h.editText(ctx, chatID, msgID, prompt, &kb) } // handleSecondarySelection processes the secondary format selection. func (h *Handler) handleSecondarySelection(ctx context.Context, chatID int64, msgID int, userID int64, formatID string) { ss := h.getUserStepState(chatID) if ss == nil { slog.Warn("nil stepState in handleSecondarySelection", "chat_id", chatID, "user_id", userID) return } if formatID == "skip" || formatID == "video" || formatID == "audio" { // Handle the skip / add-video decisions from the audio-only prompt if formatID == "video" { h.showSecondaryFormats(ctx, chatID, msgID, userID, ss, "video") return } // skip — just proceed h.advanceToLanguage(ctx, chatID, msgID, userID, ss) return } ss.SecondaryID = formatID for i := range ss.MediaInfo.Formats { f := &ss.MediaInfo.Formats[i] if f.ID == formatID { ss.SecondaryFmt = f break } } h.advanceToLanguage(ctx, chatID, msgID, userID, ss) } // advanceToLanguage moves to language selection, or starts the download if only one language. func (h *Handler) advanceToLanguage(ctx context.Context, chatID int64, msgID int, userID int64, ss *stepState) { ss.Step = 2 h.setUserStepState(chatID, ss) if len(ss.MediaInfo.Languages) > 1 { kb := languageKeyboard(ss.MediaInfo.Languages) h.editText(ctx, chatID, msgID, "Select language:", &kb) return } h.startDownload(ctx, chatID, msgID, userID, ss) } func (ss *stepState) determineMediaType() domain.MediaType { if ss.MainFormat == nil { return domain.MediaTypeVideoAudio } if ss.MainFormat.IsAudioOnly && ss.SecondaryFmt == nil { return domain.MediaTypeAudio } return domain.MediaTypeVideoAudio } func (ss *stepState) buildFormatString() string { if ss.MainFormatID == "best" { return "" } if ss.SecondaryFmt != nil { return ss.MainFormatID + "+" + ss.SecondaryID } return ss.MainFormatID } // 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 { slog.Warn("nil stepState in handleLanguageSelection", "chat_id", chatID, "user_id", userID) return } ss.Language = language h.startDownload(ctx, chatID, msgID, userID, ss) } // startDownload creates the download job and begins the download process. func (h *Handler) startDownload(ctx context.Context, chatID int64, msgID int, userID int64, ss *stepState) { mediaType := ss.determineMediaType() formatString := ss.buildFormatString() format := &domain.Format{ID: ss.MainFormatID} if ss.MainFormat != nil { format = ss.MainFormat } fileSize := format.Filesize if fileSize > 0 && fileSize > h.maxFileSize { h.editText(ctx, chatID, msgID, fmt.Sprintf("File too large (%s). Max: %s", formatBytes(fileSize), formatBytes(h.maxFileSize)), nil) return } dlCtx, dlCancel := context.WithCancel(context.Background()) job := &domain.DownloadJob{ ID: h.generateJobID(), UserID: userID, ChatID: chatID, MessageID: msgID, URL: ss.URL, MediaType: mediaType, SelectedFormat: format, FormatString: formatString, Language: ss.Language, Status: domain.StatusDownloading, Title: ss.MediaInfo.Title, FileSize: fileSize, MaxFileSize: h.maxFileSize, Done: make(chan struct{}), } h.editText(ctx, chatID, msgID, "Processing...", nil) h.setActiveJob(userID, job) h.setCancelFunc(userID, dlCancel) h.clearStepState(chatID) go h.runDownload(dlCtx, job) } // runDownload manages the download lifecycle, progress, and completion. func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) { defer func() { if job.Done != nil { close(job.Done) } }() select { case h.downloadSem <- struct{}{}: case <-ctx.Done(): return } defer func() { <-h.downloadSem }() kb := progressKeyboard(0, "", "") progressHeader := "Downloading: " + job.Title if job.ProgressPrefix != "" { progressHeader = job.ProgressPrefix + "\n" + progressHeader } msg, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{ ChatID: job.ChatID, Text: progressHeader, 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 { errMsg := err.Error() if job.StderrLog != "" { if lines := strings.SplitN(job.StderrLog, "\n", 2); len(lines) > 0 && lines[0] != "" { errMsg = strings.TrimSpace(lines[0]) } } h.editText(ctx, job.ChatID, job.MessageID, fmt.Sprintf("Download failed: %s", errMsg), nil) job.Status = domain.StatusFailed h.clearActiveJob(job.UserID) return } h.trackDownloadProgress(ctx, job, updates) if job.Status == domain.StatusCompleted { h.handleDownloadCompletion(ctx, job) } h.clearActiveJob(job.UserID) } // trackDownloadProgress reads progress updates from yt-dlp and updates the Telegram message. func (h *Handler) trackDownloadProgress(ctx context.Context, job *domain.DownloadJob, updates <-chan *domain.DownloadJob) { lastProgress := 0.0 for { select { case upd, ok := <-updates: if !ok { return } slog.Info("progress update", "pct", upd.Progress, "speed", upd.Speed, "status", upd.Status) h.activeJobsMu.Lock() currentJob := h.activeJobs[job.UserID] h.activeJobsMu.Unlock() if currentJob == nil { return } if upd.Status == domain.StatusCancelled { h.editText(ctx, job.ChatID, job.MessageID, "Download cancelled.", nil) return } if upd.Status == domain.StatusFailed { h.editText(ctx, job.ChatID, job.MessageID, "Download failed.", nil) return } if upd.Status == domain.StatusCompleted { return } if lastProgress == 0 { lastProgress = upd.Progress } else { diff := upd.Progress - lastProgress // Skip small changes (< threshold) in either direction to avoid spamming edits. // HLS streams often emit a fake 100% initial estimate then real values. if diff >= 0 && diff < domain.ProgressThreshold { continue } if diff < 0 && -diff < domain.ProgressThreshold { continue } lastProgress = upd.Progress } 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 := "Downloading: " + job.Title if job.ProgressPrefix != "" { text = job.ProgressPrefix + "\n" + text } h.editText(ctx, job.ChatID, job.MessageID, text, &kb) case <-ctx.Done(): dl.CancelDownload(job) for range updates { } return } } } // handleDownloadCompletion sends the downloaded file and records history. func (h *Handler) handleDownloadCompletion(ctx context.Context, job *domain.DownloadJob) { pattern := filepath.Join(h.downloadDir, job.ID+"_*") files, err := filepath.Glob(pattern) if err != nil || len(files) == 0 { h.editText(ctx, job.ChatID, job.MessageID, "Download completed but file not found.", nil) return } f := files[0] fileInfo, err := os.Stat(f) if err != nil { h.editText(ctx, job.ChatID, job.MessageID, "Download completed but file not found.", nil) return } 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 } if !job.QuotaApplied && job.FileSize > 0 { job.QuotaApplied = true h.applyQuota(job.UserID, job.FileSize) } h.sendDocument(ctx, job.ChatID, f, job.Title) os.Remove(f) if err := h.Repo.AddHistory(job.UserID, &domain.DownloadRecord{ URL: job.URL, Title: job.Title, MediaType: job.MediaType.String(), FormatID: job.SelectedFormat.ID, FileSize: job.FileSize, Status: "completed", }); err != nil { slog.Error("add history failed", "user_id", job.UserID, "error", err) } h.editText(ctx, job.ChatID, job.MessageID, "Download completed!", 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) h.cancelFuncsMu.Lock() if cancel, ok := h.cancelFuncs[userID]; ok { cancel() delete(h.cancelFuncs, userID) } h.cancelFuncsMu.Unlock() job := h.getActiveJob(userID) if job != nil { dl.CancelDownload(job) h.activeJobsMu.Lock() delete(h.activeJobs, userID) h.activeJobsMu.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 { slog.Warn("nil stepState in handleBackStep", "chat_id", chatID, "user_id", userID) h.handleDownloadPrompt(ctx, chatID, msgID, cbID, userID) return } ss.Step-- switch ss.Step { case -1: h.clearStepState(chatID) h.handleDownloadPrompt(ctx, chatID, msgID, cbID, userID) case 0: ss.FormatIDs, _ = buildFormatList(ss.MediaInfo.Formats) labels := make([]string, len(ss.FormatIDs)) for i, id := range ss.FormatIDs { if id == "best" { labels[i] = "Best" } else { f := findFormatByID(ss.MediaInfo.Formats, id) if f != nil { labels[i] = formatLabelForDisplay(*f) } else { labels[i] = id } } } h.setUserStepState(chatID, ss) kb := formatKeyboard("format_", labels, ss.FormatIDs) h.editText(ctx, chatID, msgID, "Select format:", &kb) case 1: // Re-show secondary format list or main format with secondary info cleared ss.SecondaryFmt = nil ss.SecondaryID = "" if ss.MainFormat != nil && ss.MainFormat.HasVideo && !ss.MainFormat.HasAudio { h.showSecondaryFormats(ctx, chatID, msgID, userID, ss, "audio") } else { h.advanceToLanguage(ctx, chatID, msgID, userID, ss) } 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) } } } // handleDownloadCommand handles the /download command, prompting for a URL. func (h *Handler) handleDownloadCommand(ctx context.Context, msg *models.Message, userID int64) { h.promptForInput(ctx, msg.Chat.ID, 0, userID, PendingURL, "Send me the URL to download:", nil) } // handleDownloadPrompt shows the download prompt (asks for URL). func (h *Handler) handleDownloadPrompt(ctx context.Context, chatID int64, msgID int, cbID string, userID int64) { h.promptForInput(ctx, chatID, msgID, userID, PendingURL, "Send me the URL to download:", nil) } // buildFormatList returns a deduplicated format ID list with "best" first, then video, then audio. func buildFormatList(formats []domain.Format) ([]string, error) { deduped := dl.DeduplicateResolutions(formats) ids := make([]string, 0, len(deduped)+1) ids = append(ids, "best") for _, f := range deduped { if f.HasVideo { ids = append(ids, f.ID) } } for _, f := range deduped { if !f.HasVideo { ids = append(ids, f.ID) } } return ids, nil } // findFormatByID returns the format with the given ID from the list, or nil. func findFormatByID(formats []domain.Format, id string) *domain.Format { for i := range formats { if formats[i].ID == id { return &formats[i] } } return nil } 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 } // Audio-only with no bitrate: show format ID + extension if f.IsAudioOnly { label := f.ID if f.Extension != "" { label += " (" + f.Extension + ")" } if f.Filesize > 0 { label += fmt.Sprintf(" (%s)", formatBytes(f.Filesize)) } return label } return f.ID } // sendMediaPreview sends a thumbnail and metadata info before the selection flow. func (h *Handler) sendMediaPreview(ctx context.Context, chatID int64, info *domain.MediaInfo) { caption := info.Title if info.Uploader != "" || info.Duration > 0 { caption += "\n" if info.Uploader != "" { caption += fmt.Sprintf("Channel: %s", info.Uploader) } if info.Duration > 0 { if info.Uploader != "" { caption += " | " } caption += fmt.Sprintf("Duration: %s", formatDuration(info.Duration)) } } if info.Thumbnail != "" { thumbPath, err := downloadFile(info.Thumbnail, h.downloadDir) if err != nil { slog.Warn("download thumbnail failed, falling back to text", "url", info.Thumbnail, "error", err) h.sendText(ctx, chatID, caption) return } defer os.Remove(thumbPath) f, err := os.Open(thumbPath) if err != nil { slog.Warn("open thumbnail failed, falling back to text", "error", err) h.sendText(ctx, chatID, caption) return } defer f.Close() photo := &models.InputFileUpload{Filename: "thumb.jpg", Data: f} _, err = h.Bot.SendPhoto(ctx, &tgbot.SendPhotoParams{ ChatID: chatID, Photo: photo, Caption: caption, }) if err == nil { return } slog.Warn("send thumbnail photo failed, falling back to text", "error", err) } h.sendText(ctx, chatID, caption) } // downloadFile downloads a URL to a temp file in the given directory and returns the path. func downloadFile(url, dir string) (string, error) { tr := &http.Transport{ Proxy: http.ProxyFromEnvironment, } client := &http.Client{ Timeout: 30 * time.Second, Transport: tr, } resp, err := client.Get(url) if err != nil { return "", fmt.Errorf("http get: %w", err) } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { return "", fmt.Errorf("unexpected status: %s", resp.Status) } if err := os.MkdirAll(dir, 0755); err != nil { return "", fmt.Errorf("create download dir: %w", err) } f, err := os.CreateTemp(dir, "thumb_*") if err != nil { return "", fmt.Errorf("create temp file: %w", err) } _, err = io.Copy(f, resp.Body) if err != nil { f.Close() os.Remove(f.Name()) return "", fmt.Errorf("write temp file: %w", err) } if err := f.Close(); err != nil { os.Remove(f.Name()) return "", fmt.Errorf("close temp file: %w", err) } return f.Name(), nil } func formatDuration(seconds float64) string { if seconds <= 0 { return "0:00" } d := time.Duration(seconds) * time.Second h := int(d.Hours()) m := int(math.Mod(d.Minutes(), 60)) s := int(math.Mod(d.Seconds(), 60)) if h > 0 { return fmt.Sprintf("%d:%02d:%02d", h, m, s) } return fmt.Sprintf("%d:%02d", m, s) }