ytui/main.go

package main

import (
	"bufio"
	_ "embed"
	"flag"
	"fmt"
	"html/template"
	"io/fs"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"

	"gitlab.com/balki/ytui/db"
	"gitlab.com/balki/ytui/pubsub"
	"golang.org/x/net/websocket"
)

//go:embed templates/index.html
var page string

var (
	ytdlCmd    = []string{"youtube-dl"}
	videosPath = "./vids"
	videosUrl  = "/vids"
	cachePath  = "./cache"
	dbPath     = "./db.json"
	port       = 8080
)

var d *db.Db

func parse() {
	ytcmd := ytdlCmd[0]
	flag.StringVar(&ytcmd, "ytdlcmd", ytcmd, "youtube-dl command seperated by spaces")
	flag.StringVar(&videosPath, "videospath", videosPath, "Path where videos are saved")
	flag.StringVar(&videosUrl, "videosurl", videosUrl, "Prefix of the url, i.e. https://domain.com/<this var>/<video filename>")
	flag.StringVar(&cachePath, "cachepath", cachePath, "Path where temporary download files are saved")
	flag.StringVar(&dbPath, "dbpath", dbPath, "Path where downloaded info is saved")
	flag.IntVar(&port, "port", port, "Port to listen on")

	flag.Parse()
	if ytcmd != ytdlCmd[0] {
		ytdlCmd = strings.Fields(ytcmd)
	}
	os.MkdirAll(cachePath, 0755)
	os.MkdirAll(videosPath, 0755)
	os.MkdirAll(path.Dir(dbPath), 0755)
}

