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

	EncodedAudioCodec      = "pcm_s16le"
	EncodedAudioFormat     = "s16le"
	EncodedAudioSampleRate = 48000

	thumbnailWidth  = 30
	thumbnailHeight = 100
	videoItag       = 18
)

// 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
}

func NewDownloader(youtubeClient YoutubeClient) *Downloader {
	return &Downloader{youtubeClient: youtubeClient}
}

type audioResult struct {
	*media.Audio
	err error
}

// videoMediaSet represents the video part of a media.MediaSet:
type videoMediaSet struct {
	bytes int64
}

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.MediaSet{ID: videoID, Source: "youtube"}

	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.VideoPath(), 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) {
	var format *youtubev2.Format
	for _, candidate := range video.Formats.WithAudioChannels() {
		if format == nil || (candidate.ContentLength > 0 && candidate.ContentLength < format.ContentLength) {
			candidate := candidate
			format = &candidate
		}
	}
	if format == nil {
		return nil, errors.New("error selecting audio format: no format available")
	}
	log.Printf("selected audio format: %+v", format)

	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", EncodedAudioFormat, "-ar", strconv.Itoa(EncodedAudioSampleRate), "-acodec", EncodedAudioCodec, "-")
	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(seconds int) (int, int) {
	x := int(math.Floor(math.Sqrt(float64(seconds))))
	if x*x < seconds {
		return x + 1, x
	}
	return x, x
}

func (d *Downloader) downloadVideo(ctx context.Context, video *youtubev2.Video, outPath, thumbnailOutPath string) (*media.Video, error) {
	// TODO: check if iTag 18 always exists, and works in a good variety of browsers.
	format := video.Formats.FindByItag(videoItag)
	log.Printf("selected video format: %+v", format)

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

	videoFile, err := os.Create(outPath)
	if err != nil {
		return nil, fmt.Errorf("error creating video file: %v", err)
	}

	streamReader := io.TeeReader(stream, videoFile)
	durationMsecs, err := strconv.Atoi(format.ApproxDurationMs)
	if err != nil {
		return nil, fmt.Errorf("could not parse video duration: %s", err)
	}
	durationSecs := durationMsecs / 1000
	gridSizeX, gridSizeY := thumbnailGridSize(durationSecs)

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

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

	duration := time.Duration(durationMsecs) * time.Millisecond

	return &media.Video{
		Bytes:           format.ContentLength,
		ThumbnailWidth:  thumbnailWidth,
		ThumbnailHeight: thumbnailHeight,
		Duration:        duration,
	}, nil
}