refactor: redesign format selection with unified list and independent audio/video picking

- Replace type/quality selection flow with unified format list (video+audio combined) - Add independent audio/video format picking: video-only formats prompt for audio track, audio-only formats offer optional video addition (skip when content is audio-only) - Fix formatKeyboard to use format IDs as callback data instead of display labels - Remove dead mediaTypeKeyboard function - Fix double-close panic in readProgress (caller is sole owner of updates channel) - Add FormatString field to DownloadJob for composite format specs (e.g. '137+140') - Change quota calculation from sum of all format sizes to max single format size - Fix 'Delete all data' button to use danger style with proper text - Clean up routePrefixedCallback: remove type_/quality_, add format_/secondary_ routes - Fix back navigation to rebuild format IDs list
2026-06-25 15:37:47 +03:30
parent 7de79b50e1
commit e65ade6c3c
7 changed files with 286 additions and 198 deletions
--- a/internal/handler/download.go
+++ b/internal/handler/download.go
@@ -17,21 +17,20 @@ import (
 	"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
+	Step          int
+	MainFormatID  string
+	MainFormat    *domain.Format
+	SecondaryID   string
+	SecondaryFmt  *domain.Format
 	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()
@@ -58,23 +57,22 @@ func (h *Handler) handleURLInput(ctx context.Context, chatID int64, text string,
 		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") {
+		switch {
+		case 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") {
+		case 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") {
+		case strings.Contains(errMsg, "live"):
 			h.sendText(ctx, chatID, "Cannot download live streams.")
-		} else {
+		default:
 			h.sendText(ctx, chatID, fmt.Sprintf("Error: %s", errMsg))
 		}
 		return
@@ -90,7 +88,6 @@ func (h *Handler) handleURLInput(ctx context.Context, chatID int64, text string,
 		return
 	}

-	// Check quota before starting flow
 	q, err := h.Repo.GetOrCreateQuotas(userID)
 	if err == nil {
 		q = h.resetQuotasIfNeeded(q)
@@ -112,124 +109,190 @@ func (h *Handler) handleURLInput(ctx context.Context, chatID int64, text string,
 		MediaInfo: info,
 		Step:      0,
 	}
-	h.setUserStepState(chatID, ss)

-	// Show thumbnail and metadata preview
 	h.sendMediaPreview(ctx, chatID, info)

-	if info.IsAudioOnlyContent() {
-		// Audio-only sites: skip type selection, go straight to quality
-		ss.MediaType = domain.MediaTypeAudio
-		ss.Step = 1
-		filtered := dl.FilterFormatsByType(info.Formats, domain.MediaTypeAudio)
-		if len(filtered) == 0 {
-			h.sendText(ctx, chatID, "No audio formats available for this URL.")
-			return
+	// Build unified format list: video formats first, then audio
+	deduped := dl.DeduplicateResolutions(info.Formats)
+	var videoFmts, audioFmts []domain.Format
+	for _, f := range deduped {
+		if f.HasVideo {
+			videoFmts = append(videoFmts, f)
+		} else {
+			audioFmts = append(audioFmts, f)
 		}
-		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)
-		}
-		ss.FormatIDs = append([]string{"best"}, ss.FormatIDs...)
-		labels = append([]string{"Best"}, labels...)
-		h.setUserStepState(chatID, ss)
-		kb := formatKeyboard("quality_", labels)
-		h.sendWithKB(ctx, chatID, "Select audio quality:", kb)
-		return
+	}
+	ss.FormatIDs = make([]string, 0, len(deduped)+1)
+	labels := make([]string, 0, len(deduped)+1)
+	ss.FormatIDs = append(ss.FormatIDs, "best")
+	labels = append(labels, "Best")
+	for _, f := range videoFmts {
+		ss.FormatIDs = append(ss.FormatIDs, f.ID)
+		labels = append(labels, formatLabelForDisplay(f))
+	}
+	for _, f := range audioFmts {
+		ss.FormatIDs = append(ss.FormatIDs, f.ID)
+		labels = append(labels, formatLabelForDisplay(f))
 	}

-	kb := mediaTypeKeyboard()
-	h.sendWithKB(ctx, chatID, fmt.Sprintf("Title: %s\n\nSelect media type:", info.Title), kb)
+	h.setUserStepState(chatID, ss)
+	kb := formatKeyboard("format_", labels, ss.FormatIDs)
+	h.sendWithKB(ctx, chatID, "Select format:", kb)
 }

-// handleTypeSelection processes the media type choice.
-func (h *Handler) handleTypeSelection(ctx context.Context, chatID int64, msgID int, userID int64, choice string) {
+// 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 {
 		return
 	}
-	switch choice {
-	case "audio":
-		ss.MediaType = domain.MediaTypeAudio
-	case "video":
-		ss.MediaType = domain.MediaTypeVideo
-	case "both":
-		ss.MediaType = domain.MediaTypeVideoAudio
-	default:
+	ss.MainFormatID = formatID
+
+	if formatID == "best" {
+		ss.MainFormat = &domain.Format{ID: "best", HasVideo: true, HasAudio: true}
+		h.advanceToLanguage(ctx, chatID, msgID, userID, ss)
 		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{
+	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)
-	ss.FormatIDs = make([]string, len(deduped))
+	ids := make([]string, len(deduped))
 	labels := make([]string, len(deduped))
 	for i, f := range deduped {
-		ss.FormatIDs[i] = f.ID
+		ids[i] = f.ID
 		labels[i] = formatLabelForDisplay(f)
 	}
-	// Add "Best" option first
-	ss.FormatIDs = append([]string{"best"}, ss.FormatIDs...)
-	labels = append([]string{"Best"}, labels...)
-
+	ss.FormatIDs = ids
 	h.setUserStepState(chatID, ss)
-	kb := formatKeyboard("quality_", labels)
-	h.editText(ctx, chatID, msgID, fmt.Sprintf("Title: %s\n\nSelect quality:", ss.MediaInfo.Title), &kb)
+
+	prompt := fmt.Sprintf("Select %s format to combine:", kind)
+	kb := formatKeyboard("secondary_", labels, ids)
+	h.editText(ctx, chatID, msgID, prompt, &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) {
+// 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 {
 		return
 	}
-	ss.Quality = quality
-	ss.Step = 2

-	// Check if we need language selection
+	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 skips to container if N/A.
+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.advanceToContainer(ctx, chatID, msgID, userID, ss)
+}

-	// Skip to container selection
+// advanceToContainer shows the container picker.
+func (h *Handler) advanceToContainer(ctx context.Context, chatID int64, msgID int, userID int64, ss *stepState) {
 	ss.Step = 3
-	containers := dl.ContainerOptions(ss.MediaType)
+	h.setUserStepState(chatID, ss)
+
+	containers := dl.ContainerOptions(ss.determineMediaType())
 	kb := containerKeyboard(containers)
 	h.editText(ctx, chatID, msgID, "Select container format:", &kb)
 }

+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)
@@ -237,15 +300,10 @@ func (h *Handler) handleLanguageSelection(ctx context.Context, chatID int64, msg
 		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)
+	h.advanceToContainer(ctx, chatID, msgID, userID, ss)
 }

-// handleContainerSelection processes the container format choice and starts the download.
+// handleContainerSelection processes the container 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 {
@@ -253,24 +311,23 @@ func (h *Handler) handleContainerSelection(ctx context.Context, chatID int64, ms
 	}
 	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
-		}
+	mediaType := ss.determineMediaType()
+	formatString := ss.buildFormatString()
+
+	format := &domain.Format{ID: ss.MainFormatID}
+	if ss.MainFormat != nil {
+		format = ss.MainFormat
 	}

