From 07b2c00c1447bb90d1d58b1ffec9afffe42abfee Mon Sep 17 00:00:00 2001
From: Dan Jones <danjones@goodevilgenius.org>
Date: Tue, 29 Aug 2023 20:18:27 -0500
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Select=20audio=20stream?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cmd/root.go      |  2 ++
 config/config.go | 16 +++++++---
 ffmpeg/probe.go  | 82 ++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 95 insertions(+), 5 deletions(-)

diff --git a/cmd/root.go b/cmd/root.go
index 396bc72..c4de1ff 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -22,6 +22,8 @@ var rootCmd = &cobra.Command{
 		file := files.PickRandomFile()
 		out := ffmpeg.ProbeFile(file)
 		fmt.Printf("%+v\n", out)
+		stream := ffmpeg.SelectAudioStream(out)
+		fmt.Printf("Stream %d of %s: %+v\n", stream, out.Format.Path, out.Streams[stream])
 	},
 }
 
diff --git a/config/config.go b/config/config.go
index c3c76f5..3831883 100644
--- a/config/config.go
+++ b/config/config.go
@@ -10,11 +10,14 @@ import (
 )
 
 type Config struct {
-	Source           string   `toml:"source"`
-	SourceExtensions []string `toml:"source_ext"`
-	Recurse          bool     `toml:"recurse"`
-	SavePath         string   `toml:"save_path"`
-	MusicPath        string   `toml:"music_path"`
+	Source           string            `toml:"source"`
+	SourceExtensions []string          `toml:"source_ext"`
+	Recurse          bool              `toml:"recurse"`
+	SavePath         string            `toml:"save_path"`
+	MusicPath        string            `toml:"music_path"`
+	Codec            string            `toml:"codec"`
+	AllowedCodecs    []string          `toml:"allowed_codec"`
+	CodecExt         map[string]string `toml:"codec_ext"`
 }
 
 var config Config
@@ -26,6 +29,9 @@ func newConfig() Config {
 		Recurse:          true,
 		SavePath:         filepath.Join(xdg.UserDirs.Documents, "StripBeats"),
 		MusicPath:        xdg.UserDirs.Music,
+		Codec:            "opus",
+		AllowedCodecs:    []string{"aac", "opus"},
+		CodecExt:         map[string]string{"aac": "m4a", "opus": "opus"},
 	}
 }
 
diff --git a/ffmpeg/probe.go b/ffmpeg/probe.go
index b599de2..5d4ffa1 100644
--- a/ffmpeg/probe.go
+++ b/ffmpeg/probe.go
@@ -1,7 +1,9 @@
 package ffmpeg
 
 import (
+	"codeberg.org/danjones000/strip-beats/config"
 	"encoding/json"
+	"github.com/akrennmair/slice"
 	ffmpeg "github.com/u2takey/ffmpeg-go"
 )
 
@@ -10,6 +12,56 @@ type Probe struct {
 	Streams []Stream
 }
 
+func (p Probe) GetAudioStreams() []*Stream {
+	return slice.Map(slice.Filter(p.Streams, func(st Stream) bool {
+		return st.CodecType == "audio"
+	}), func(s Stream) *Stream {
+		return &s
+	})
+}
+
+func (p Probe) GetFirstAudio() *Stream {
+	st := p.GetAudioStreams()
+	if len(st) == 0 {
+		return nil
+	}
+	return st[0]
+}
+
+func (p Probe) GetAudioCount() int {
+	st := p.GetAudioStreams()
+	return len(st)
+}
+
+func (p Probe) GetFirstPreferredAudio() *Stream {
+	sts := slice.Filter(p.GetAudioStreams(), func(st *Stream) bool {
+		return st.isWantedCodec()
+	})
+
+	if len(sts) < 1 {
+		return nil
+	}
+
+	return sts[0]
+}
+
+func (p Probe) GetFirstAcceptableAudio() *Stream {
+	pref := p.GetFirstPreferredAudio()
+	if pref != nil {
+		return pref
+	}
+
+	sts := slice.Filter(p.GetAudioStreams(), func(st *Stream) bool {
+		return st.isAcceptableCodec()
+	})
+
+	if len(sts) < 1 {
+		return nil
+	}
+
+	return sts[0]
+}
+
 type Stream struct {
 	Index         int
 	CodecName     string `json:"codec_name"`
@@ -25,6 +77,18 @@ type Stream struct {
 	Tags          Tags
 }
 
+func (st Stream) isWantedCodec() bool {
+	wantedcodec := config.GetConfig().Codec
+	return st.CodecName == wantedcodec
+}
+
+func (st Stream) isAcceptableCodec() bool {
+	codecs := config.GetConfig().AllowedCodecs
+	return slice.ReduceWithInitialValue(codecs, false, func(acc bool, codec string) bool {
+		return acc || st.CodecName == codec
+	})
+}
+
 type Format struct {
 	Path           string  `json:"filename"`
 	Streams        int     `json:"nb_streams"`
@@ -66,3 +130,21 @@ func ProbeFile(path string) Probe {
 	}
 	return ret
 }
+
+func SelectAudioStream(pr Probe) int {
+	count := pr.GetAudioCount()
+	if count < 1 {
+		return -1
+	}
+
+	if count == 1 {
+		return pr.GetFirstAudio().Index
+	}
+
+	pref := pr.GetFirstAcceptableAudio()
+	if pref != nil {
+		return pref.Index
+	}
+
+	return pr.GetFirstAudio().Index
+}