mp-qvyun/backend/modules/commands/discover/discover_medias.go

package discover

import (
	"bytes"
	"crypto/md5"
	"fmt"
	"io"
	"math"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

	"backend/common/media_store"
	"backend/modules/medias"
	"backend/pkg/path"

	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cast"
)

// @provider
type DiscoverMedias struct {
	mediasSvc *medias.Service
	log       *log.Entry `inject:"false"`
}

// Prepare
func (d *DiscoverMedias) Prepare() error {
	d.log = log.WithField("module", "DiscoverMedias")
	return nil
}

func (d *DiscoverMedias) RunE(from, to string) error {
	d.log.Infof("Discover medias from: %s to: %s", from, to)
	if from == "" || to == "" {
		return errors.New("from or to is empty")
	}

	videos, err := d.globVideos(from)
	if err != nil {
		return errors.Wrapf(err, "glob videos: %s", from)
	}

	if err := d.ensureDirectory(to); err != nil {
		return errors.Wrapf(err, "ensure directory: %s", to)
	}

	store, err := media_store.NewStore(to)
	if err != nil {
		return errors.Wrapf(err, "new store: %s", to)
	}

	defer func() {
		d.runCleanup(to)
	}()

	for _, video := range videos {
		md5, err := d.getFileMD5(video)
		if err != nil {
			return errors.Wrapf(err, "get file md5: %s", video)
		}
		name := filepath.Base(video)[0:strings.LastIndex(filepath.Base(video), ".")]

		if info, ok := store.HashExists(md5); ok {
			info.Name = name
			store.Update(info)
			if err := store.Save(to); err != nil {
				return errors.Wrapf(err, "save store: %s", to)
			}
			continue
		}

		info, err := d.processVideo(video, filepath.Join(to, md5))
		if err != nil {
			return errors.Wrapf(err, "process video: %s", video)
		}
		info.Hash = md5
		info.Name = name

		store = store.Append(info)
		d.log.Infof("store: %+v", store)
		if err := store.Save(to); err != nil {
			return errors.Wrapf(err, "save store: %s", to)
		}
	}

	return nil
}

func (d *DiscoverMedias) processVideo(video string, to string) (media_store.VideoInfo, error) {
	var info media_store.VideoInfo

	if err := d.ensureDirectory(to); err != nil {
		return info, errors.Wrapf(err, "ensure directory: %s", to)
	}

	// extract poster
	posterFile := filepath.Join(to, "poster.jpg")
	if err := d.ffmpegVideoToPoster(video, posterFile); err != nil {
		return info, errors.Wrapf(err, "ffmpeg video to poster: %s", video)
	}

	// extract audio from video
	audioFile := filepath.Join(to, "audio.mp3")
	if err := d.extractAudio(video, audioFile); err != nil {
		return info, errors.Wrapf(err, "extract audio %s from %s", audioFile, video)
	}
	defer os.Remove(audioFile)

	// ffmpeg video to m3u8
	if err := d.ffmpegVideoToM3U8(video, to); err != nil {
		return info, errors.Wrapf(err, "ffmpeg video to m3u8: %s", video)
	}

	// ffmpeg audio to m3u8
	if err := d.ffmpegAudioToM3U8(audioFile, to); err != nil {
		return info, errors.Wrapf(err, "ffmpeg audio to m3u8: %s", audioFile)
	}

	// get media duration
	duration, err := d.getMediaDuration(video)
	if err != nil {
		return info, errors.Wrapf(err, "get media duration: %s", video)
	}
	info.Duration = duration

	return info, nil
}

func (d *DiscoverMedias) getFileMD5(filePath string) (string, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return "", errors.Wrapf(err, "open file: %s", filePath)
	}
	defer file.Close()

	hash := md5.New()
	if _, err := io.Copy(hash, file); err != nil {
		return "", errors.Wrapf(err, "copy file to hash: %s", filePath)
	}

	return fmt.Sprintf("%x", hash.Sum(nil)), nil
}

// extractAudio extract audio from video
func (d *DiscoverMedias) extractAudio(video string, output string) error {
	args := []string{
		"-i", video,
		"-vn",
		"-acodec", "libmp3lame",
		"-ar", "44100",
		"-b:a", "128k",
		"-q:a", "2",
		output,
	}

	d.log.Infof("cmd: ffmpeg %s", strings.Join(args, " "))
	logs, err := exec.Command("ffmpeg", args...).CombinedOutput()
	if err != nil {
		return errors.Wrapf(err, "extract audio: %s", video)
	}
	d.log.Infof("extract audio: %s", logs)

	return nil
}

// ensureDirectory ensure the directory exists
func (d *DiscoverMedias) ensureDirectory(dir string) error {
	st, err := os.Stat(dir)
	if os.IsNotExist(err) {
		if err := os.MkdirAll(dir, os.ModePerm); err != nil {
			return errors.Wrapf(err, "mkdir: %s", dir)
		}
		return nil
	}

	if !st.IsDir() {
		os.RemoveAll(dir)
		if err := os.MkdirAll(dir, os.ModePerm); err != nil {
			return errors.Wrapf(err, "mkdir: %s", dir)
		}
	}

	return nil
}

