yt-dlp-bot/app/ytdlp.go

package main

import (
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/lrstanley/go-ytdlp"
)

var ytdlpBinary = os.Getenv("YTDLP_BIN")

func GetFormats(url string) (*FormatOptions, error) {
	dl := ytdlp.New().
		SetExecutable(ytdlpBinary).
		SkipDownload().
		DumpJSON()

	result, err := dl.Run(context.TODO(), url)
	if err != nil {
		return nil, err
	}

	// Parse the JSON output
	var info struct {
		Formats []struct {
			VCodec       *string  `json:"vcodec"`
			ACodec       *string  `json:"acodec"`
			NeedsTesting *bool    `json:"__needs_testing"`
			Height       *int     `json:"height"`
			Resolution   string   `json:"resolution"`
			FormatID     string   `json:"format_id"`
			Format       string   `json:"format"`
			Ext          string   `json:"ext"`
			TBR          *float64 `json:"tbr"`
			Language     *string  `json:"language"`
			LanguagePref *int     `json:"language_preference"`
			URL          *string  `json:"url"`
			Protocol     *string  `json:"protocol"`
		} `json:"formats"`
	}

	if err := json.Unmarshal([]byte(result.Stdout), &info); err != nil {
		return nil, err
	}

	formatOpts := &FormatOptions{
		VideoOptions: []VideoOption{},
		AudioOptions: []AudioOption{},
	}

	for _, fmt := range info.Formats {
		// Skip formats that need testing
		if fmt.NeedsTesting != nil && *fmt.NeedsTesting {
			continue
		}

		// Skip SABR formats (https://github.com/yt-dlp/yt-dlp/issues/12482)
		if fmt.URL == nil || *fmt.URL == "" {
			continue
		}

		// Video-only: has video codec but no audio codec
		if fmt.VCodec != nil && *fmt.VCodec != "none" &&
			fmt.ACodec != nil && *fmt.ACodec == "none" {
			formatOpts.VideoOptions = append(formatOpts.VideoOptions, VideoOption{
				Height:     fmt.Height,
				Resolution: fmt.Resolution,
				FormatID:   fmt.FormatID,
				Ext:        fmt.Ext,
				TBR:        fmt.TBR,
			})
		}

		// Audio-only: has audio codec but no video codec
		if fmt.ACodec != nil && *fmt.ACodec != "none" &&
			fmt.VCodec != nil && *fmt.VCodec == "none" {
			audioOpt := AudioOption{
				Format:   fmt.Format,
				FormatID: fmt.FormatID,
				Ext:      fmt.Ext,
				TBR:      fmt.TBR,
			}

			// Use language if available, otherwise use language_preference
			if fmt.Language != nil {
				audioOpt.Language = fmt.Language
			}

			formatOpts.AudioOptions = append(formatOpts.AudioOptions, audioOpt)
		}
	}

	return formatOpts, nil
}

func dirSize(path string) int {
	var total int64
	filepath.Walk(path, func(_ string, info os.FileInfo, err error) error {
		if err == nil && !info.IsDir() {
			total += info.Size()
		}
		return nil
	})
	return int(total)
}

func DownloadVideo(out_dir, temp_dir, url string, opts DownloadOptions, progressChan chan<- ProgressUpdate) {
	var mu sync.Mutex
	currentPhase := "downloading"
	currentFilename := ""

	// Poll the temp directory for actual bytes-on-disk progress.
	// The yt-dlp progress callback only tracks phase/filename since
	// DownloadedBytes from the callback is unreliable for DASH streams.
	var wg sync.WaitGroup
	wg.Add(1)
	stopPoll := make(chan struct{})
	go func() {
		defer wg.Done()
		ticker := time.NewTicker(500 * time.Millisecond)
		defer ticker.Stop()
		for {
			select {
			case <-stopPoll:
				return
			case <-ticker.C:
				mu.Lock()
				phase := currentPhase
				filename := currentFilename
				mu.Unlock()
				progressChan <- ProgressUpdate{
					Phase:           phase,
					DownloadedBytes: dirSize(temp_dir),
					Filename:        filename,
				}
			}
		}
	}()

	defer func() {
		close(stopPoll)
		wg.Wait()
		close(progressChan)
	}()

	homePath := "home:" + out_dir
	tempPath := "temp:" + temp_dir

	dl := ytdlp.New().
		SetExecutable(ytdlpBinary).
		Paths(homePath).
		Paths(tempPath).
		RecodeVideo("mp4").
		ProgressFunc(100*time.Millisecond, func(prog ytdlp.ProgressUpdate) {
			if prog.Status == ytdlp.ProgressStatusFinished ||
				prog.Status == ytdlp.ProgressStatusStarting ||
				prog.Status == ytdlp.ProgressStatusError {
				return
			}

			phase := "downloading"
			if prog.Status == ytdlp.ProgressStatusPostProcessing {
				phase = "post-processing"
			}

			mu.Lock()
			currentPhase = phase
			currentFilename = prog.Filename
			mu.Unlock()
		}).
		Output("%(title)s.%(ext)s")

	// Set format if both video and audio format IDs are provided
	if opts.VideoFormatID != "" && opts.AudioFormatID != "" {
		dl = dl.Format(opts.VideoFormatID + "+" + opts.AudioFormatID)
	} else if opts.VideoFormatID != "" {
		dl = dl.Format(opts.VideoFormatID)
	} else if opts.AudioFormatID != "" {
		dl = dl.Format(opts.AudioFormatID)
	} else {
		// Default format selection if none specified
		dl = dl.FormatSort("res,ext:mp4:m4a")
	}

	if opts.EmbedThumbnail {
		dl = dl.EmbedThumbnail()
	}

	if opts.IncludeSubtitles {
		dl = dl.CompatOptions("no-keep-subs")
		dl = dl.EmbedSubs()
		dl = dl.SubLangs("en,en*")
		dl = dl.WriteAutoSubs()
		dl = dl.WriteSubs()
	}

	_, err := dl.Run(context.TODO(), url)
	if err != nil {
		panic(err)
	}
}