package main

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

	"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,
				},
			},
		},
	}

	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, temp_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()
}