You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

156 lines
3.3 KiB

package main
import (
"fmt"
"io"
"os"
"runtime"
"strings"
"github.com/charmbracelet/bubbles/spinner"
tea "github.com/charmbracelet/bubbletea"
"github.com/kkdai/youtube/v2"
)
var client youtube.Client
func dl(i int, id string) (int, string) {
video, err := client.GetVideo(id)
if err != nil {
return -1, id
}
idx := 0
format := "mp4"
for i,f := range video.Formats {
if f.Quality != "" && f.AudioQuality != "" {
format = strings.Split(format, ";")[0]
idx = i
break
}
}
stream, _, err := client.GetStream(video, &video.Formats[idx])
if err != nil {
return -1, id
}
file, err := os.Create(video.Title + "." + format)
if err != nil {
return -1, id
}
defer file.Close()
_, err = io.Copy(file, stream)
if err != nil {
return -1, id
}
return 1, id
}
func worker(jobs <-chan struct{int;string}, results chan<- struct{int;string}) {
for j := range jobs {
ret, id := dl(j.int, j.string)
results <- struct{int;string}{ret, id}
}
}
type job struct {
name string
state int
}
type model struct {
spin spinner.Model
jobs []job
idcount int
completed int
results chan struct{int;string}
}
func initModel() model {
spin := spinner.NewModel()
spin.Spinner = spinner.Line
jobs := make([]job, 0)
return model{spin, jobs, 0, 0, nil}
}
func (m model) Init() tea.Cmd {
go func() {
for m.completed < m.idcount {
res := <-m.results
for i,j := range m.jobs {
if j.name == res.string {
m.jobs[i].state = res.int
}
}
m.completed++
}
}()
return spinner.Tick
}
func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
var cmd tea.Cmd = nil
switch msg := msg.(type) {
case tea.KeyMsg:
switch msg.String() {
case "ctrl+c", "q":
return m, tea.Quit
}
default:
m.spin, cmd = m.spin.Update(msg)
}
return m, cmd
}
func (m model) View() string {
out := "Jobs:\n"
for _,job := range m.jobs {
if job.name == "" {continue}
statestr := m.spin.View()
if job.state == 1 {
statestr = "✔️"
} else if job.state == -1 {
statestr = "❌"
}
out += statestr + " " + job.name + "\n"
}
return out
}
func main() {
args := os.Args
vidids := make([]string, 1)
if len(args) > 1 {
vidids = args[1:]
} else {
// :trollhd:
vidids = append(vidids, "https://youtu.be/dQw4w9WgXcQ")
}
client = youtube.Client{}
idslen := len(vidids)
jobs := make(chan struct{int;string}, idslen)
results := make(chan struct{int;string}, idslen)
for i := 0; i < runtime.NumCPU(); i++ {
go worker(jobs, results)
}
initmodel := initModel()
initmodel.idcount = idslen
initmodel.results = results
for idx,id := range vidids {
initmodel.jobs = append(initmodel.jobs, job{id,0})
jobs <- struct{int;string}{idx,id}
}
close(jobs)
p := tea.NewProgram(initmodel)
if err := p.Start(); err != nil {
fmt.Printf("Error starting tea: %v", err)
os.Exit(1)
}
}