clipper/backend/youtube/youtube.go

package youtube

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"math"
	"os"
	"os/exec"
	"strconv"
	"sync"
	"time"

	"git.netflux.io/rob/clipper/media"

	youtubev2 "github.com/kkdai/youtube/v2"
)

const (
	SizeOfInt16 = 2

	rawAudioCodec      = "pcm_s16le"
	rawAudioFormat     = "s16le"
	rawAudioSampleRate = 48000

	thumbnailPrescaleWidth  = -1
	thumbnailPrescaleHeight = 120
	thumbnailWidth          = 30
	thumbnailHeight         = 100
)

// YoutubeClient wraps the youtube.Client client.
type YoutubeClient interface {
	GetVideoContext(context.Context, string) (*youtubev2.Video, error)
	GetStreamContext(context.Context, *youtubev2.Video, *youtubev2.Format) (io.ReadCloser, int64, error)
}

// Downloader downloads a set of Youtube media for a given video ID, including
// separate audio and video files and a JSON metadata file. Additionally, it
// also renders the downloaded audio file as a raw audio file.
type Downloader struct {
	youtubeClient YoutubeClient
}

// NewDownloader returns a new *Downloader.
func NewDownloader(youtubeClient YoutubeClient) *Downloader {
	return &Downloader{youtubeClient: youtubeClient}
}

type audioResult struct {
	*media.Audio
	err error
}

type videoResult struct {
	*media.Video
	err error
}

// Download downloads the relevant audio and video files for the provided
// Youtube video ID. If successful, a *media.MediaSet struct containing
// metadata about the downloaded items is returned.
func (d *Downloader) Download(ctx context.Context, videoID string) (*media.MediaSet, error) {
	var video *youtubev2.Video
	video, err := d.youtubeClient.GetVideoContext(ctx, videoID)
	if err != nil {
		return nil, fmt.Errorf("error fetching video: %v", err)
	}

	mediaSet := media.NewMediaSet(videoID)

	audioResultChan := make(chan audioResult, 1)
	videoResultChan := make(chan videoResult, 1)

	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer close(audioResultChan)
		audio, audioErr := d.downloadAudio(ctx, video, mediaSet.EncodedAudioPath(), mediaSet.RawAudioPath())
		result := audioResult{audio, audioErr}
		audioResultChan <- result
		wg.Done()
	}()
	go func() {
		defer close(videoResultChan)
		video, videoErr := d.downloadVideo(ctx, video, mediaSet.ThumbnailPath())
		result := videoResult{video, videoErr}
		videoResultChan <- result
		wg.Done()
	}()

	wg.Wait()

	audioResult := <-audioResultChan
	videoResult := <-videoResultChan

	if err = audioResult.err; err != nil {
		return nil, fmt.Errorf("error downloading audio: %v", err)
	}
	if err = videoResult.err; err != nil {
		return nil, fmt.Errorf("error downloading video: %v", err)
	}

	mediaSet.Audio = *audioResult.Audio
	mediaSet.Video = *videoResult.Video

	metadataFile, err := os.Create(mediaSet.MetadataPath())
	if err != nil {
		return nil, fmt.Errorf("error opening metadata file: %v", err)
	}

	if err = json.NewEncoder(metadataFile).Encode(mediaSet); err != nil {
		return nil, fmt.Errorf("error encoding metadata: %v", err)
	}

	if err = metadataFile.Close(); err != nil {
		return nil, fmt.Errorf("error writing metadata file: %v", err)
	}

	log.Println("finished downloading mediaset")

	return mediaSet, nil
}

