YTSFlix_Go/engine/movietorrent.go

package engine

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/anacrolix/torrent"
	"github.com/asticode/go-astisub"
	"github.com/odwrtw/yifysubs"
	"github.com/sirupsen/logrus"
	"gopkg.in/h2non/filetype.v1"
)

type MovieTorrent struct {
	FileName  string
	MIME      string
	Subtitles map[string]string
	torrent   *torrent.Torrent
	file      *torrent.File
	engine    *Engine
}

func NewMovieTorrent(file *torrent.File, t *torrent.Torrent) *MovieTorrent {
	fileName := filepath.Base(file.Path())
	ext := strings.SplitAfter(path.Ext(fileName), ".")[1]
	MIME := filetype.GetType(ext).MIME.Value
	return &MovieTorrent{
		torrent:  t,
		file:     file,
		FileName: fileName,
		MIME:     MIME,
	}
}

type Progress struct {
	Downloaded    uint64
	Size          uint64
	DownloadSpeed uint64
	Percentage    int
}

func (m *MovieTorrent) NewReader() torrent.Reader {
	reader := m.file.NewReader()
	reader.SetReadahead(m.file.Length() / 100)
	reader.SetResponsive()
	return reader
}

func (m *MovieTorrent) Progress() chan Progress {
	pc := make(chan Progress)
	var p int64 = 0
	size := m.torrent.Info().TotalLength()
	c := m.torrent.Closed()
	go func() {
		defer close(pc)
		for {
			select {
			case <-c:
				logrus.Debugf("%s stopped", m.FileName)
				return
			default:
				if p >= size {
					logrus.Debugf("%s completed", m.FileName)
					return
				}
				currentProgress := m.torrent.BytesCompleted()
				percent := int(float64(currentProgress) / float64(size) * 100)
				pc <- Progress{
					Downloaded:    uint64(currentProgress),
					Size:          uint64(size),
					Percentage:    percent,
					DownloadSpeed: uint64(currentProgress - p),
				}
				p = currentProgress
				time.Sleep(time.Second)

			}
		}
	}()
	return pc
}

func (m *MovieTorrent) Cancel() {
	logrus.Infof("Stopping torrent: %s", m.FileName)
	m.torrent.Drop()
}

func (m *MovieTorrent) DownloadSubtitles(imdbID string) error {
	logrus.Infof("Getting subtitles for: %s", m.FileName)
	//Create a client
	client := yifysubs.New("http://yifysubtitles.com")

	//Search subtitles
	subtitles, err := client.Search(imdbID)
	if err != nil {
		return err
	}
	m.Subtitles = map[string]string{}
	basePath := filepath.Join(m.engine.StoragePath, filepath.Dir(m.file.Path()))
	logrus.Debugf("Subtitles base path: %s", basePath)
	if _, err := os.Stat(basePath); os.IsNotExist(err) {
		os.Mkdir(basePath, os.ModePerm)
	}
	var wg sync.WaitGroup
	wg.Add(len(subtitles))
	mx := sync.RWMutex{}
	for _, subtitle := range subtitles {
		go func(sub *yifysubs.Subtitle) {
			defer wg.Done()
			name := fmt.Sprintf("%s.vtt", sub.Lang)
			ref := filepath.Join(filepath.Dir(m.file.Path()), name)
			p := filepath.Join(basePath, name)
			if _, err := os.Stat(p); err == nil {
				mx.Lock()
				m.Subtitles[name] = ref
				mx.Unlock()
				return
			}
			file, err := os.Create(p)
			if err != nil {
				logrus.Error(err)
				return
			}
			defer file.Close()
			defer sub.Close()
			s, _ := astisub.ReadFromSRT(sub)

			if err := s.WriteToWebVTT(file); err != nil {
				logrus.Debugf("error subtitle: %v", err)
				os.Remove(p)
				return
			}
			mx.Lock()
			m.Subtitles[name] = ref
			mx.Unlock()

		}(subtitle)
	}
	wg.Wait()
	return nil
}