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

package discover

import (
	"encoding/json"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"backend/modules/medias"

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

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

type MediaUuidMap struct {
	UUID string
	Name string
}

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

	mapFile := filepath.Join(to, "map.json")
	d.log.Infof("read in from: %s", mapFile)
	if _, err := os.Stat(mapFile); err != nil {
		if os.IsNotExist(err) {
			if err := os.WriteFile(mapFile, []byte("[]"), os.ModePerm); err != nil {
				return errors.Wrapf(err, "write file: %s", mapFile)
			}
		}
	}

	b, err := os.ReadFile(mapFile)
	if err != nil {
		return errors.Wrapf(err, "read file: %s", mapFile)
	}

	var store []MediaUuidMap
	err = json.Unmarshal(b, &store)
	if err != nil {
		return errors.Wrapf(err, "unmarshal json: %s", mapFile)
	}
	defer func() {
		d.runCleanup(to, store)
	}()

	storeMap := make(map[string]string)
	for _, item := range store {
		storeMap[item.Name] = item.UUID
	}

	dirs, err := d.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 _, ok := storeMap[dir]; ok {
			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)
			}
		}

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

		store = append(store, MediaUuidMap{
			UUID: uuid,
			Name: dir,
		})

		b, err := json.Marshal(store)
		if err != nil {
			return errors.Wrapf(err, "marshal json: %s", mapFile)
		}
		if err := os.WriteFile(mapFile, b, os.ModePerm); err != nil {
			return errors.Wrapf(err, "write file: %s", mapFile)
		}
		storeMap[dir] = uuid
		d.log.Infof("Store map: %s", storeMap)
	}

	return nil
}

func (d *DiscoverMedias) getSubDirs(root string) ([]string, error) {
	fd, err := os.Open(root)
	if err != nil {
		return nil, errors.Wrapf(err, "open root directory: %s", root)
	}
	defer fd.Close()

	entries, err := fd.Readdir(-1)
	if err != nil {
		return nil, errors.Wrapf(err, "read root directory: %s", root)
	}

	var paths []string
	for _, entry := range entries {
		if entry.IsDir() {
			paths = append(paths, entry.Name())
		}
	}

	return paths, 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 []MediaUuidMap) {
	uuids := lo.Map(store, func(item MediaUuidMap, _ int) string {
		return item.UUID
	})

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

	for _, dir := range dirs {
		if !lo.Contains(uuids, dir) {
			d.log.Infof("Remove dir: %s", dir)
			if err := os.RemoveAll(filepath.Join(to, dir)); err != nil {
				d.log.Errorf("Remove dir: %s", dir)
			}
		}
	}
}