youtubedl.go

package youtubedl

import (
	"context"
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"github.com/asticode/go-astisub"
	log "github.com/sirupsen/logrus"
	"gitlab.ethz.ch/chgerber/MessageComposition/src/pkg/util"
	"gitlab.ethz.ch/chgerber/annotation/v2"
	"gitlab.ethz.ch/chgerber/monitor"
	"go.mongodb.org/mongo-driver/mongo"
	"google.golang.org/api/googleapi/transport"
	"google.golang.org/api/option"
	"google.golang.org/api/youtube/v3"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
)

// Subtitle is a wrapper type for YouTube subtitles that can be parsed by astisub
type Subtitle struct {
	*astisub.Subtitles
}

// Annotations turns the subtitles into annotations with the src field set as given
// Method of the annotation.Caption IF
func (s *Subtitle) Annotations(src string) []*annotation.Annotation {
	return annotation.SubToAnnotation(s.Subtitles, src)
}

// SubtitleFormat type describes a subtitle type (e.g. `vtt`)
type SubtitleFormat string

func (sf SubtitleFormat) String() string {
	return string(sf)
}

const (
	// SubVTT caption of type vtt
	SubVTT SubtitleFormat = "vtt"
	// SubSRT caption of type srt
	SubSRT SubtitleFormat = "srt"
	// SubSSA caption of type ssa
	SubSSA SubtitleFormat = "ssa"
	// SubSTL caption of type stl
	SubSTL SubtitleFormat = "stl"
	// SubTTML caption of type ttml
	SubTTML SubtitleFormat = "ttml"
	// SubSRV3 caption of the youtube xml timedtext type with file extension .srv3
	SubSRV3 SubtitleFormat = "srv3"
)

// SupportedSubs is a list in preferred order of subtitle types that are supported to be parsed
var SupportedSubs = []SubtitleFormat{
	SubSRV3,
	SubSRT,
	SubTTML,
	SubSSA,
	SubSTL,
	//SubVTT TODO Timing does not work with YouTube subs of this type. -> Fix parsing
}

// ErrYoutubeDLLangNotSupported is thrown when language not supported
type ErrYoutubeDLLangNotSupported string

func (e ErrYoutubeDLLangNotSupported) Error() string {
	return string(e)
}

// ErrYoutubeVideoUnavailable is thrown when language not supported
type ErrYoutubeVideoUnavailable string

func (e ErrYoutubeVideoUnavailable) Error() string {
	return string(e)
}

// ErrYoutubeDL are errors with related to the youtubedl pkg
type ErrYoutubeDL string

func (e ErrYoutubeDL) Error() string {
	return string(e)
}

func runYoutubeDL(args []string) (output []byte, err error) {
	cmd := exec.Command("youtube-dl", args...)
	log.Debug(args)
	output, err = cmd.CombinedOutput()
	log.Trace(string(output))
	if err != nil && strings.Contains(string(output), "video is unavailable") {
		return nil, ErrYoutubeVideoUnavailable(string(output))
	}
	return output, err
}

func getVideoInfo(id string) (info map[string]interface{}, err error) {

	args := []string{
		"--skip-download",
		"-J",
		"https://youtu.be/" + id,
	}

	output, err := runYoutubeDL(args)
	if err != nil {
		return info, err
	}

	lines := strings.Split(strings.Trim(string(output), "\n\t "), "\n")

	if len(lines) == 0 {
		return nil, ErrYoutubeDL("youtube-dl response decoding error")
	}

	err = json.Unmarshal([]byte(lines[len(lines)-1]), &info)
	if err != nil {
		return info, err
	}

	return info, err
}

// SubList stores two  maps[language][sub-formats]
// Automatic Captions are the automatically generated youtube captions
// Subtitles are the proper Subtitles
type SubList struct {
	AutomaticCaptions map[string][]string
	Subtitles         map[string][]string
}

// NewSubList creates a new reference to a SubList with empty members.
func NewSubList() *SubList {
	return &SubList{
		AutomaticCaptions: make(map[string][]string),
		Subtitles:         make(map[string][]string),
	}
}