func (d *Downloader) downloadAudio(ctx context.Context, video *youtubev2.Video, outPath, rawOutPath string) (*media.Audio, error) {
	if len(video.Formats) == 0 {
		return nil, errors.New("error selecting audio format: no format available")
	}
	format := SortAudio(video.Formats)[0]
	log.Printf("selected audio format: %s", FormatDebugString(&format, false))

	stream, _, err := d.youtubeClient.GetStreamContext(ctx, video, &format)
	if err != nil {
		return nil, fmt.Errorf("error fetching audio stream: %v", err)
	}

	rawAudioFile, err := os.Create(rawOutPath)
	if err != nil {
		return nil, fmt.Errorf("error creating raw audio file: %v", err)
	}

	encodedAudioFile, err := os.Create(outPath)
	if err != nil {
		return nil, fmt.Errorf("error creating encoded audio file: %v", err)
	}
	streamReader := io.TeeReader(stream, encodedAudioFile)

	var errOut bytes.Buffer
	cmd := exec.CommandContext(ctx, "ffmpeg", "-i", "-", "-f", rawAudioFormat, "-ar", strconv.Itoa(rawAudioSampleRate), "-acodec", rawAudioCodec, "-")
	cmd.Stdin = streamReader
	cmd.Stdout = rawAudioFile
	cmd.Stderr = &errOut

	if err = cmd.Run(); err != nil {
		log.Println(errOut.String())
		return nil, fmt.Errorf("error processing audio: %v", err)
	}

	if err = rawAudioFile.Close(); err != nil {
		return nil, fmt.Errorf("error writing raw audio file: %v", err)
	}

	rawAudioFile, err = os.Open(rawOutPath)
	if err != nil {
		return nil, fmt.Errorf("error opening raw audio file: %v", err)
	}

	fi, err := rawAudioFile.Stat()
	if err != nil {
		return nil, fmt.Errorf("error reading raw audio file info: %v", err)
	}

	numFrames := fi.Size() / int64(SizeOfInt16) / int64(format.AudioChannels)
	sampleRate, err := strconv.Atoi(format.AudioSampleRate)
	if err != nil {
		return nil, fmt.Errorf("invalid samplerate: %s", format.AudioSampleRate)
	}

	return &media.Audio{
		Bytes:      fi.Size(),
		Channels:   format.AudioChannels,
		Frames:     numFrames,
		SampleRate: sampleRate,
	}, nil
}

func thumbnailGridSize(msecs int) (int, int) {
	secs := msecs / 1000
	x := int(math.Floor(math.Sqrt(float64(secs))))
	if x*x < secs {
		return x + 1, x
	}
	return x, x
}

func (d *Downloader) downloadVideo(ctx context.Context, video *youtubev2.Video, thumbnailOutPath string) (*media.Video, error) {
	if len(video.Formats) == 0 {
		return nil, errors.New("error selecting audio format: no format available")
	}
	format := SortVideo(video.Formats)[0]
	log.Printf("selected video format: %s", FormatDebugString(&format, false))

	stream, _, err := d.youtubeClient.GetStreamContext(ctx, video, &format)
	if err != nil {
		return nil, fmt.Errorf("error fetching video stream: %v", err)
	}

	durationMsecs, err := strconv.Atoi(format.ApproxDurationMs)
	if err != nil {
		return nil, fmt.Errorf("could not parse video duration: %s", err)
	}
	gridSizeX, gridSizeY := thumbnailGridSize(durationMsecs)

	var errOut bytes.Buffer
	cmd := exec.CommandContext(
		ctx,
		"ffmpeg",
		"-i",
		"-",
		"-vf",
		fmt.Sprintf("fps=1,scale=%d:%d,crop=%d:%d,tile=%dx%d", thumbnailPrescaleWidth, thumbnailPrescaleHeight, thumbnailWidth, thumbnailHeight, gridSizeX, gridSizeY),
		"-f",
		"image2pipe",
		"-vsync",
		"0",
		thumbnailOutPath,
	)
	cmd.Stdin = stream
	cmd.Stderr = &errOut

	if err = cmd.Run(); err != nil {
		log.Println(errOut.String())
		return nil, fmt.Errorf("error processing video: %v", err)
	}

	return &media.Video{
		Bytes:           format.ContentLength,
		ThumbnailWidth:  thumbnailWidth,
		ThumbnailHeight: thumbnailHeight,
		Duration:        time.Duration(durationMsecs) * time.Millisecond,
	}, nil
}