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 Client struct {
	binary string
}

func New() (*Client, error) {
	return &Client{binary: "yt-dlp"}, nil
}

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 DownloadPlaylist(ctx context.Context, url string, dir string) error {
// 	args := []string{
// 		"--no-simulate", "-j",
// 		"--progress", "--newline", "--progress-template", progressTemplate,
// 		"-o", path.Join(dir, "%(title)s.%(ext)s"),
// 		url,
// 	}

// 	group, ctx := errgroup.WithContext(ctx)

// 	pr, w := io.Pipe()
// 	cmd := exec.CommandContext(ctx, "yt-dlp", args...)
// 	cmd.Stdout = w

// 	r := io.TeeReader(pr, os.Stdout)

// 	group.Go(func() error {
// 		reader := bufio.NewReader(r)
// 		for {
// 			line, err := reader.ReadString('\n')
// 			if err != nil {
// 				if err == io.EOF {
// 					return nil
// 				}
// 				return err
// 			}
// 			line = strings.Trim(line, " \r\t")
// 			if len(line) == 0 {
// 				continue
// 			}
// 			if strings.HasPrefix(line, "{") {
// 				item := &PlaylistEntry{}
// 				err = json.Unmarshal([]byte(line), &item)
// 				if err != nil {
// 					return err
// 				}
// 			} else if body, ok := strings.CutPrefix(line, "%"); ok {
// 				p := &DownloadProgress{}
// 				err = json.Unmarshal([]byte(body), &p)
// 				if err != nil {
// 					return err
// 				}
// 			} else {
// 				return fmt.Errorf("Failed to parse output, unkonow first symbol: %v", string([]rune(line)[0]))
// 			}
// 		}
// 	})

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

// 	return group.Wait()
// }

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
}