yt-dlp-bot/app/main.go

package main

import (
	"fmt"
	"log"
	"os"
	"os/signal"
	"sort"
	"syscall"

	"github.com/bwmarrin/discordgo"
)

func main() {
	out_dir := os.Getenv("OUT_PATH")
	temp_dir := os.Getenv("TEMP_PATH")
	bot_token := os.Getenv("DISCORD_TOKEN")

	if out_dir == "" {
		panic("No output dir specified")
	}

	if temp_dir == "" {
		panic("No temp 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,
				},
			},
		},
		{
			Name:                     "version",
			Description:              "Show application version",
			DefaultMemberPermissions: &defaultMemberPermissions,
			Contexts:                 &[]discordgo.InteractionContextType{interactionPrivateChannel},
		},
		{
			Name: "download video",
			Type: discordgo.MessageApplicationCommand,
			IntegrationTypes: &[]discordgo.ApplicationIntegrationType{
				discordgo.ApplicationIntegrationUserInstall,
			},

			Contexts: &[]discordgo.InteractionContextType{
				discordgo.InteractionContextBotDM,
				discordgo.InteractionContextPrivateChannel,
			},
		},
	}

	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 requester user ID
			state.Requester = i.User.ID

			// 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 and name
			state.VideoFormatID = selectedValues[0]
			for _, vOpt := range state.FormatOptions.VideoOptions {
				if vOpt.FormatID == selectedValues[0] {
					label := fmt.Sprintf("%s (%s", vOpt.Resolution, vOpt.Ext)
					if vOpt.TBR != nil {
						label += fmt.Sprintf(", %.0fkbps", *vOpt.TBR)
					}
					label += ")"
					state.VideoFormatName = label
					break
				}
			}
			setInteractionState(i.Interaction.Token, state)

			// Build audio format options
			audioMenuOptions := buildAudioMenuOptions(state.FormatOptions.AudioOptions)

			// Build video format options (to keep them visible but disabled)
			videoMenuOptions := buildVideoMenuOptions(state.FormatOptions.VideoOptions)

			// 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.InteractionResponseUpdateMessage,
					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 and name
			state.AudioFormatID = audioFormatID
			for _, aOpt := range state.FormatOptions.AudioOptions {
				if aOpt.FormatID == audioFormatID {
					label := aOpt.Format
					if aOpt.Language != nil {
						label += fmt.Sprintf(" [%s]", *aOpt.Language)
					}
					if aOpt.TBR != nil {
						label += fmt.Sprintf(" (%.0fkbps)", *aOpt.TBR)
					}
					state.AudioFormatName = label
					break
				}
			}

			response := ""
			if state.URL != "" {
				// Respond immediately to prevent timeout
				response = fmt.Sprintf("%s **Starting download**", loading_emoji)

				// Clean up state before responding
				deleteInteractionState(i.Interaction.Token)
			} else {
				response = "I don't see a video here :("
			}
			err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
				Type: discordgo.InteractionResponseUpdateMessage,
				Data: &discordgo.InteractionResponseData{
					Content: response,
					Flags:   discordgo.MessageFlagsEphemeral,
				},
			})
			if err != nil {
				log.Printf("Error: %v", err)
			}
			if state.URL != "" {
				go startAsyncDownload(s, i, state.Requester, state.URL, state.VideoFormatID, state.VideoFormatName, state.AudioFormatID, state.AudioFormatName, out_dir, temp_dir)
			}
		},
	}

	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: fmt.Sprintf("%s Fetching available formats...", loading_emoji),
					Flags:   discordgo.MessageFlagsEphemeral,
				},
			})
			if err != nil {
				log.Printf("Error: %v", err)
				return
			}

			// Fetch formats asynchronously
			go fetchAndShowFormats(s, i, url)
		},
		"download video": func(s *discordgo.Session, i *discordgo.InteractionCreate) {
			data := i.ApplicationCommandData()
			targetMsg, ok := data.Resolved.Messages[data.TargetID]
			if !ok || targetMsg == nil {
				s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
					Type: discordgo.InteractionResponseChannelMessageWithSource,
					Data: &discordgo.InteractionResponseData{
						Content: "Error: Could not find the target message",
						Flags:   discordgo.MessageFlagsEphemeral,
					},
				})
				return
			}
			url := extractURLFromString(targetMsg.Content)
			if url == "" {
				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: fmt.Sprintf("%s Fetching available formats...", loading_emoji),
					Flags:   discordgo.MessageFlagsEphemeral,
				},
			})
			if err != nil {
				log.Printf("Error: %v", err)
				return
			}

			// Fetch formats asynchronously
			go fetchAndShowFormats(s, i, url)
		},
		"version": func(s *discordgo.Session, i *discordgo.InteractionCreate) {
			s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
				Type: discordgo.InteractionResponseChannelMessageWithSource,
				Data: &discordgo.InteractionResponseData{
					Content: "[yt-dlp-bot](https://git.dubyatp.xyz/williamp/yt-dlp-bot) by dubyatp",
					Flags:   discordgo.MessageFlagsEphemeral,
				},
			})
		},
	}

	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)
			}
		case discordgo.InteractionMessageComponent:
			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()
}

