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

//go:embed assets
var assetsFS embed.FS

var (
	ytdlCmd    = []string{"youtube-dl"}
	videosPath = "./vids"
	videosURL  = "/vids"
	cachePath  = "./cache"
	assetsPath = "./assets"
	saveAssets = false
	dbPath     = "./db.json"
	port       = 8080
	title      = "Youtube Downloader"
)

var d *db.Db

func parse() {
	ytcmd := ytdlCmd[0]
	flag.StringVar(&ytcmd, "ytdlcmd", ytcmd, "youtube-dl command seperated by spaces")
	flag.StringVar(&title, "title", title, "Title of the page")
	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.StringVar(&assetsPath, "assetspath", assetsPath, "Path where css files are saved and served")
	flag.BoolVar(&saveAssets, "saveassets", saveAssets, "Should the assets be saved in dir, so can be served by web server")
	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)
	if saveAssets {
		os.MkdirAll(assetsPath, 0755)
	}
}

func main() {
	log.Print("Youtube UI")
	log.SetFlags(log.Flags() | log.Lshortfile)
	parse()
	tmpl, err := template.New("page").Parse(page)
	if err != nil {
		log.Panic(err)
	}
	d, err = db.Load(dbPath)
	if err != nil {
		log.Panic(err)
	}
	defer d.Save()

	assets := []string{
		"assets/bootstrap.min.css",
		"assets/bootstrap.min.js",
	}

	if saveAssets {
		for _, assetFile := range assets {
			fname := strings.TrimPrefix(assetFile, "assets/")
			contents, err := assetsFS.ReadFile(assetFile)
			if err != nil {
				log.Panic(err)
			}
			destination := path.Join(assetsPath, fname)
			err = os.WriteFile(destination, contents, 0644)
			if err != nil {
				log.Printf("failed writing asset, assetFile: %s, destination: %s\n", assetFile, destination)
			}
		}
	}

	for _, assetFile := range assets {
		http.Handle("/"+assetFile, http.FileServer(http.FS(assetsFS)))
	}

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		switch r.Method {
		case http.MethodGet:
			td := struct {
				Title     string
				VidoesURL string
				Items     []db.Item
			}{title, videosURL, nil}
			d.Run(func(jd *db.Jdb) {
				//reverse order
				for _, item := range jd.Items {
					td.Items = append([]db.Item{item}, td.Items...)
				}
			})
			if err := tmpl.Execute(w, td); err != nil {
				log.Panic(err)
			}
		case http.MethodPost:
			yturl := r.PostFormValue("youtube_url")
			if yturl == "" {
				log.Printf("yturl empty, postform:%v\n", r.PostForm)
				w.WriteHeader(http.StatusBadRequest)
				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)
		default:
			w.WriteHeader(http.StatusForbidden)
		}
	})

	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.StatusNotFound)
			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
					}
				}
			}
		}
		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", 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", fname, source, destination, err)
	}
	return fname, nil
}