diff --git a/app/main.go b/app/main.go index 58bc138..a40de2e 100644 --- a/app/main.go +++ b/app/main.go @@ -5,6 +5,7 @@ import ( "log" "os" "os/signal" + "sort" "sync" "syscall" "time" @@ -14,7 +15,10 @@ import ( // InteractionState holds the state for a specific interaction type InteractionState struct { - URL string + URL string + FormatOptions *FormatOptions + VideoFormatID string + AudioFormatID string } // DownloadResult represents the result of an async download operation @@ -27,7 +31,7 @@ type DownloadResult struct { } // startAsyncDownload initiates a download in a goroutine and handles progress updates -func startAsyncDownload(s *discordgo.Session, i *discordgo.InteractionCreate, url, audioFormat, outputDir string) { +func startAsyncDownload(s *discordgo.Session, i *discordgo.InteractionCreate, url, videoFormatID, audioFormatID, outputDir string) { progressChan := make(chan ProgressUpdate, 1) resultChan := make(chan DownloadResult, 1) @@ -41,21 +45,26 @@ func startAsyncDownload(s *discordgo.Session, i *discordgo.InteractionCreate, ur Success: false, Message: fmt.Sprintf("Download failed: %v", r), URL: url, - Format: audioFormat, + Format: fmt.Sprintf("video: %s, audio: %s", videoFormatID, audioFormatID), Error: fmt.Errorf("%v", r), } } }() // Call downloadVideo (it panics on error instead of returning error) - downloadVideo(outputDir, url, DownloadOptions{EmbedThumbnail: true, IncludeSubtitles: true}, progressChan) + downloadVideo(outputDir, url, DownloadOptions{ + EmbedThumbnail: true, + IncludeSubtitles: true, + VideoFormatID: videoFormatID, + AudioFormatID: audioFormatID, + }, progressChan) // If we reach here, download was successful resultChan <- DownloadResult{ Success: true, Message: "Video Downloaded Successfully!", URL: url, - Format: audioFormat, + Format: fmt.Sprintf("video: %s, audio: %s", videoFormatID, audioFormatID), Error: nil, } }() @@ -64,7 +73,7 @@ func startAsyncDownload(s *discordgo.Session, i *discordgo.InteractionCreate, ur go func() { // First update the original ephemeral message with "Processing..." _, err := s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{ - Content: ptr("🔄 Processing download...\nURL: " + url + "\nAudio: " + audioFormat), + Content: ptr(fmt.Sprintf("🔄 Processing download...\nURL: %s\nVideo: %s\nAudio: %s", url, videoFormatID, audioFormatID)), }) if err != nil { log.Printf("Error updating interaction: %v", err) @@ -188,7 +197,109 @@ func main() { var componentHandlers = map[string]func(s *discordgo.Session, i *discordgo.InteractionCreate){ "video_select": func(s *discordgo.Session, i *discordgo.InteractionCreate) { - // Update components + // Get interaction state with fallback search + var state *InteractionState + + // First try current token + state = getInteractionState(i.Interaction.Token) + + // If not found, try to find state associated with this interaction context + if state == nil { + interactionStatesMutex.RLock() + for token, s := range interactionStates { + if s != nil && s.FormatOptions != nil { + // Found a state, transfer it to current token + state = s + // Copy to current token and clean up old one + interactionStatesMutex.RUnlock() + setInteractionState(i.Interaction.Token, state) + deleteInteractionState(token) + break + } + } + if state == nil { + interactionStatesMutex.RUnlock() + } + } + if state == nil || state.FormatOptions == nil { + s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ + Type: discordgo.InteractionResponseChannelMessageWithSource, + Data: &discordgo.InteractionResponseData{ + Content: "Error: Session expired. Please start over with /download.", + Flags: discordgo.MessageFlagsEphemeral, + }, + }) + return + } + + // Get selected video format + selectedValues := i.MessageComponentData().Values + if len(selectedValues) == 0 { + s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ + Type: discordgo.InteractionResponseChannelMessageWithSource, + Data: &discordgo.InteractionResponseData{ + Content: "Error: No video format selected", + Flags: discordgo.MessageFlagsEphemeral, + }, + }) + return + } + + // Store selected video format ID + state.VideoFormatID = selectedValues[0] + setInteractionState(i.Interaction.Token, state) + + // Build audio format options + audioMenuOptions := []discordgo.SelectMenuOption{} + for _, aOpt := range state.FormatOptions.AudioOptions { + label := aOpt.Format + if aOpt.Language != nil { + label += fmt.Sprintf(" [%s]", *aOpt.Language) + } + if aOpt.TBR != nil { + label += fmt.Sprintf(" (%.0fkbps)", *aOpt.TBR) + } + + // Discord has a 100 char limit on labels + if len(label) > 100 { + label = label[:97] + "..." + } + + audioMenuOptions = append(audioMenuOptions, discordgo.SelectMenuOption{ + Label: label, + Value: aOpt.FormatID, + }) + + // Discord has a limit of 25 options per select menu + if len(audioMenuOptions) >= 25 { + break + } + } + + // Build video format options (to keep them visible but disabled) + videoMenuOptions := []discordgo.SelectMenuOption{} + for _, vOpt := range state.FormatOptions.VideoOptions { + label := fmt.Sprintf("%s (%s", vOpt.Resolution, vOpt.Ext) + if vOpt.TBR != nil { + label += fmt.Sprintf(", %.0fkbps", *vOpt.TBR) + } + label += ")" + + if len(label) > 100 { + label = label[:97] + "..." + } + + videoMenuOptions = append(videoMenuOptions, discordgo.SelectMenuOption{ + Label: label, + Value: vOpt.FormatID, + }) + + if len(videoMenuOptions) >= 25 { + break + } + } + + // Update components - disable video select, enable audio select updatedComponents := []discordgo.MessageComponent{ discordgo.ActionsRow{ Components: []discordgo.MessageComponent{ @@ -197,20 +308,7 @@ func main() { Placeholder: "Choose a video format...", MaxValues: 1, Disabled: true, - Options: []discordgo.SelectMenuOption{ - { - Label: "1920x1080 1080p (mp4, 1910kbps)", - Value: "1080p", - }, - { - Label: "1280x720 720p (mp4, 700kbps)", - Value: "720p", - }, - { - Label: "640x480 480p (mp4, 300kbps)", - Value: "480p", - }, - }, + Options: videoMenuOptions, }, }, }, @@ -221,16 +319,7 @@ func main() { Placeholder: "Choose an audio format...", MaxValues: 1, Disabled: false, - Options: []discordgo.SelectMenuOption{ - { - Label: "Medium", - Value: "medium", - }, - { - Label: "Low", - Value: "low", - }, - }, + Options: audioMenuOptions, }, }, }, @@ -285,21 +374,25 @@ func main() { // Get selected audio format selectedValues := i.MessageComponentData().Values - audioFormat := "medium" // default + audioFormatID := "" if len(selectedValues) > 0 { - audioFormat = selectedValues[0] + audioFormatID = selectedValues[0] } + // Store selected audio format ID + state.AudioFormatID = audioFormatID + response := "" if state.URL != "" { // Respond immediately to prevent timeout - response = "🚀 Starting download...\nURL: " + state.URL + "\nAudio: " + audioFormat + "\n\nYou'll receive an update when the download completes!" + response = fmt.Sprintf("🚀 Starting download...\nURL: %s\nVideo: %s\nAudio: %s\n\nYou'll receive an update when the download completes!", + state.URL, state.VideoFormatID, state.AudioFormatID) // Start async download after responding go func() { // Small delay to ensure response is sent first time.Sleep(100 * time.Millisecond) - startAsyncDownload(s, i, state.URL, audioFormat, out_dir) + startAsyncDownload(s, i, state.URL, state.VideoFormatID, state.AudioFormatID, out_dir) }() // Clean up state after starting download @@ -342,36 +435,144 @@ func main() { return } - // Store URL in interaction state (we'll use the response message ID later) - // For now, store with the current token, then update after getting message ID - setInteractionState(i.Interaction.Token, &InteractionState{URL: url}) - + // Send initial "fetching formats" response err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{ Type: discordgo.InteractionResponseChannelMessageWithSource, Data: &discordgo.InteractionResponseData{ - Content: "Select a video format:", + Content: "🔍 Fetching available formats...", Flags: discordgo.MessageFlagsEphemeral, - Components: []discordgo.MessageComponent{ + }, + }) + if err != nil { + log.Printf("Error: %v", err) + return + } + + // Fetch formats asynchronously + go func() { + formatOptions, err := GetFormats(url) + if err != nil { + _, err = s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{ + Content: ptr("❌ Error fetching formats: " + err.Error()), + }) + if err != nil { + log.Printf("Error updating interaction: %v", err) + } + return + } + + // Sort video formats: highest resolution first, then by bitrate + sort.Slice(formatOptions.VideoOptions, func(i, j int) bool { + // Compare by height (descending) + heightI := 0 + if formatOptions.VideoOptions[i].Height != nil { + heightI = *formatOptions.VideoOptions[i].Height + } + heightJ := 0 + if formatOptions.VideoOptions[j].Height != nil { + heightJ = *formatOptions.VideoOptions[j].Height + } + + if heightI != heightJ { + return heightI > heightJ + } + + // If heights are equal, compare by TBR (descending) + tbrI := 0.0 + if formatOptions.VideoOptions[i].TBR != nil { + tbrI = *formatOptions.VideoOptions[i].TBR + } + tbrJ := 0.0 + if formatOptions.VideoOptions[j].TBR != nil { + tbrJ = *formatOptions.VideoOptions[j].TBR + } + + return tbrI > tbrJ + }) + + // Sort audio formats: highest bitrate first + sort.Slice(formatOptions.AudioOptions, func(i, j int) bool { + tbrI := 0.0 + if formatOptions.AudioOptions[i].TBR != nil { + tbrI = *formatOptions.AudioOptions[i].TBR + } + tbrJ := 0.0 + if formatOptions.AudioOptions[j].TBR != nil { + tbrJ = *formatOptions.AudioOptions[j].TBR + } + + return tbrI > tbrJ + }) + + // Build video format options for Discord select menu + videoMenuOptions := []discordgo.SelectMenuOption{} + for _, vOpt := range formatOptions.VideoOptions { + label := fmt.Sprintf("%s (%s", vOpt.Resolution, vOpt.Ext) + if vOpt.TBR != nil { + label += fmt.Sprintf(", %.0fkbps", *vOpt.TBR) + } + label += ")" + + // Discord has a 100 char limit on labels + if len(label) > 100 { + label = label[:97] + "..." + } + + videoMenuOptions = append(videoMenuOptions, discordgo.SelectMenuOption{ + Label: label, + Value: vOpt.FormatID, + }) + + // Discord has a limit of 25 options per select menu + if len(videoMenuOptions) >= 25 { + break + } + } + + // Build audio format options for Discord select menu + audioMenuOptions := []discordgo.SelectMenuOption{} + for _, aOpt := range formatOptions.AudioOptions { + label := aOpt.Format + if aOpt.Language != nil { + label += fmt.Sprintf(" [%s]", *aOpt.Language) + } + if aOpt.TBR != nil { + label += fmt.Sprintf(" (%.0fkbps)", *aOpt.TBR) + } + + // Discord has a 100 char limit on labels + if len(label) > 100 { + label = label[:97] + "..." + } + + audioMenuOptions = append(audioMenuOptions, discordgo.SelectMenuOption{ + Label: label, + Value: aOpt.FormatID, + }) + + // Discord has a limit of 25 options per select menu + if len(audioMenuOptions) >= 25 { + break + } + } + + // Store format options in interaction state + setInteractionState(i.Interaction.Token, &InteractionState{ + URL: url, + FormatOptions: formatOptions, + }) + + // Update message with format selection menus + _, err = s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{ + Content: ptr("Select a video format:"), + Components: &[]discordgo.MessageComponent{ discordgo.ActionsRow{ Components: []discordgo.MessageComponent{ discordgo.SelectMenu{ CustomID: "video_select", Placeholder: "Choose a video format...", MaxValues: 1, - Options: []discordgo.SelectMenuOption{ - { - Label: "1920x1080 1080p (mp4, 1910kbps)", - Value: "1080p", - }, - { - Label: "1280x720 720p (mp4, 700kbps)", - Value: "720p", - }, - { - Label: "640x480 480p (mp4, 300kbps)", - Value: "480p", - }, - }, + Options: videoMenuOptions, }, }, }, @@ -382,25 +583,16 @@ func main() { Placeholder: "Choose an audio format...", MaxValues: 1, Disabled: true, - Options: []discordgo.SelectMenuOption{ - { - Label: "Medium", - Value: "medium", - }, - { - Label: "Low", - Value: "low", - }, - }, + Options: audioMenuOptions, }, }, }, }, - }, - }) - if err != nil { - log.Printf("Error: %v", err) - } + }) + if err != nil { + log.Printf("Error updating interaction: %v", err) + } + }() }, } diff --git a/app/ytdlp.go b/app/ytdlp.go index 41c1fd5..3d3a5ee 100644 --- a/app/ytdlp.go +++ b/app/ytdlp.go @@ -2,6 +2,7 @@ package main import ( "context" + "encoding/json" "time" "github.com/lrstanley/go-ytdlp" @@ -10,6 +11,29 @@ import ( type DownloadOptions struct { EmbedThumbnail bool IncludeSubtitles bool + VideoFormatID string + AudioFormatID string +} + +type VideoOption struct { + Height *int `json:"height,omitempty"` + Resolution string `json:"resolution,omitempty"` + FormatID string `json:"format_id"` + Ext string `json:"ext"` + TBR *float64 `json:"tbr,omitempty"` +} + +type AudioOption struct { + Format string `json:"format"` + FormatID string `json:"format_id"` + Ext string `json:"ext"` + TBR *float64 `json:"tbr,omitempty"` + Language *string `json:"language,omitempty"` +} + +type FormatOptions struct { + VideoOptions []VideoOption `json:"video_options"` + AudioOptions []AudioOption `json:"audio_options"` } type ProgressUpdate struct { @@ -20,6 +44,89 @@ type ProgressUpdate struct { Phase string } +func GetFormats(url string) (*FormatOptions, error) { + dl := ytdlp.New(). + SkipDownload(). + DumpJSON() + + result, err := dl.Run(context.TODO(), url) + if err != nil { + return nil, err + } + + // Parse the JSON output + var info struct { + Formats []struct { + VCodec *string `json:"vcodec"` + ACodec *string `json:"acodec"` + NeedsTesting *bool `json:"__needs_testing"` + Height *int `json:"height"` + Resolution string `json:"resolution"` + FormatID string `json:"format_id"` + Format string `json:"format"` + Ext string `json:"ext"` + TBR *float64 `json:"tbr"` + Language *string `json:"language"` + LanguagePref *int `json:"language_preference"` + URL *string `json:"url"` + Protocol *string `json:"protocol"` + } `json:"formats"` + } + + if err := json.Unmarshal([]byte(result.Stdout), &info); err != nil { + return nil, err + } + + formatOpts := &FormatOptions{ + VideoOptions: []VideoOption{}, + AudioOptions: []AudioOption{}, + } + + for _, fmt := range info.Formats { + // Skip formats that need testing + if fmt.NeedsTesting != nil && *fmt.NeedsTesting { + continue + } + + // Skip SABR formats (https://github.com/yt-dlp/yt-dlp/issues/12482) + if fmt.URL == nil || *fmt.URL == "" { + continue + } + + // Video-only: has video codec but no audio codec + if fmt.VCodec != nil && *fmt.VCodec != "none" && + fmt.ACodec != nil && *fmt.ACodec == "none" { + formatOpts.VideoOptions = append(formatOpts.VideoOptions, VideoOption{ + Height: fmt.Height, + Resolution: fmt.Resolution, + FormatID: fmt.FormatID, + Ext: fmt.Ext, + TBR: fmt.TBR, + }) + } + + // Audio-only: has audio codec but no video codec + if fmt.ACodec != nil && *fmt.ACodec != "none" && + fmt.VCodec != nil && *fmt.VCodec == "none" { + audioOpt := AudioOption{ + Format: fmt.Format, + FormatID: fmt.FormatID, + Ext: fmt.Ext, + TBR: fmt.TBR, + } + + // Use language if available, otherwise use language_preference + if fmt.Language != nil { + audioOpt.Language = fmt.Language + } + + formatOpts.AudioOptions = append(formatOpts.AudioOptions, audioOpt) + } + } + + return formatOpts, nil +} + func downloadVideo(out_dir, url string, opts DownloadOptions, progressChan chan<- ProgressUpdate) { defer close(progressChan) @@ -27,7 +134,6 @@ func downloadVideo(out_dir, url string, opts DownloadOptions, progressChan chan< dl := ytdlp.New(). SetWorkDir(out_dir). - FormatSort("res,ext:mp4:m4a"). RecodeVideo("mp4"). ProgressFunc(100*time.Millisecond, func(prog ytdlp.ProgressUpdate) { // Detect phase transition -- differentiate "downloading" as the main download @@ -54,6 +160,18 @@ func downloadVideo(out_dir, url string, opts DownloadOptions, progressChan chan< }). Output("%(title)s.%(ext)s") + // Set format if both video and audio format IDs are provided + if opts.VideoFormatID != "" && opts.AudioFormatID != "" { + dl = dl.Format(opts.VideoFormatID + "+" + opts.AudioFormatID) + } else if opts.VideoFormatID != "" { + dl = dl.Format(opts.VideoFormatID) + } else if opts.AudioFormatID != "" { + dl = dl.Format(opts.AudioFormatID) + } else { + // Default format selection if none specified + dl = dl.FormatSort("res,ext:mp4:m4a") + } + if opts.EmbedThumbnail { dl = dl.EmbedThumbnail() }