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

package discover

import (
	"bytes"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"

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

	"github.com/google/uuid"
	"github.com/pkg/errors"
	log "github.com/sirupsen/logrus"
)

// @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")
	}

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

	dirs, err := path.GetSubDirs(from)
	if err != nil {
		return errors.Wrapf(err, "get sub dirs: %s", from)
	}

	for _, dir := range dirs {
		d.log.Infof("Discover medias in: %s", dir)

		if exists := store.Exists(dir); exists {
			d.log.Infof("Skip dir: %s", dir)
			continue
		}

		dirPath := filepath.Join(from, dir)

		uuid := uuid.New().String()
		to := filepath.Join(to, uuid)

		videos := d.globVideo(dirPath)
		if len(videos) > 0 {
			video := videos[0]
			d.log.Infof("video: %s", video)
			if err := d.ffmpegVideoToM3U8(video, to); err != nil {
				return errors.Wrapf(err, "ffmpeg video to m3u8: %s", video)
			}

			d.log.Infof("snapshot", video)
			posterPath := filepath.Join(to, "poster.jpg")
			if err := d.ffmpegVideoToPoster(video, posterPath, 5); err != nil {
				return errors.Wrapf(err, "ffmpeg snapshot : %s", video)
			}
		}

		audios := d.globAudio(dirPath)
		if len(audios) > 0 {
			audio := audios[0]
			d.log.Infof("audio: %s", audio)
			if err := d.ffmpegAudioToM3U8(audio, to); err != nil {
				return errors.Wrapf(err, "ffmpeg audio to m3u8: %s", audio)
			}
		}

		price, err := d.getPrice(dirPath)
		if err != nil {
			return errors.Wrapf(err, "get price: %s", dirPath)
		}

		store.Add(uuid, dir, price)
		if err := store.Save(to); err != nil {
			return errors.Wrapf(err, "store save: %s", to)
		}
	}

	return nil
}

func (d *DiscoverMedias) globVideo(root string) []string {
	files := []string{}

	exts := []string{"*.mp4", "*.MP4", "*.mov", "*.MOV"}

	for _, ext := range exts {
		fs, _ := filepath.Glob(filepath.Join(root, ext))
		files = append(files, fs...)
	}

	return files
}

func (d *DiscoverMedias) globAudio(root string) []string {
	files := []string{}

	exts := []string{"*.mp3", "*.MP3", "*.wav", "*.WAV"}

	for _, ext := range exts {
		fs, _ := filepath.Glob(filepath.Join(root, ext))
		files = append(files, fs...)
	}

	return files
}

// 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",
		"-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 store.Exists(dir) {
			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, seconds int) error {
	// ffmpeg -i input_video.mp4 -ss N -vframes 1 -vf "scale=width:height" output_image.jpg
	args := []string{
		"-i", video,
		"-ss", fmt.Sprintf("00:%02d:00", seconds),
		"-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
}