YTSFlix_Go/handler/handler.go

package handler

import (
	"fmt"
	"git.adphi.net/Adphi/ytsflix/templates"
	"git.adphi.net/Adphi/ytsflix/ytsclient"
	"github.com/dustin/go-humanize"
	"github.com/gorilla/mux"
	"github.com/sirupsen/logrus"
	"gopkg.in/h2non/filetype.v1"
	"io"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"
)

var (
	_ io.Reader = (*os.File)(nil)
)

type Handler struct {
	engine      *Engine
	clients     map[string]int
	yts         *ytsclient.Client
	mutex       sync.Mutex
	storagePath string
}

type HomeData struct {
	Categories map[string]map[string]Movie
}

type Movie struct {
	Title    string
	Year     int
	Cover    string
	Link     string
	Synopsis string
	Genre    []string
	Youtube  string
	Casting  []Cast
}

type Cast struct {
	Name  string
	Image string
}

type MovieData struct {
	Title     string
	Link      string
	Subtitles map[string]string
}

func NewHandler(storagePath string) (*Handler, error) {
	h := Handler{storagePath: storagePath}
	var err error
	h.engine, err = NewEngine(storagePath)
	if err != nil {
		return nil, err
	}
	h.yts = ytsclient.NewClient("https://yts.am/api/v2")
	h.clients = map[string]int{}
	return &h, nil
}

func (h *Handler) Close() {
	h.engine.Close()
}

func (h *Handler) Home(writer http.ResponseWriter, request *http.Request) {
	hd := HomeData{Categories: map[string]map[string]Movie{}}
	var mu sync.Mutex
	var wg sync.WaitGroup
	wg.Add(len(ytsclient.Genres))
	for _, g := range ytsclient.Genres {
		go func(genre ytsclient.Genre) {
			defer wg.Done()
			res, err := h.yts.List(&ytsclient.ListParams{
				Genre:   genre,
				Limit:   50,
				SortBy:  ytsclient.Year,
				OrderBy: ytsclient.Desc,
				Quality: ytsclient.Quality1080p,
			})
			if err != nil {
				sendError(writer, err)
				return
			}
			ms := map[string]Movie{}
			for _, m := range res {
				ms[m.Title] = Movie{
					Link:  fmt.Sprintf("/movie/%d", m.ID),
					Cover: m.MediumCoverImage,
				}
			}
			if len(ms) >= 10 {
				mu.Lock()
				hd.Categories[string(genre)] = ms
				mu.Unlock()
			}
		}(g)
	}
	wg.Wait()
	t := templates.HomeTemplate()
	writer.Header().Set("Content-Type", "text/html; charset=utf-8")
	writer.WriteHeader(200)
	t.Execute(writer, hd)

	//fmt.Fprint(writer, html)
}

func (h *Handler) Movie(w http.ResponseWriter, r *http.Request) {
	id := mux.Vars(r)["id"]
	m, err := h.yts.Get(id, nil)
	if err != nil {
		sendError(w, err)
	}
	go func() { h.engine.Prepare(getMagnet(m)) }()
	t := templates.MovieTemplate()
	var cast []Cast
	for _, c := range m.Cast {
		cast = append(cast, Cast{Name: c.Name, Image: c.URLSmallImage})
	}
	md := Movie{
		Title:    m.Title,
		Year:     m.Year,
		Genre:    m.Genres,
		Cover:    m.MediumCoverImage,
		Link:     fmt.Sprintf("/torrent/%d", m.ID),
		Synopsis: m.DescriptionFull,
		Youtube:  fmt.Sprintf("https://www.youtube.com/embed/%s", m.YtTrailerCode),
		Casting:  cast,
	}
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	t.Execute(w, md)
}