-	// Create the download job
 	job := &domain.DownloadJob{
 		ID:             h.generateJobID(),
 		UserID:         userID,
 		ChatID:         chatID,
 		MessageID:      msgID,
 		URL:            ss.URL,
-		MediaType:      ss.MediaType,
+		MediaType:      mediaType,
 		SelectedFormat: format,
+		FormatString:   formatString,
 		Container:      container,
 		Language:       ss.Language,
 		Status:         domain.StatusDownloading,
@@ -281,13 +338,11 @@ func (h *Handler) handleContainerSelection(ctx context.Context, chatID int64, ms
 	h.setActiveJob(userID, job)
 	h.clearStepState(chatID)

-	// Start download in background
 	go h.runDownload(context.Background(), job)
 }

-// runDownload manages the download lifecycle, progress updates, and completion.
+// runDownload manages the download lifecycle, progress, and completion.
 func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) {
-	// Show initial progress message
 	kb := progressKeyboard(0, "", "")
 	msg, err := h.Bot.SendMessage(ctx, &tgbot.SendMessageParams{
 		ChatID:      job.ChatID,
@@ -316,7 +371,6 @@ func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) {
 		currentJob := h.activeJobs[job.UserID]
 		h.mu.Unlock()
 		if currentJob == nil {
-			// Job was cancelled
 			return
 		}

@@ -336,11 +390,9 @@ func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) {
 			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)
@@ -352,17 +404,15 @@ func (h *Handler) runDownload(ctx context.Context, job *domain.DownloadJob) {
 		}
 	}

-	// Download completed
 	h.mu.Lock()
 	delete(h.activeJobs, job.UserID)
 	h.mu.Unlock()

 	if job.Status == domain.StatusCompleted {
-		// Find the downloaded file using job ID prefix
 		pattern := filepath.Join(h.downloadDir, job.ID+"_*")
 		files, err := filepath.Glob(pattern)
 		if err == nil && len(files) > 0 {
-			f := files[0] // take the first match
+			f := files[0]
 			fileInfo, err := os.Stat(f)
 			if err == nil {
 				if fileInfo.Size() > h.maxFileSize {
@@ -448,30 +498,41 @@ func (h *Handler) handleBackStep(ctx context.Context, chatID int64, msgID int, c
 	ss.Step--
 	switch ss.Step {
 	case -1:
-		// Go back to URL input
 		h.clearStepState(chatID)
 		h.handleDownloadPrompt(ctx, chatID, msgID, cbID, userID)
 	case 0:
-		if ss.MediaInfo != nil && ss.MediaInfo.IsAudioOnlyContent() {
-			// Audio-only: go back to URL input directly
-			h.clearStepState(chatID)
-			h.handleDownloadPrompt(ctx, chatID, msgID, cbID, userID)
-			return
+		// Re-show format list
+		deduped := dl.DeduplicateResolutions(ss.MediaInfo.Formats)
+		var videoFmts, audioFmts []domain.Format
+		for _, f := range deduped {
+			if f.HasVideo {
+				videoFmts = append(videoFmts, f)
+			} else {
+				audioFmts = append(audioFmts, f)
+			}
 		}
-		h.editText(ctx, chatID, msgID, "Select media type:", &models.InlineKeyboardMarkup{
-			InlineKeyboard: mediaTypeKeyboard().InlineKeyboard,
-		})
+		ss.FormatIDs = []string{"best"}
+		labels := []string{"Best"}
+		for _, f := range videoFmts {
+			ss.FormatIDs = append(ss.FormatIDs, f.ID)
+			labels = append(labels, formatLabelForDisplay(f))
+		}
+		for _, f := range audioFmts {
+			ss.FormatIDs = append(ss.FormatIDs, f.ID)
+			labels = append(labels, formatLabelForDisplay(f))
+		}
+		h.setUserStepState(chatID, ss)
+		kb := formatKeyboard("format_", labels, ss.FormatIDs)
+		h.editText(ctx, chatID, msgID, "Select format:", &kb)
 	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)
+		// 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)
 		}
-		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)
@@ -495,6 +556,24 @@ func (h *Handler) handleDownloadPrompt(ctx context.Context, chatID int64, msgID
 		"Send me the URL to download:", 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
+	}
+	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