tstor/pkg/ytdlp/playlist.go

package ytdlp

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"io"
	"os"
	"os/exec"
	"strings"

	"github.com/royalcat/ctxprogress"
	"golang.org/x/sync/errgroup"
)

type PlaylistEntry struct {
	ID                 string       `json:"id"`
	Uploader           string       `json:"uploader"`
	UploaderID         string       `json:"uploader_id"`
	UploadDate         string       `json:"upload_date"`
	Title              string       `json:"title"`
	Thumbnail          string       `json:"thumbnail"`
	Duration           int64        `json:"duration"`
	LikeCount          int64        `json:"like_count"`
	DislikeCount       int64        `json:"dislike_count"`
	CommentCount       int64        `json:"comment_count"`
	Formats            []Format     `json:"formats"`
	AgeLimit           int64        `json:"age_limit"`
	Tags               []string     `json:"tags"`
	Categories         []string     `json:"categories"`
	Cast               []any        `json:"cast"`
	Subtitles          Subtitles    `json:"subtitles"`
	Thumbnails         []Thumbnail  `json:"thumbnails"`
	Timestamp          int64        `json:"timestamp"`
	ViewCount          int64        `json:"view_count"`
	WebpageURL         string       `json:"webpage_url"`
	OriginalURL        string       `json:"original_url"`
	WebpageURLBasename string       `json:"webpage_url_basename"`
	WebpageURLDomain   string       `json:"webpage_url_domain"`
	Extractor          string       `json:"extractor"`
	ExtractorKey       string       `json:"extractor_key"`
	PlaylistCount      int64        `json:"playlist_count"`
	Playlist           string       `json:"playlist"`
	PlaylistID         string       `json:"playlist_id"`
	PlaylistTitle      string       `json:"playlist_title"`
	PlaylistUploader   string       `json:"playlist_uploader"`
	PlaylistUploaderID string       `json:"playlist_uploader_id"`
	NEntries           int64        `json:"n_entries"`
	PlaylistIndex      int64        `json:"playlist_index"`
	PlaylistAutonumber int64        `json:"playlist_autonumber"`
	DisplayID          string       `json:"display_id"`
	Fulltitle          string       `json:"fulltitle"`
	DurationString     string       `json:"duration_string"`
	ReleaseYear        int          `json:"release_year"`
	Epoch              int64        `json:"epoch"`
	FormatID           string       `json:"format_id"`
	URL                string       `json:"url"`
	ManifestURL        string       `json:"manifest_url"`
	Tbr                float64      `json:"tbr"`
	EXT                EXT          `json:"ext"`
	FPS                float64      `json:"fps"`
	Protocol           Protocol     `json:"protocol"`
	VideoHasDRM        bool         `json:"has_drm"`
	Width              int64        `json:"width"`
	Height             int64        `json:"height"`
	Vcodec             string       `json:"vcodec"`
	Acodec             string       `json:"acodec"`
	DynamicRange       DynamicRange `json:"dynamic_range"`
	Resolution         string       `json:"resolution"`
	AspectRatio        float64      `json:"aspect_ratio"`
	HTTPHeaders        HTTPHeaders  `json:"http_headers"`
	VideoEXT           EXT          `json:"video_ext"`
	AudioEXT           AudioEXT     `json:"audio_ext"`
	Format             string       `json:"format"`
	Filename           string       `json:"_filename"`
	VideoFilename      string       `json:"filename"`
	Type               string       `json:"_type"`
	Version            Version      `json:"_version"`
}

// Progress implements ctxprogress.Progress.
func (p PlaylistEntry) Progress() (current int, total int) {
	return int(p.PlaylistIndex), int(p.PlaylistCount)
}

func (p PlaylistEntry) Url() string {
	if p.URL != "" {
		return p.URL
	}
	if p.WebpageURL != "" {
		return p.WebpageURL
	}
	if p.OriginalURL != "" {
		return p.OriginalURL
	}

	return ""
}

func (yt *Client) Playlist(ctx context.Context, url string) ([]PlaylistEntry, error) {
	group, ctx := errgroup.WithContext(ctx)
	w, lines, err := lineReader(group)
	if err != nil {
		return nil, err
	}

	cmd := exec.CommandContext(ctx, yt.binary, "-j", url)

	cmd.Stdout = w

	group.Go(func() error {
		err := cmd.Run()
		if err != nil {
			return err
		}
		return w.Close()
	})

	playlists := []PlaylistEntry{}
	for line := range lines {
		entry := PlaylistEntry{}
		err = json.Unmarshal([]byte(line), &entry)
		if err != nil {
			return nil, err
		}

		playlists = append(playlists, entry)
	}

	return playlists, nil
}

func lineReader(group *errgroup.Group) (io.WriteCloser, <-chan string, error) {
	lines := make(chan string)
	var r io.Reader
	r, w := io.Pipe()

	r = io.TeeReader(r, os.Stdout)

	group.Go(func() error {
		defer close(lines)

		reader := bufio.NewReader(r)
		for {
			line, err := reader.ReadString('\n')
			if err != nil {
				if err == io.EOF {
					return w.Close()
				}
				return errors.Join(err, w.Close())
			}
			line = strings.Trim(line, " \r\t")

			if line == "" {
				continue
			}

			lines <- line
		}
	})

	return w, lines, nil
}

var _ ctxprogress.Progress = (*PlaylistEntry)(nil)

var _ ctxprogress.Progress = (*DownloadProgress)(nil)

func parseProgress(line string) (ctxprogress.Progress, error) {
	line = strings.Trim(line, " \r\t")

	p := &DownloadProgress{}
	err := json.Unmarshal([]byte(line), p)
	if err != nil {
		return nil, err
	}

	return p, nil
}