func (h *Handler) Search(writer http.ResponseWriter, request *http.Request) {
	vars := mux.Vars(request)
	s := vars["search"]
	ms, err := h.yts.Search(s, &ytsclient.ListParams{Quality: ytsclient.Quality1080p})
	if err != nil {
		sendError(writer, err)
		return
	}
	html := "<ul>"
	for _, m := range ms {
		html += fmt.Sprintf("<li><a href=\"/movie/%d\">%s</a></li>", m.ID, m.Title)
	}
	html += "</ul>"
	writer.Header().Set("Content-Type", "text/html; charset=utf-8")
	writer.WriteHeader(200)
	fmt.Fprint(writer, html)
}

func (h *Handler) Torrents(writer http.ResponseWriter, request *http.Request) {
	id := mux.Vars(request)["id"]
	m, err := h.yts.Get(id, nil)
	if err != nil {
		sendError(writer, err)
		return
	}
	magnet := getMagnet(m)
	mt, err := h.engine.Download(magnet)
	if err != nil {
		sendError(writer, err)
		return
	}
	if err := mt.DownloadSubtitles(m.ImdbCode); err != nil {
		logrus.Error(err)
	}
	cp := mt.Progress()
	go func() {
		for p := range cp {
			hSize := humanize.Bytes(uint64(p.Size))
			downloadSpeed := humanize.Bytes(p.DownloadSpeed) + "/s"
			complete := humanize.Bytes(p.Downloaded)
			logrus.Debugf("%s: Progress: %s / %s  %s", mt.FileName, complete, hSize, downloadSpeed)
		}
	}()
	link := fmt.Sprintf("/watch/%s", mt.FileName)

	ss := map[string]string{}
	for s, v := range mt.Subtitles {
		ss[strings.Replace(s, ".vtt", "", -1)] = fmt.Sprintf("/subs/%s", v)
	}
	md := MovieData{Title: mt.FileName, Link: link, Subtitles: ss}

	writer.Header().Set("Content-Type", "text/html; charset=utf-8")
	writer.WriteHeader(200)
	t := templates.WatchTemplate()
	t.Execute(writer, md)
}

func (h *Handler) Serve(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	var mt *MovieTorrent
	var err error
	movie := vars["movie"]
	h.mutex.Lock()
	h.clients[movie]++
	h.mutex.Unlock()
	if mt, err = h.engine.Get(movie); err != nil {
		http.Error(w, "file not found", http.StatusBadRequest)
		return
	}
	reader := mt.NewReader()
	defer reader.Close()
	defer func() {
		h.mutex.Lock()
		h.clients[movie]--
		h.mutex.Unlock()
		time.Sleep(5 * time.Second)
		if h.clients[movie] == 0 {
			logrus.Infof("The client closed the connection. Cleaning up: %s.", movie)
			mt.Cancel()
		}
	}()
	w.Header().Set("Connection", "Keep-Alive")
	w.Header().Set("Content-Type", mt.MIME)
	http.ServeContent(w, r, mt.FileName, time.Now(), reader)
}

func (h *Handler) Sub(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	m := vars["movie"]
	s := vars["sub"]
	logrus.Debugf("requesting subs: movie: %s, sub: %s", vars["movie"], vars["sub"])
	f, err := os.Open(fmt.Sprintf("%s/%s/%s", h.storagePath, m, s))
	if err != nil {
		sendError(w, err)
		return
	}
	mime := filetype.GetType(".vtt").MIME.Value
	w.Header().Add("Context-Type", mime)
	http.ServeContent(w, r, vars["sub"], time.Now(), f)
}

func sendError(writer http.ResponseWriter, err error) {
	logrus.Error(err)
	http.Error(writer, err.Error(), http.StatusInternalServerError)
}

func getMagnet(m ytsclient.Movie) string {
	var to *ytsclient.Torrent
	for _, t := range m.Torrents {
		if t.Quality == string(ytsclient.Quality1080p) {
			to = &t
			break
		}
	}
	if to == nil {
		to = &m.Torrents[0]
	}
	magnet, _ := m.Magnet(*to)
	return magnet
}