639 lines
18 KiB
Go
639 lines
18 KiB
Go
package main
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"os/signal"
|
|
"sort"
|
|
"sync"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/bwmarrin/discordgo"
|
|
)
|
|
|
|
// InteractionState holds the state for a specific interaction
|
|
type InteractionState struct {
|
|
URL string
|
|
FormatOptions *FormatOptions
|
|
VideoFormatID string
|
|
AudioFormatID string
|
|
}
|
|
|
|
// DownloadResult represents the result of an async download operation
|
|
type DownloadResult struct {
|
|
Success bool
|
|
Message string
|
|
URL string
|
|
Format string
|
|
Error error
|
|
}
|
|
|
|
// startAsyncDownload initiates a download in a goroutine and handles progress updates
|
|
func startAsyncDownload(s *discordgo.Session, i *discordgo.InteractionCreate, url, videoFormatID, audioFormatID, outputDir string) {
|
|
progressChan := make(chan ProgressUpdate, 1)
|
|
resultChan := make(chan DownloadResult, 1)
|
|
|
|
// Start download in goroutine
|
|
go func() {
|
|
defer close(resultChan)
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
// Handle panic from downloadVideo
|
|
resultChan <- DownloadResult{
|
|
Success: false,
|
|
Message: fmt.Sprintf("Download failed: %v", r),
|
|
URL: url,
|
|
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,
|
|
VideoFormatID: videoFormatID,
|
|
AudioFormatID: audioFormatID,
|
|
}, progressChan)
|
|
|
|
// If we reach here, download was successful
|
|
resultChan <- DownloadResult{
|
|
Success: true,
|
|
Message: "Video Downloaded Successfully!",
|
|
URL: url,
|
|
Format: fmt.Sprintf("video: %s, audio: %s", videoFormatID, audioFormatID),
|
|
Error: nil,
|
|
}
|
|
}()
|
|
|
|
// Handle progress and results asynchronously
|
|
go func() {
|
|
// First update the original ephemeral message with "Processing..."
|
|
_, err := s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{
|
|
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)
|
|
}
|
|
|
|
for {
|
|
select {
|
|
case prog, ok := <-progressChan:
|
|
if !ok {
|
|
progressChan = nil
|
|
continue
|
|
}
|
|
// Update message w/ phase and real time progress
|
|
phaseEmoji := "⏬"
|
|
if prog.Phase == "post-processing" {
|
|
phaseEmoji = "⚙️"
|
|
}
|
|
content := fmt.Sprintf("%s %s\n%s @ %s [eta: %s]\n📄 %s",
|
|
phaseEmoji,
|
|
prog.Phase,
|
|
prog.Status,
|
|
prog.Percent,
|
|
prog.ETA,
|
|
prog.Filename)
|
|
_, err := s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{
|
|
Content: ptr(content),
|
|
})
|
|
if err != nil {
|
|
log.Printf("Error updating progress: %v", err)
|
|
}
|
|
case result := <-resultChan:
|
|
// Handle completion
|
|
if result.Success {
|
|
_, err = s.FollowupMessageCreate(i.Interaction, false, &discordgo.WebhookParams{
|
|
Content: "📥 Video downloaded: " + result.URL,
|
|
})
|
|
if err != nil {
|
|
log.Printf("Error updating interaction: %v", err)
|
|
}
|
|
} else {
|
|
_, err = s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{
|
|
Content: ptr("❌ Download failed: " + result.Message),
|
|
})
|
|
if err != nil {
|
|
log.Printf("Error updating interaction: %v", err)
|
|
}
|
|
}
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
// Helper function to create string pointer
|
|
func ptr(s string) *string {
|
|
return &s
|
|
}
|
|
|
|
// Global state management
|
|
var (
|
|
interactionStates = make(map[string]*InteractionState)
|
|
interactionStatesMutex = sync.RWMutex{}
|
|
)
|
|
|
|
func getInteractionState(token string) *InteractionState {
|
|
interactionStatesMutex.RLock()
|
|
defer interactionStatesMutex.RUnlock()
|
|
return interactionStates[token]
|
|
}
|
|
|
|
func setInteractionState(token string, state *InteractionState) {
|
|
interactionStatesMutex.Lock()
|
|
defer interactionStatesMutex.Unlock()
|
|
interactionStates[token] = state
|
|
}
|
|
|
|
func deleteInteractionState(token string) {
|
|
interactionStatesMutex.Lock()
|
|
defer interactionStatesMutex.Unlock()
|
|
delete(interactionStates, token)
|
|
}
|
|
|
|
func main() {
|
|
|
|
out_dir := os.Getenv("OUT_PATH")
|
|
bot_token := os.Getenv("DISCORD_TOKEN")
|
|
|
|
if out_dir == "" {
|
|
panic("No output dir specified")
|
|
}
|
|
|
|
s, err := discordgo.New("Bot " + bot_token)
|
|
if err != nil {
|
|
log.Fatalf("Invalid bot parameters: %v", err)
|
|
}
|
|
|
|
err = s.Open()
|
|
if err != nil {
|
|
log.Fatalf("Error opening connection: %v", err)
|
|
}
|
|
|
|
var defaultMemberPermissions int64 = discordgo.PermissionAllText
|
|
var interactionPrivateChannel = discordgo.InteractionContextPrivateChannel
|
|
|
|
var commands = []*discordgo.ApplicationCommand{
|
|
{
|
|
Name: "download",
|
|
Description: "Download video and save it to 'youtube-vids",
|
|
DefaultMemberPermissions: &defaultMemberPermissions,
|
|
Contexts: &[]discordgo.InteractionContextType{interactionPrivateChannel},
|
|
Options: []*discordgo.ApplicationCommandOption{
|
|
{
|
|
Name: "url",
|
|
Description: "URL",
|
|
Type: discordgo.ApplicationCommandOptionString,
|
|
Required: true,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
var componentHandlers = map[string]func(s *discordgo.Session, i *discordgo.InteractionCreate){
|
|
"video_select": func(s *discordgo.Session, i *discordgo.InteractionCreate) {
|
|
// 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{
|
|
discordgo.SelectMenu{
|
|
CustomID: "video_select",
|
|
Placeholder: "Choose a video format...",
|
|
MaxValues: 1,
|
|
Disabled: true,
|
|
Options: videoMenuOptions,
|
|
},
|
|
},
|
|
},
|
|
discordgo.ActionsRow{
|
|
Components: []discordgo.MessageComponent{
|
|
discordgo.SelectMenu{
|
|
CustomID: "audio_select",
|
|
Placeholder: "Choose an audio format...",
|
|
MaxValues: 1,
|
|
Disabled: false,
|
|
Options: audioMenuOptions,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseUpdateMessage,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: "Now select an audio format:",
|
|
Components: updatedComponents,
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
if err != nil {
|
|
log.Printf("Error: %v", err)
|
|
}
|
|
},
|
|
"audio_select": func(s *discordgo.Session, i *discordgo.InteractionCreate) {
|
|
// Get URL from 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.URL != "" {
|
|
// 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 {
|
|
s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseChannelMessageWithSource,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: "Error: No video URL found. Please start over with /download.",
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
// Get selected audio format
|
|
selectedValues := i.MessageComponentData().Values
|
|
audioFormatID := ""
|
|
if len(selectedValues) > 0 {
|
|
audioFormatID = selectedValues[0]
|
|
}
|
|
|
|
// Store selected audio format ID
|
|
state.AudioFormatID = audioFormatID
|
|
|
|
response := ""
|
|
if state.URL != "" {
|
|
// Respond immediately to prevent timeout
|
|
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, state.VideoFormatID, state.AudioFormatID, out_dir)
|
|
}()
|
|
|
|
// Clean up state after starting download
|
|
deleteInteractionState(i.Interaction.Token)
|
|
} else {
|
|
response = "I don't see a video here :("
|
|
}
|
|
err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseChannelMessageWithSource,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: response,
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
if err != nil {
|
|
log.Printf("Error: %v", err)
|
|
}
|
|
},
|
|
}
|
|
|
|
var commandHandlers = map[string]func(s *discordgo.Session, i *discordgo.InteractionCreate){
|
|
"download": func(s *discordgo.Session, i *discordgo.InteractionCreate) {
|
|
options := i.ApplicationCommandData().Options
|
|
optionMap := make(map[string]*discordgo.ApplicationCommandInteractionDataOption, len(options))
|
|
for _, opt := range options {
|
|
optionMap[opt.Name] = opt
|
|
}
|
|
|
|
var url string
|
|
if option, ok := optionMap["url"]; ok {
|
|
url = option.StringValue()
|
|
} else {
|
|
s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseChannelMessageWithSource,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: "Error: No URL provided",
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
return
|
|
}
|
|
|
|
// Send initial "fetching formats" response
|
|
err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
|
|
Type: discordgo.InteractionResponseChannelMessageWithSource,
|
|
Data: &discordgo.InteractionResponseData{
|
|
Content: "🔍 Fetching available formats...",
|
|
Flags: discordgo.MessageFlagsEphemeral,
|
|
},
|
|
})
|
|
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: videoMenuOptions,
|
|
},
|
|
},
|
|
},
|
|
discordgo.ActionsRow{
|
|
Components: []discordgo.MessageComponent{
|
|
discordgo.SelectMenu{
|
|
CustomID: "audio_select",
|
|
Placeholder: "Choose an audio format...",
|
|
MaxValues: 1,
|
|
Disabled: true,
|
|
Options: audioMenuOptions,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
if err != nil {
|
|
log.Printf("Error updating interaction: %v", err)
|
|
}
|
|
}()
|
|
},
|
|
}
|
|
|
|
s.AddHandler(func(s *discordgo.Session, i *discordgo.InteractionCreate) {
|
|
switch i.Type {
|
|
case discordgo.InteractionApplicationCommand:
|
|
if h, ok := commandHandlers[i.ApplicationCommandData().Name]; ok {
|
|
h(s, i)
|
|
}
|
|
if h, ok := componentHandlers[i.ApplicationCommandData().Name]; ok {
|
|
h(s, i)
|
|
}
|
|
case discordgo.InteractionMessageComponent:
|
|
if h, ok := commandHandlers[i.MessageComponentData().CustomID]; ok {
|
|
h(s, i)
|
|
}
|
|
if h, ok := componentHandlers[i.MessageComponentData().CustomID]; ok {
|
|
h(s, i)
|
|
}
|
|
}
|
|
})
|
|
|
|
log.Println("Adding commands")
|
|
registeredCommands := make([]*discordgo.ApplicationCommand, len(commands))
|
|
for i, v := range commands {
|
|
cmd, err := s.ApplicationCommandCreate(s.State.User.ID, "", v)
|
|
if err != nil {
|
|
log.Panicf("Cannot create '%v' command: %v", v.Name, err)
|
|
}
|
|
registeredCommands[i] = cmd
|
|
}
|
|
|
|
log.Println("Bot is now running. Press CTRL+C to exit")
|
|
sc := make(chan os.Signal, 1)
|
|
signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
|
|
<-sc
|
|
|
|
s.Close()
|
|
|
|
//var url string = "https://www.youtube.com/watch?v=WpBWSFF03eI"
|
|
|
|
//downloadVideo(out_dir, url, DownloadOptions{EmbedThumbnail: true, IncludeSubtitles: true})
|
|
}
|