✨ Select audio stream

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
+}