clipper/backend/youtube/youtube.go

package youtube

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"strconv"

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

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

const (
	SizeOfInt16 = 2

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

// YoutubeClient wraps the youtube.Client client.
type YoutubeClient interface {
	GetVideo(string) (*youtubev2.Video, error)
	GetStream(*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}
}

// 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(videoID string) (*media.MediaSet, error) {
	video, err := d.youtubeClient.GetVideo(videoID)
	if err != nil {
		return nil, fmt.Errorf("error fetching video: %v", err)
	}

	// TODO: improve selection of audio and video format.
	// Perhaps download both separately?
	var format *youtubev2.Format
	for i := range video.Formats {
		candidate := video.Formats[i]
		if candidate.FPS == 0 || candidate.AudioChannels == 0 {
			continue
		}
		if format == nil || (candidate.ContentLength > 0 && candidate.ContentLength < format.ContentLength) {
			format = &candidate
		}
	}
	if format == nil {
		return nil, errors.New("error selecting format: no format available")
	}
	log.Printf("selected format: %+v", format)

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

	mediaSet := media.MediaSet{ID: videoID, Source: "youtube"}

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

	encodedAudioFile, err := os.Create(mediaSet.EncodedAudioPath())
	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.Command("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(mediaSet.RawAudioPath())
	if err != nil {
		return nil, fmt.Errorf("error reading file: %v", err)
	}

	fi, err := rawAudioFile.Stat()
	if err != nil {
		return nil, fmt.Errorf("error reading 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)
	}

	mediaSet.Bytes = fi.Size()
	mediaSet.Channels = format.AudioChannels
	mediaSet.Frames = numFrames
	mediaSet.SampleRate = sampleRate

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

	return &mediaSet, nil
}