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"
)

// Progress implements ctxprogress.Progress.
func (p Entry) 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) ([]Entry, 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 := []Entry{}
	for line := range lines {
		entry := Entry{}
		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 = (*Entry)(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
}