func listSubs(id string) (*SubList, error) {

	subs := NewSubList()

	args := []string{
		"--skip-download",
		"--list-subs",
		"https://youtu.be/" + id,
	}
	output, err := runYoutubeDL(args)
	if err != nil {
		return nil, err
	}

	lines := strings.Split(string(output), "\n")

	section := ""
	for _, line := range lines {
		if strings.Contains(strings.ToLower(line), "youtube") {
			continue
		}
		if strings.Contains(strings.ToLower(line), "available automatic captions") {
			section = "automatic captions"
			continue
		}
		if strings.Contains(strings.ToLower(line), "available subtitles") {
			section = "available subtitles"
			continue
		}
		if strings.Contains(strings.ToLower(line), "language formats") {
			continue
		}
		if section != "" {
			space := regexp.MustCompile(`\s+`)
			line := space.ReplaceAllString(line, " ")
			line = strings.ReplaceAll(line, ",", "")
			fields := strings.Split(string(line), " ")
			if len(fields) >= 2 {
				if section == "automatic captions" {
					subs.AutomaticCaptions[fields[0]] = fields[1:]

				}
				if section == "available subtitles" {
					subs.Subtitles[fields[0]] = fields[1:]
				}
			}
		}
	}

	return subs, nil
}

// findSub return nil when subtitle in the requested language is not available in any format
func findSub(id string, language string) (formats []string, err error) {
	subList, err := listSubs(id)
	if err != nil {
		return nil, err
	}
	if formats, ok := subList.Subtitles[language]; ok {
		return formats, nil
	}
	return nil, nil
}

//  findAutomaticCaption return nil when subtitle not available in any format
func findAutomaticCaption(id string, language string) (formats []string, err error) {
	subList, err := listSubs(id)
	if err != nil {
		return nil, err
	}
	if formats, ok := subList.AutomaticCaptions[language]; ok {
		return formats, nil
	}
	return nil, nil
}

// downloadSubtitle downloads the subtitle of the given youtube video in the given language and format and saves it as file to $file.$language.$format.
func downloadSubtitle(auto bool, id string, language string, format string, file string) error {
	subType := "--write-sub"
	if auto {
		subType = "--write-auto-sub"
	}

	args := []string{
		"--skip-download",
		subType,
		"--sub-lang",
		language,
		"--sub-format",
		format,
		"https://youtu.be/" + id,
		"-o",
		file,
	}

	output, err := runYoutubeDL(args)
	if err != nil {
		return err
	}
	if strings.Contains(strings.ToLower(string(output)), "subtitles not available") {
		return ErrYoutubeDL(fmt.Sprintf("subtitles language %s not available", language))
	}
	if strings.Contains(strings.ToLower(string(output)), "no subtitle format found") {
		return ErrYoutubeDL(fmt.Sprintf("subtitle format %s not found", format))

		//TODO delete the alternative subtitle file that it downloads automatically
	}

	return nil
}

// GetSubtitle downloads the subtitle of the video id in the passed language.
// Throws an error if subtitle not found.
// Downlaods the automatic caption when auto is true.
func GetSubtitle(auto bool, id string, language string) (sub annotation.Caption, err error) {

	var formats []string
	if auto == false {
		formats, err = findSub(id, language)
		if err != nil {
			return nil, err
		}

		if formats == nil {
			return nil, ErrYoutubeDLLangNotSupported(fmt.Sprintf("No subtitle with language %s for video %s", language, id))
		}
	} else {
		formats, err = findAutomaticCaption(id, language)
		if err != nil {
			return nil, err
		}

		if formats == nil {
			return nil, ErrYoutubeDLLangNotSupported(fmt.Sprintf("No automatic caption with language %s for video %s", language, id))
		}
	}

	log.WithFields(log.Fields{
		"videoID":  id,
		"language": language,
		"auto":     auto,
		"types":    formats,
	}).Tracef("Available subtitle types for this video")

	for _, preferredSubType := range SupportedSubs {

		for _, format := range formats {

			if preferredSubType.String() == format {

				filename := filepath.Join("/tmp", "sub-"+id)
				file := filename + "." + language + "." + format

				defer os.Remove(file)
				err := downloadSubtitle(auto, id, language, format, filename)
				if err != nil {
					log.Warn(err)
					continue
				}

				sub, err := parseSub(file, preferredSubType)
				if err != nil {
					log.Warn(err)
					continue
				}

				return sub, nil
			}
		}

		log.WithFields(log.Fields{"videoID": id,
			"language": language,
			"type":     preferredSubType.String(),
		}).Debugf("Preferred subtitle type not available for this video")

	}

	return nil, ErrYoutubeDL("No available subtitle is supported")
}

