Files
yt-dlp-bot/app/main.go

428 lines
12 KiB
Go

package main
import (
"fmt"
"log"
"os"
"os/signal"
"sync"
"syscall"
"time"
"github.com/bwmarrin/discordgo"
)
// InteractionState holds the state for a specific interaction
type InteractionState struct {
URL 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, audioFormat, outputDir string) {
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: audioFormat,
Error: fmt.Errorf("%v", r),
}
}
}()
// Call downloadVideo (it panics on error instead of returning error)
downloadVideo(outputDir, url, DownloadOptions{EmbedThumbnail: true, IncludeSubtitles: true})
// If we reach here, download was successful
resultChan <- DownloadResult{
Success: true,
Message: "Video Downloaded Successfully!",
URL: url,
Format: audioFormat,
Error: nil,
}
}()
// Handle results asynchronously
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),
})
if err != nil {
log.Printf("Error updating interaction: %v", err)
}
result := <-resultChan
if result.Success {
// Update ephemeral message with completion status
_, err = s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{
Content: ptr("✅ Download completed successfully!\nURL: " + result.URL + "\nAudio: " + result.Format),
})
if err != nil {
log.Printf("Error updating interaction: %v", err)
}
// Send non-ephemeral completion message
_, err = s.FollowupMessageCreate(i.Interaction, false, &discordgo.WebhookParams{
Content: "📥 Video downloaded: " + result.URL,
})
if err != nil {
log.Printf("Error sending public completion message: %v", err)
}
} else {
// Update ephemeral message with error
_, err = s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{
Content: ptr("❌ Download failed: " + result.Message + "\nURL: " + result.URL + "\nAudio: " + result.Format),
})
if err != nil {
log.Printf("Error updating interaction: %v", err)
}
}
}()
}
// 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) {
// Update components
updatedComponents := []discordgo.MessageComponent{
discordgo.ActionsRow{
Components: []discordgo.MessageComponent{
discordgo.SelectMenu{
CustomID: "video_select",
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",
},
},
},
},
},
discordgo.ActionsRow{
Components: []discordgo.MessageComponent{
discordgo.SelectMenu{
CustomID: "audio_select",
Placeholder: "Choose an audio format...",
MaxValues: 1,
Disabled: false,
Options: []discordgo.SelectMenuOption{
{
Label: "Medium",
Value: "medium",
},
{
Label: "Low",
Value: "low",
},
},
},
},
},
}
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
audioFormat := "medium" // default
if len(selectedValues) > 0 {
audioFormat = selectedValues[0]
}
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!"
// 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)
}()
// 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
}
// 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})
err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
Type: discordgo.InteractionResponseChannelMessageWithSource,
Data: &discordgo.InteractionResponseData{
Content: "Select a video format:",
Flags: discordgo.MessageFlagsEphemeral,
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",
},
},
},
},
},
discordgo.ActionsRow{
Components: []discordgo.MessageComponent{
discordgo.SelectMenu{
CustomID: "audio_select",
Placeholder: "Choose an audio format...",
MaxValues: 1,
Disabled: true,
Options: []discordgo.SelectMenuOption{
{
Label: "Medium",
Value: "medium",
},
{
Label: "Low",
Value: "low",
},
},
},
},
},
},
},
})
if err != nil {
log.Printf("Error: %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})
}