Load config from file or cli.

1 files changed, 137 insertions, 17 deletions

diff --git a/main.go b/main.go
index 1cf56ad..792df07 100644
--- a/main.go
+++ b/main.go
@@ -1,30 +1,150 @@
 package main
 
 import (
+	"bilibili-livestream-archiver/common"
 	"bilibili-livestream-archiver/recording"
 	"context"
 	"fmt"
+	"github.com/akamensky/argparse"
+	"github.com/spf13/viper"
+	"os"
+	"os/signal"
+	"sync"
 )
 
-func main() {
-	task := recording.TaskConfig{
-		RoomId: 7777,
-		Transport: recording.TransportConfig{
-			SocketTimeoutSeconds: 10,
-			RetryIntervalSeconds: 5,
-			MaxRetryTimes:        5,
+var globalConfig *GlobalConfig
+
+func getTasks() (tasks []recording.TaskConfig) {
+	var err error
+	parser := argparse.NewParser(
+		"slbr",
+		"Record bilibili live streams",
+	)
+	defer func() {
+		if err != nil {
+			fmt.Printf("ERROR: %v.\n", err)
+			fmt.Print(parser.Usage(""))
+			os.Exit(0)
+		}
+	}()
+	configFilePtr := parser.String(
+		"c", "config",
+		&argparse.Options{
+			Required: false,
+			Help:     "Specify which configuration file to use. JSON, TOML and YAML are all supported.",
 		},
-		Download: recording.DownloadConfig{
-			SaveDirectory:    ".",
-			FileNameTemplate: "",
+	)
+	rooms := parser.IntList(
+		"s", "room",
+		&argparse.Options{
+			Required: false,
+			Help:     "The room id to record. Set this to run without config file",
 		},
+	)
+	saveToPtr := parser.String(
+		"o", "save-to",
+		&argparse.Options{
+			Required: false,
+			Help:     "Specify which configuration file to use",
+		},
+	)
+	err = parser.Parse(os.Args)
+	if err != nil {
+		return
+	}
+
+	fromCli := len(*rooms) > 0
+	fromFile := *configFilePtr != ""
+
+	if fromCli == fromFile {
+		err = fmt.Errorf("cannot specify room id argument and config file at the same time")
+		return
+	}
+
+	if !fromCli && !fromFile {
+		err = fmt.Errorf("no task specified")
+		return
+	}
+
+	if fromFile {
+		configFile := *configFilePtr
+		fmt.Printf("Config file: %v\n", configFile)
+		var file *os.File
+		file, err = os.Open(configFile)
+		if err != nil {
+			err = fmt.Errorf("cannot open config file \"%v\": %w", configFile, err)
+			return
+		}
+		err = viper.ReadConfig(file)
+		if err != nil {
+			err = fmt.Errorf("cannot read config file \"%v\": %w", configFile, err)
+			return
+		}
+		var gc GlobalConfig
+		err = viper.Unmarshal(&gc)
+		if err != nil {
+			err = fmt.Errorf("cannot parse config file \"%v\": %w", configFile, err)
+			return
+		}
+		globalConfig = &gc
+		return globalConfig.Tasks
 	}
+
+	// generate task list from cli
+	taskCount := len(*rooms)
+	tasks = make([]recording.TaskConfig, taskCount)
+	saveTo := common.Zeroable[string](*saveToPtr).OrElse(".")
+	for i := 0; i < taskCount; i++ {
+		tasks[i] = recording.TaskConfig{
+			RoomId:    common.RoomId((*rooms)[i]),
+			Transport: recording.DefaultTransportConfig(),
+			Download: recording.DownloadConfig{
+				SaveDirectory: saveTo,
+			},
+		}
+	}
+
+	return
+}
+
+func main() {
+	tasks := getTasks()
+
+	fmt.Println("Record tasks:")
+	for _, task := range tasks {
+		fmt.Println(task.String())
+	}
+	fmt.Println("========================")
+
 	chResult := make(chan recording.TaskResult)
-	go recording.RunTask(
-		context.Background(),
-		&task,
-		chResult,
-	)
-	result := <-chResult
-	fmt.Println(result.Error)
+	wg := sync.WaitGroup{}
+	ctx, cancelTasks := context.WithCancel(context.Background())
+	for _, task := range tasks {
+		wg.Add(1)
+		go recording.RunTask(
+			ctx,
+			&wg,
+			&task,
+			chResult,
+		)
+	}
+
+	chSigInt := make(chan os.Signal)
+	signal.Notify(chSigInt, os.Interrupt)
+	for {
+		select {
+		case <-chSigInt:
+			fmt.Println("Stopping...")
+			cancelTasks()
+		case result := <-chResult:
+			err := result.Error
+			if err != nil {
+				fmt.Printf("A task stopped with an error (room %v): %v\n",
+					result.Task.RoomId, result.Error)
+			} else {
+				fmt.Printf("Task stopped (room %v): %v\n",
+					result.Task.RoomId, result.Task.String())
+			}
+		}
+	}
 }