// GetVideoURL retrieves the videoplayback url of the video and audio stream
func GetVideoURL(id string) (video *url.URL, audio *url.URL, err error) {

	args := []string{
		"-g",
		"https://www.youtube.com/watch?v=" + id,
	}

	// run command
	cmd := exec.Command("youtube-dl", args...)
	out, err := cmd.Output()
	if err != nil {
		return nil, nil, err
	}
	links := strings.Split(strings.Trim(string(out), "\n\t "), "\n")
	if len(links) != 2 {
		return nil, nil, ErrYoutubeDL(fmt.Sprintf("Expected two links but got %v", len(links)))
	}
	video, err = url.Parse(links[0])
	if err != nil {
		return nil, nil, err
	}
	audio, err = url.Parse(links[1])
	if err != nil {
		return nil, nil, err
	}
	return video, audio, nil
}

// GetPlaylistVideos returns a list of youtube video id's that are in the given playlist.
// The response is limited to the first 50 videos.
func GetPlaylistVideos(playListID string, apiKey string) (members []string, err error) {
	client := &http.Client{
		Transport: &transport.APIKey{Key: apiKey},
	}

	service, err := youtube.NewService(context.Background(), option.WithHTTPClient(client))
	if err != nil {
		return nil, err
	}
	// Make the API call to YouTube.
	call := service.PlaylistItems.List("id,contentDetails").
		PlaylistId(playListID).
		MaxResults(50)

	response, err := call.Do()
	if err != nil {
		return nil, err
	}

	for _, item := range response.Items {
		members = append(members, item.ContentDetails.VideoId)
	}

	return members, nil
}

// LoadYouTubeVideos downloads the subtitles of the specified YouTube videos and creates annotations in the specified mongo db collection
// user supplied YouTube subtitles are preferred over auto-generated ones
func LoadYouTubeVideos(videoIDs []string, language string, collection *mongo.Collection) error {
	for _, v := range videoIDs {
		err := loadYouTubeAnnotationsVideo(v, language, collection)
		if err != nil {
			log.Warn(err)
			continue
		}
	}
	return nil
}

func loadYouTubeAnnotationsVideo(videoID string, language string, collection *mongo.Collection) error {
	var subs annotation.Caption
	err := errors.New("")

	// priorize manual subtitle before downloading automatic youtube caption (autogenerated)
	subs, err = GetSubtitle(false, videoID, language)
	if err != nil {
		_, ok := err.(ErrYoutubeDLLangNotSupported)
		if ok {
			log.WithFields(log.Fields{"videoID": videoID, "language": language}).Trace("No manual subtitle found")
			// download automatic caption
			subs, err = GetSubtitle(true, videoID, language)
			if err != nil {
				return err
			}

		} else {
			return err
		}
	}

	annotations := subs.Annotations("youtube://" + videoID)

	log.WithFields(log.Fields{"videoID": videoID, "language": language, "quantity": len(annotations)}).Trace("Subtitle items found")

	err = annotation.UploadToDB(annotations, collection)
	if err != nil {
		return err
	}

	return nil
}

// LoadYouTubeAnnotationsPlaylist downloads the subtitles of the specified YouTube playlist and loads them to the specified mongo db collection
// user supplied YouTube subtitles are preferred over auto-generated ones
func LoadYouTubeAnnotationsPlaylist(playlistID string, language string, collection *mongo.Collection) error {
	defer monitor.Elapsed()()

	members, err := GetPlaylistVideos(playlistID, os.Getenv("YOUTUBE_DATA_API_KEY"))
	if err != nil {
		return err
	}
	log.WithFields(log.Fields{"playlistID": playlistID, "language": language, "quantity": len(members)}).Trace("Videos found in playlist")

	for _, videoID := range members {
		err := loadYouTubeAnnotationsVideo(videoID, language, collection)
		if err != nil {
			log.Warning(err)
			continue
		}
	}

	return nil
}

// TimedText describes a parsed YouTube caption file type of .srv3 format (xml)
type TimedText struct {
	XMLName xml.Name `xml:"timedtext"`
	Format  string   `xml:"format,attr,omitempty"`
	Body    Body     `xml:"body"`
}

// Body contains all the paragraphs
type Body struct {
	Paragraphs []Paragraph `xml:"p"`
}

