strip-beats/media/probe.go

package media

import (
	"encoding/json"

	"codeberg.org/danjones000/strip-beats/config"
	"codeberg.org/danjones000/strip-beats/utils"
	"dario.cat/mergo"
	"github.com/akrennmair/slice"
	ffmpeg "github.com/u2takey/ffmpeg-go"
)

type Probe struct {
	Format  Format
	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]
}

func (pr Probe) ShortPath() string {
	return utils.GetShortPath(pr.Format.Path)
}

func (pr Probe) WantedAudioStream() *Stream {
	count := pr.GetAudioCount()
	if count < 1 {
		return nil
	}

	if count == 1 {
		return pr.GetFirstAudio()
	}

	pref := pr.GetFirstAcceptableAudio()
	if pref != nil {
		return pref
	}

	return pr.GetFirstAudio()
}

type Stream struct {
	Index         int
	CodecName     string `json:"codec_name"`
	CodecLongName string `json:"codec_long_name"`
	Profile       string
	CodecType     string `json:"codec_type"`
	CodecTag      string `json:"codec_tag_string"`
	Width         int
	Height        int
	StartTime     float64 `json:"start_time,string"`
	Duration      float64 `json:",string"`
	DurationTs    int     `json:"duration_ts"`
	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"`
	FormatName     string  `json:"format_name"`
	FormatLongName string  `json:"format_long_name"`
	StartTime      float64 `json:"start_time,string"`
	Duration       float64 `json:",string"`
	Size           int     `json:",string"`
	BitRate        int     `json:"bit_rate,string"`
	Score          int     `json:"probe_score"`
	Tags           Tags
}

type Tags struct {
	MajorBrand  string `json:"major_brand"`
	Title       string
	Artist      string
	AlbumArtist string `json:"album_artist"`
	Album       string
	Date        string
	Encoder     string
	Comment     string
	Description string
	Composer    string
	Genre       string
	Disc        string
	Track       string
}

func (pr Probe) FullTags() Tags {
	t := pr.Format.Tags
	s := pr.WantedAudioStream()
	if s != nil {
		mergo.Merge(&t, s.Tags)
	}
	for _, st := range pr.Streams {
		mergo.Merge(&t, st.Tags)
	}
	return t
}

func (pr Probe) Title() string {
	return pr.FullTags().Title
}

func ProbeFile(path string) Probe {
	out, err := ffmpeg.Probe(path)
	if err != nil {
		panic(err)
	}
	ret := Probe{}
	err = json.Unmarshal([]byte(out), &ret)
	if err != nil {
		panic(err)
	}
	return ret
}