447 lines
12 KiB
Go
447 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) {
|
|
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: audioFormat,
|
|
Error: fmt.Errorf("%v", r),
|
|
}
|
|
}
|
|
}()
|
|
|
|
// Call downloadVideo (it panics on error instead of returning error)
|
|
downloadVideo(outputDir, url, DownloadOptions{EmbedThumbnail: true, IncludeSubtitles: true}, progressChan)
|
|
|
|
// If we reach here, download was successful
|
|
resultChan <- DownloadResult{
|
|
Success: true,
|
|
Message: "Video Downloaded Successfully!",
|
|
URL: url,
|
|
Format: audioFormat,
|
|
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("🔄 Processing download...\nURL: " + url + "\nAudio: " + audioFormat),
|
|
})
|
|
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) {
|
|
// 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})
|
|
}
|