func (d *DiscoverMedias) ffmpegVideoToM3U8(input string, output string) error {
	output = filepath.Join(output, "video")
	if err := d.ensureDirectory(output); err != nil {
		return errors.Wrapf(err, "ensure directory: %s", output)
	}

	args := []string{
		"-i", input,
		"-c:v", "libx264",
		"-c:a", "aac",
		"-strict",
		"-2",
		"-vf", "scale=-720:576",
		"-f", "hls",
		"-hls_time", "10",
		"-hls_list_size", "0",
		"-hls_segment_filename", filepath.Join(output, "%d.ts"),
		filepath.Join(output, "index.m3u8"),
	}

	log.Infof("cmd: ffmpeg %s", strings.Join(args, " "))
	logs, err := exec.Command("ffmpeg", args...).CombinedOutput()
	if err != nil {
		return errors.Wrapf(err, "ffmpeg video to m3u8: %s", input)
	}

	d.log.Infof("ffmpeg video to m3u8: %s", logs)
	return nil
}

func (d *DiscoverMedias) ffmpegAudioToM3U8(input string, output string) error {
	output = filepath.Join(output, "audio")
	if err := d.ensureDirectory(output); err != nil {
		return errors.Wrapf(err, "ensure directory: %s", output)
	}

	args := []string{
		"-i", input,
		"-c:a", "aac",
		"-strict",
		"-2",
		"-f", "hls",
		"-hls_time", "10",
		"-hls_list_size", "0",
		"-hls_segment_filename", filepath.Join(output, "%d.ts"),
		filepath.Join(output, "index.m3u8"),
	}

	log.Infof("cmd: ffmpeg %s", strings.Join(args, " "))
	logs, err := exec.Command("ffmpeg", args...).CombinedOutput()
	if err != nil {
		return errors.Wrapf(err, "ffmpeg audio to m3u8: %s", input)
	}

	d.log.Infof("ffmpeg audio to m3u8: %s", logs)
	return nil
}

func (d *DiscoverMedias) runCleanup(to string) {
	store, err := media_store.NewStore(to)
	if err != nil {
		d.log.Errorf("new store: %s", to)
		return
	}

	dirs, err := path.GetSubDirs(to)
	if err != nil {
		d.log.Errorf("get sub dirs: %s", to)
		return
	}

	for _, dir := range dirs {
		if _, ok := store.HashExists(dir); ok {
			continue
		}

		d.log.Infof("Remove dir: %s", dir)
		if err := os.RemoveAll(filepath.Join(to, dir)); err != nil {
			d.log.Errorf("Remove dir: %s", dir)
		}
	}
}

// getSnapshot get the snapshot of target seconds of a video
func (d *DiscoverMedias) ffmpegVideoToPoster(video string, output string) error {
	// ffmpeg -i input_video.mp4 -ss N -vframes 1 -vf "scale=width:height" output_image.jpg
	args := []string{
		"-i", video,
		"-ss", "00:00:01",
		"-vframes", "1",
		"-vf", "scale=640:360",
		output,
	}

	d.log.Infof("cmd: ffmpeg %s", strings.Join(args, " "))
	logs, err := exec.Command("ffmpeg", args...).CombinedOutput()
	if err != nil {
		return errors.Wrapf(err, "get snapshot: %s", video)
	}

	d.log.Infof("get snapshot: %s", logs)
	return nil
}

func (d *DiscoverMedias) getPrice(dir string) (uint, error) {
	price, err := os.ReadFile(filepath.Join(dir, "price.txt"))
	if err != nil {
		return 0, errors.Wrapf(err, "read price: %s", dir)
	}

	if len(price) == 0 {
		return 0, fmt.Errorf("%s, price is empty", dir)
	}

	price = bytes.TrimSpace(price)

	p, err := strconv.Atoi(string(price))
	if err != nil {
		return 0, fmt.Errorf(dir, ", price is not a number")
	}

	return uint(p), nil
}

// getMediaDuration get the duration of a media file
func (d *DiscoverMedias) getMediaDuration(file string) (int64, error) {
	// ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 input_video.mp4
	args := []string{
		"-v", "error",
		"-show_entries", "format=duration",
		"-of", "default=noprint_wrappers=1:nokey=1",
		file,
	}

	log.Infof("cmd: ffprobe %s", strings.Join(args, " "))
	out, err := exec.Command("ffprobe", args...).CombinedOutput()
	if err != nil {
		return 0, errors.Wrapf(err, "get media duration: %s", file)
	}

	duration, err := cast.ToFloat64E(strings.TrimSpace(string(out)))
	if err != nil {
		return 0, errors.Wrapf(err, "get media duration: %s", file)
	}

	return int64(math.Floor(duration)), nil
}

// getMedias get the medias in the directory
func (d *DiscoverMedias) globVideos(dir string) ([]string, error) {
	// glob *.mp4 in dir
	return filepath.Glob(filepath.Join(dir, "*.mp4"))
}