func buildVideoMenuOptions(videoOptions []VideoOption) []discordgo.SelectMenuOption {
	opts := make([]discordgo.SelectMenuOption, 0, 25)
	for _, vOpt := range 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] + "..."
		}
		opts = append(opts, discordgo.SelectMenuOption{
			Label: label,
			Value: vOpt.FormatID,
		})
		if len(opts) >= 25 {
			break
		}
	}
	return opts
}

func buildAudioMenuOptions(audioOptions []AudioOption) []discordgo.SelectMenuOption {
	opts := make([]discordgo.SelectMenuOption, 0, 25)
	for _, aOpt := range audioOptions {
		label := aOpt.Format
		if aOpt.Language != nil {
			label += fmt.Sprintf(" [%s]", *aOpt.Language)
		}
		if aOpt.TBR != nil {
			label += fmt.Sprintf(" (%.0fkbps)", *aOpt.TBR)
		}
		if len(label) > 100 {
			label = label[:97] + "..."
		}
		opts = append(opts, discordgo.SelectMenuOption{
			Label: label,
			Value: aOpt.FormatID,
		})
		if len(opts) >= 25 {
			break
		}
	}
	return opts
}

func fetchAndShowFormats(s *discordgo.Session, i *discordgo.InteractionCreate, url string) {
	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(x, y int) bool {
		heightX, heightY := 0, 0
		if formatOptions.VideoOptions[x].Height != nil {
			heightX = *formatOptions.VideoOptions[x].Height
		}
		if formatOptions.VideoOptions[y].Height != nil {
			heightY = *formatOptions.VideoOptions[y].Height
		}
		if heightX != heightY {
			return heightX > heightY
		}
		tbrX, tbrY := 0.0, 0.0
		if formatOptions.VideoOptions[x].TBR != nil {
			tbrX = *formatOptions.VideoOptions[x].TBR
		}
		if formatOptions.VideoOptions[y].TBR != nil {
			tbrY = *formatOptions.VideoOptions[y].TBR
		}
		return tbrX > tbrY
	})

	// Sort audio formats: highest bitrate first
	sort.Slice(formatOptions.AudioOptions, func(x, y int) bool {
		tbrX, tbrY := 0.0, 0.0
		if formatOptions.AudioOptions[x].TBR != nil {
			tbrX = *formatOptions.AudioOptions[x].TBR
		}
		if formatOptions.AudioOptions[y].TBR != nil {
			tbrY = *formatOptions.AudioOptions[y].TBR
		}
		return tbrX > tbrY
	})

	videoMenuOptions := buildVideoMenuOptions(formatOptions.VideoOptions)
	audioMenuOptions := buildAudioMenuOptions(formatOptions.AudioOptions)

	if len(videoMenuOptions) == 0 || len(audioMenuOptions) == 0 {
		_, err = s.InteractionResponseEdit(i.Interaction, &discordgo.WebhookEdit{
			Content: ptr("❌ No separate video/audio streams found for this URL. The source may only provide combined formats."),
		})
		if err != nil {
			log.Printf("Error updating interaction: %v", err)
		}
		return
	}

	setInteractionState(i.Interaction.Token, &InteractionState{
		URL:           url,
		FormatOptions: formatOptions,
	})

	_, 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)
	}
}