// Paragraph is one caption line (sequence of words=segments)
// Time [ms] is relative to the start of the caption (t=0)
// Duration ??
type Paragraph struct {
	Time     int       `xml:"t,attr"`
	Duration int       `xml:"d,attr"`
	Value    string    `xml:",chardata"`
	Segments []Segment `xml:"s,omitempty"`
}

// Segment is one words/term with the Time [ms] (relative to the Paragraph.Time) and Duration [ms]
type Segment struct {
	Time     int    `xml:"t,attr,omitempty"`
	Duration int    `xml:"ac,attr,"`
	Value    string `xml:",chardata"`
}

// SegmentsText returns the space delimited concatenation of all segment values
func (p *Paragraph) SegmentsText() string {
	var texts []string
	for _, segment := range p.Segments {
		texts = append(texts, strings.TrimSpace(segment.Value))
	}
	return util.WordsToString(texts)
}

// parseSRV decodes a youtubedl subtitle of type .srv3 into a TimedStruct
func parseSRV3(file []byte) (*TimedText, error) {
	var timedText TimedText

	err := xml.Unmarshal(file, &timedText)
	if err != nil {
		return nil, err
	}
	return &timedText, nil
}

// Annotations turns the TimedText into annotations with the src field set as given
// Method of the annotation.Caption IF
// Args default values
// 1. minWindowLength := 1
// 2. maxWindowLength := 5
func (s *TimedText) Annotations(src string) []*annotation.Annotation {

	return fromTimedText(s, src)
}

// fromTimedText creates one Annotation for each Paragraph
// TODO find smart way which segments to combine into one annotation (e.g. when t_pause > x || len(segments) > 10)
func fromTimedText(subtitle *TimedText, src string) []*annotation.Annotation {

	var annos []*annotation.Annotation

	count := 0
	for paragraphNum, paragraph := range subtitle.Body.Paragraphs {

		// if non-empty paragraph
		if len(paragraph.Segments) > 0 || strings.Join(strings.Fields(paragraph.Value), " ") != "" {
			count++

			// Create annotation for the paragraph
			a := annotation.Annotation{Src: src, Subtitle: annotation.Subtitle{Count: count}}

			// handle when current paragraph subtitle stays displayed even when next one appears
			var endTime int
			if paragraphNum < len(subtitle.Body.Paragraphs)-1 {
				if paragraph.Time+paragraph.Duration > subtitle.Body.Paragraphs[paragraphNum+1].Time {
					endTime = subtitle.Body.Paragraphs[paragraphNum+1].Time
				} else {
					endTime = paragraph.Time + paragraph.Duration
				}

			} else {
				endTime = paragraph.Time + paragraph.Duration
			}

			a.Subtitle.Start = paragraph.Time
			a.Subtitle.End = endTime

			// Set segments of annotation
			// Handle automatic capti0xc00046af50,ons with segments (word granularity timing)
			var segments []annotation.TimedText
			for segmentIdx, segment := range paragraph.Segments {

				// Compute the start time of the current window
				startTime := paragraph.Time + segment.Time

				var endTime int
				// if current segment not last one
				if segmentIdx < len(paragraph.Segments)-1 {
					endTime = paragraph.Time + paragraph.Segments[segmentIdx+1].Time
				} else {
					endTime = a.Subtitle.End
				}

				s := annotation.TimedText{
					Text:  strings.TrimSpace(segment.Value),
					Start: startTime,
					End:   endTime,
				}

				segments = append(segments, s)
			}

			// set paragraph value if non-empty and no segments present and
			if len(paragraph.Segments) == 0 {
				a.Subtitle.Text = paragraph.Value
			} else {
				a.Subtitle.Text = paragraph.SegmentsText()
			}

			a.Subtitle.Segments = segments

			annos = append(annos, &a)
		}
	}

	return annos
}

func parseSub(file string, format SubtitleFormat) (annotation.Caption, error) {

	if format == SubSRV3 {

		// opes file
		xmlFile, err := os.Open(file)
		if err != nil {
			return nil, err
		}
		defer xmlFile.Close()

		// read file
		byteValue, err := ioutil.ReadAll(xmlFile)
		if err != nil {
			return nil, err
		}

		// parse srv3
		timedText, err := parseSRV3(byteValue)
		if err != nil {
			return nil, err
		}

		return timedText, nil

	}

	sub, err := astisub.OpenFile(file)
	if err != nil {
		return nil, err
	}

	return &Subtitle{sub}, nil
}