From 3d6cd45909493ae376af737935b65abee8e97b43 Mon Sep 17 00:00:00 2001
From: krau <71133316+krau@users.noreply.github.com>
Date: Thu, 25 Jun 2026 16:17:03 +0800
Subject: [PATCH] feat: add configuration options for video download settings

---
 config.example.toml             | 11 +++++++++
 config/viper.go                 |  4 ++++
 config/ytdlp.go                 | 13 +++++++++++
 core/tasks/ytdlp/execute.go     |  8 +++----
 core/tasks/ytdlp/format.go      | 40 +++++++++++++++++++++++++++++++++
 core/tasks/ytdlp/format_test.go | 23 +++++++++++++++++++
 6 files changed, 94 insertions(+), 5 deletions(-)
 create mode 100644 config/ytdlp.go
 create mode 100644 core/tasks/ytdlp/format.go
 create mode 100644 core/tasks/ytdlp/format_test.go

diff --git a/config.example.toml b/config.example.toml
index 11dd746..ad1bd18 100644
--- a/config.example.toml
+++ b/config.example.toml
@@ -33,6 +33,17 @@ secret = ""
 # 转存完成后删除 Aria2 下载的本地文件
 remove_after_transfer = true
 
+# yt-dlp 视频下载配置
+[ytdlp]
+# 默认下载的最高视频清晰度 (按高度限制), 如 1080, 720, 480; 0 表示不限制 (下载最佳画质)
+# 仅在使用 /ytdlp 命令且未手动指定任何参数时生效
+max_height = 1080
+# 直接指定 yt-dlp format 选择表达式, 留空则使用 max_height
+# 设置后优先级高于 max_height, 例如: "bv*[height<=720]+ba/b"
+format = ""
+# 下载后转封装的视频容器格式, 留空则不转封装. 默认 mp4
+recode = "mp4"
+
 # HTTP API 配置
 [api]
 # 启用 HTTP API
diff --git a/config/viper.go b/config/viper.go
index 43c7ffc..89cc831 100644
--- a/config/viper.go
+++ b/config/viper.go
@@ -35,6 +35,7 @@ type Config struct {
 	Storages []storage.StorageConfig `toml:"-" mapstructure:"-" json:"storages"`
 	Parser   parserConfig            `toml:"parser" mapstructure:"parser" json:"parser"`
 	Hook     hookConfig              `toml:"hook" mapstructure:"hook" json:"hook"`
+	Ytdlp    YtdlpConfig             `toml:"ytdlp" mapstructure:"ytdlp" json:"ytdlp"`
 }
 
 type aria2Config struct {
@@ -131,6 +132,9 @@ func Init(ctx context.Context, configFile ...string) error {
 		"api.host":   "0.0.0.0",
 		"api.port":   8080,
 		"api.token":  "",
+
+		// yt-dlp
+		"ytdlp.recode": "mp4",
 	}
 
 	for key, value := range defaultConfigs {
diff --git a/config/ytdlp.go b/config/ytdlp.go
new file mode 100644
index 0000000..079362a
--- /dev/null
+++ b/config/ytdlp.go
@@ -0,0 +1,13 @@
+package config
+
+type YtdlpConfig struct {
+	// MaxHeight limits the video resolution by height in pixels (e.g. 1080, 720).
+	// 0 means no limit (best available). Ignored when Format is set.
+	MaxHeight int `toml:"max_height" mapstructure:"max_height" json:"max_height"`
+	// Format is a raw yt-dlp format selector (-f). When set, it takes precedence
+	// over MaxHeight and gives the user full control.
+	Format string `toml:"format" mapstructure:"format" json:"format"`
+	// Recode is the target video container yt-dlp recodes into (e.g. mp4).
+	// Empty disables recoding.
+	Recode string `toml:"recode" mapstructure:"recode" json:"recode"`
+}
diff --git a/core/tasks/ytdlp/execute.go b/core/tasks/ytdlp/execute.go
index 20bb362..7a35d99 100644
--- a/core/tasks/ytdlp/execute.go
+++ b/core/tasks/ytdlp/execute.go
@@ -85,12 +85,10 @@ func (t *Task) downloadFiles(ctx context.Context, tempDir string) ([]string, err
 	cmd := ytdlp.New().
 		Output(filepath.Join(tempDir, "%(title)s.%(ext)s"))
 
-	// If no custom flags are provided, use default behavior
+	// Apply config-based format/quality defaults only when the user passes no
+	// custom flags. Any user flag means they take full control of yt-dlp.
 	if len(t.Flags) == 0 {
-		cmd = cmd.
-			FormatSort("res,ext:mp4:m4a").
-			RecodeVideo("mp4").
-			RestrictFilenames()
+		cmd = applyFormatConfig(cmd, config.C().Ytdlp)
 	}
 	// Note: If custom flags are provided, users have full control over format/quality
 	// The output path is always set above to ensure downloads go to the correct directory
diff --git a/core/tasks/ytdlp/format.go b/core/tasks/ytdlp/format.go
new file mode 100644
index 0000000..4aa8c88
--- /dev/null
+++ b/core/tasks/ytdlp/format.go
@@ -0,0 +1,40 @@
+package ytdlp
+
+import (
+	"strconv"
+
+	ytdlp "github.com/lrstanley/go-ytdlp"
+
+	"github.com/krau/SaveAny-Bot/config"
+)
+
+// buildFormatSelector translates a max height into a yt-dlp format selector.
+// It prefers merging the best video+audio within the height limit, then falls
+// back to a single muxed stream. An empty result means "no explicit selector".
+func buildFormatSelector(maxHeight int) string {
+	if maxHeight <= 0 {
+		return ""
+	}
+	h := strconv.Itoa(maxHeight)
+	return "bv*[height<=" + h + "]+ba/b[height<=" + h + "]/b"
+}
+
+// applyFormatConfig configures format/quality on the yt-dlp command according to
+// the ytdlp config. It is only meant to be called when the user did not supply
+// any custom flags, so config-driven defaults never conflict with user input.
+func applyFormatConfig(cmd *ytdlp.Command, cfg config.YtdlpConfig) *ytdlp.Command {
+	switch {
+	case cfg.Format != "":
+		cmd = cmd.Format(cfg.Format)
+	case cfg.MaxHeight > 0:
+		cmd = cmd.Format(buildFormatSelector(cfg.MaxHeight))
+	default:
+		// Preserve the original default: prefer highest resolution mp4/m4a.
+		cmd = cmd.FormatSort("res,ext:mp4:m4a")
+	}
+	if cfg.Recode != "" {
+		cmd = cmd.RecodeVideo(cfg.Recode)
+	}
+	cmd = cmd.RestrictFilenames()
+	return cmd
+}
diff --git a/core/tasks/ytdlp/format_test.go b/core/tasks/ytdlp/format_test.go
new file mode 100644
index 0000000..4be88d1
--- /dev/null
+++ b/core/tasks/ytdlp/format_test.go
@@ -0,0 +1,23 @@
+package ytdlp
+
+import "testing"
+
+func TestBuildFormatSelector(t *testing.T) {
+	tests := []struct {
+		name      string
+		maxHeight int
+		want      string
+	}{
+		{"no limit", 0, ""},
+		{"negative", -1, ""},
+		{"1080p", 1080, "bv*[height<=1080]+ba/b[height<=1080]/b"},
+		{"720p", 720, "bv*[height<=720]+ba/b[height<=720]/b"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := buildFormatSelector(tt.maxHeight); got != tt.want {
+				t.Errorf("buildFormatSelector(%d) = %q, want %q", tt.maxHeight, got, tt.want)
+			}
+		})
+	}
+}