func main() {
	log.Print("Youtube UI")
	log.SetFlags(log.Flags() | log.Lshortfile)
	parse()
	tmpl := template.New("page")
	tmpl = tmpl.Funcs(map[string]any{
		"vids_prefix": func() string { return videosUrl },
	})
	tmpl, err := tmpl.Parse(page)
	if err != nil {
		log.Panic(err)
	}
	d, err = db.Load(dbPath)
	if err != nil {
		log.Panic(err)
	}
	defer d.Save()
	http.Handle("/assets/", http.FileServer(http.Dir("./")))
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodGet {
			d.Run(func(d *db.Jdb) {
				d.Foo = "Hello"
				if err := tmpl.Execute(w, d); err != nil {
					log.Panic(err)
				}
			})
			return
		}
		yturl := r.PostFormValue("youtube_url")
		if yturl == "" {
			log.Printf("yturl empty, postform:%v\n", r.PostForm)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		id, isNew := d.Add(db.Item{
			Date:   time.Now().Format("2006-01-02 15:04"),
			URL:    yturl,
			Title:  "Loading",
			Status: db.NotStarted,
		})
		if isNew {
			go getTitle(id, yturl)
			go download(id, yturl)
		}
		http.Redirect(w, r, "/", http.StatusSeeOther)
	})
	http.HandleFunc("/ws/", func(w http.ResponseWriter, r *http.Request) {
		p := r.URL.Path
		idStr := strings.TrimPrefix(p, "/ws/")
		id, err := strconv.Atoi(idStr)
		if err != nil {
			log.Printf("Invalid id, %q, err:%v\n", idStr, err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		var sc <-chan string
		d.Transact(id, false, func(i *db.Item) {
			pt := i.Pt
			if pt != nil {
				sc = pt.Subscribe()
			}
		})
		var wh websocket.Handler = func(c *websocket.Conn) {
			defer c.Close()
			if sc != nil {
				for update := range sc {
					_, err := c.Write([]byte(update))
					if err != nil {
						log.Printf("err: %v\n", err)
						return
					}
				}
			}
			var status db.DownloadStatus
			var fname string
			d.Transact(id, false, func(i *db.Item) {
				status = i.Status
				fname = i.FileName
			})
			if status != db.Done {
				_, err := c.Write([]byte(status))
				if err != nil {
					log.Printf("err: %v\n", err)
				}
				return
			}
			alink := fmt.Sprintf(`<a target="_blank" href="%s/%s">Watch</a>`, videosUrl, fname)
			c.Write([]byte(alink))
		}
		wh.ServeHTTP(w, r)
	})
	http.HandleFunc("/title/", func(w http.ResponseWriter, r *http.Request) {
		p := r.URL.Path
		idStr := strings.TrimPrefix(p, "/title/")
		id, err := strconv.Atoi(idStr)
		if err != nil {
			log.Printf("Invalid id, %q, err:%v\n", idStr, err)
			w.WriteHeader(http.StatusInternalServerError)
			return
		}

		var tc <-chan struct{}
		d.Transact(id, false, func(i *db.Item) {
			tc = i.TitleChan
		})
		<-tc
		var title string
		d.Transact(id, false, func(i *db.Item) {
			title = i.Title
		})
		w.Write([]byte(title))
	})
	log.Panic(http.ListenAndServe(fmt.Sprintf(":%v", port), nil))

}

func getTitle(id int, yturl string) {
	tc := make(chan struct{})
	defer close(tc)
	d.Transact(id, false, func(i *db.Item) {
		i.TitleChan = tc
	})
	args := append(ytdlCmd, "--get-title", yturl)
	cmd := exec.Command(args[0], args[1:]...)
	var title string
	if op, err := cmd.Output(); err != nil {
		log.Printf("command failed, cmd: %v, err: %v", cmd, err)
		title = "ERROR"
	} else {
		title = string(op)
	}
	d.Transact(id, true, func(i *db.Item) {
		i.Title = title
	})
}

func download(id int, yturl string) {
	pt := pubsub.NewProgressTracker()
	d.Transact(id, true, func(i *db.Item) {
		i.Status = db.InProgress
		i.Pt = pt
	})
	pc, err := pt.Publish()
	defer close(pc)
	if err != nil {
		log.Panic(err)
	}

	//Log progress
	go func() {
		sc := pt.Subscribe()
		log.Println("Watching download progress for id: ", id)
		if sc != nil {
			for update := range sc {
				log.Println(update)
			}
		}
		log.Println("Done watching download progress for id: ", id)
	}()

	var status db.DownloadStatus
	var fname string
	if fname, err = downloadYt(id, yturl, pc); err != nil {
		log.Printf("err: %v\n", err)
		status = db.Error
	} else {
		status = db.Done
	}
	d.Transact(id, true, func(i *db.Item) {
		i.Status = status
		i.FileName = fname
	})
}

func downloadYt(id int, yturl string, pc chan<- string) (string, error) {
	pathTmpl := fmt.Sprintf("%s/video_%d.%%(ext)s", cachePath, id)
	// pathTmpl := fmt.Sprintf("%s/video_%d.mp4", cachePath, id)
	args := append(ytdlCmd, "--newline", "--output", pathTmpl, yturl)
	cmd := exec.Command(args[0], args[1:]...)
	rc, err := cmd.StdoutPipe()
	if err != nil {
		pc <- "Pre starting error"
		return "", err
	}
	defer rc.Close()
	pc <- "Starting download"
	if err := cmd.Start(); err != nil {
		pc <- "Start error"
		return "", err
	}
	br := bufio.NewReader(rc)
	for {
		line, _, err := br.ReadLine()
		if err != nil {
			break
		}
		pc <- string(line)
	}
	pc <- "Waiting to complete..."
	err = cmd.Wait()
	var fname string
	if err != nil {
		pc <- "Download Error"
		return "", fmt.Errorf("Download Error, err: %w", err)
	}
	pc <- "Download Done, renaming"
	matches, err := fs.Glob(os.DirFS(cachePath), fmt.Sprintf("video_%d.*", id))
	if err != nil {
		pc <- "Match Error"
		return "", fmt.Errorf("Glob match error, err: %w\n", err)
	}
	if len(matches) != 1 {
		pc <- "Multiple Match Error"
		return "", fmt.Errorf("Got multiple matches, count: %v", len(matches))
	}
	fname = matches[0]
	source := path.Join(cachePath, fname)
	destination := path.Join(videosPath, fname)
	if err := os.Rename(source, destination); err != nil {
		pc <- "Rename error"
		return "", fmt.Errorf("Rename error, fname: %q, source: %q, destination: %q, err: %w\n", fname, source, destination, err)
	}
	return fname, nil
}