quyun/backend/app/service/commands/compress.go

package commands

import (
	"bytes"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"quyun/app/service"

	"go.ipao.vip/atom"
	"go.ipao.vip/atom/container"

	"github.com/pkg/errors"
	"github.com/samber/lo"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"go.uber.org/dig"
)

func defaultProviders() container.Providers {
	return service.Default(container.Providers{}...)
}

func Command() atom.Option {
	return atom.Command(
		atom.Name("compress"),
		atom.Short("compress videos"),
		atom.RunE(Serve),
		atom.Providers(
			defaultProviders().With(),
		),
	)
}

type Service struct {
	dig.In
}

var compressed []compressedDB

type compressedDB struct {
	Name string `json:"name"`
	Md5  string `json:"md5"`
}

func Serve(cmd *cobra.Command, args []string) error {
	// log.SetFormatter(&log.JSONFormatter{})
	if len(args) == 0 {
		return errors.New("usage: compress path")
	}
	log.Infof("start compressing %s", args[0])

	failTimes := 0
	for {
		log.Infof("start compressing %s", args[0])

		if err := commandCompress(args[0]); err != nil {
			log.WithError(err).Errorf("compressing %s failed", args[0])
			failTimes++

			if failTimes >= 3 {
				return errors.New("failed to compress after 3 attempts")
			}
			continue
		}
		log.Infof("compressing %s done", args[0])

		time.Sleep(5 * time.Minute)
	}
}

func commandCompress(dir string) error {
	dstPath := filepath.Join(dir, "compressed")
	// check if the directory exists
	if _, err := os.Stat(dstPath); os.IsNotExist(err) {
		if err := os.MkdirAll(dstPath, 0o755); err != nil {
			return errors.Wrapf(err, "failed to create directory: %s", dstPath)
		}
	}
	if err := readDB(dir); err != nil {
		return errors.Wrap(err, "failed to read db")
	}

	dirFiles, err := os.ReadDir(dir)
	if err != nil {
		return errors.Wrapf(err, "failed to read directory: %s", dir)
	}
	//  loop through the files
	for _, file := range dirFiles {
		//  check if the file is a video
		if file.IsDir() {
			continue
		}

		log.Infof("process file %s", file.Name())

		//  get the file name
		fileName := file.Name()
		filePath := filepath.Join(dir, fileName)
		dstPath := filepath.Join(dstPath, fileName)

		fileMd5, err := calculateMD5(filePath)
		if err != nil {
			return err
		}

		if _, ok := lo.Find(compressed, func(c compressedDB) bool { return c.Md5 == fileMd5 }); ok {
			log.Infof("file %s already compressed", fileName)
			continue
		}

		//  get the file path
		if err := doCompress(filePath, dstPath); err != nil {
			log.WithError(err).Errorf("failed to compress file: %s", fileName)
			continue
		}

		compressed = append(compressed, compressedDB{Name: fileName, Md5: fileMd5})
		writeDB(dir)
		log.Infof("file %s compressed", fileName)
	}

	return nil
}

func readDB(path string) error {
	dbFile := filepath.Join(path, ".compressed.json")
	// check if the file exists
	if _, err := os.Stat(dbFile); os.IsNotExist(err) {
		// create the file
		file, err := os.Create(dbFile)
		if err != nil {
			return errors.Wrapf(err, "failed to create file: %s", dbFile)
		}
		defer file.Close()
		// write the file
		_, err = file.WriteString("[]")
		if err != nil {
			return errors.Wrapf(err, "failed to write file: %s", dbFile)
		}
	}

	// read the db
	file, err := os.ReadFile(dbFile)
	if err != nil {
		return errors.Wrapf(err, "failed to read file: %s", dbFile)
	}
	err = json.Unmarshal(file, &compressed)
	if err != nil {
		return errors.Wrapf(err, "failed to unmarshal file: %s", dbFile)
	}
	return nil
}

func writeDB(path string) error {
	dbFile := filepath.Join(path, ".compressed.json")

	// write the db
	data, err := json.Marshal(compressed)
	if err != nil {
		log.Errorf("failed to marshal file: %s", dbFile)
	}
	err = os.WriteFile(dbFile, data, 0o644)
	if err != nil {
		log.Errorf("failed to write file: %s", dbFile)
		return err
	}
	return nil
}

// calculateMD5 calculates the md5 of a file
func calculateMD5(filePath string) (string, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return "", errors.Wrapf(err, "failed to open file: %s", filePath)
	}
	defer file.Close()

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

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

// doCompress compresses the video file
func doCompress(srcPath, dstPath string) error {
	// get file ext
	ext := strings.ToLower(filepath.Ext(srcPath))
	if ext == "" {
		return errors.New("file has no extension")
	}

	// check if the file is a .mp3
	if strings.EqualFold(ext, ".mp3") {
		return compressAudio(srcPath, dstPath)
	}

	// check if the file is a video
	videoExts := []string{".mp4", ".avi", ".mkv"}
	if !lo.Contains(videoExts, ext) {
		return errors.New("file is not a video")
	}

	return compressVideo(srcPath, dstPath)
}

// copyFile copies the file from src to dst
func copyFile(src, dst string) error {
	srcFile, err := os.Open(src)
	if err != nil {
		return errors.Wrapf(err, "failed to open file: %s", src)
	}
	defer srcFile.Close()

	dstFile, err := os.Create(dst)
	if err != nil {
		return errors.Wrapf(err, "failed to create file: %s", dst)
	}
	defer dstFile.Close()

	if _, err := io.Copy(dstFile, srcFile); err != nil {
		return errors.Wrapf(err, "failed to copy file: %s", src)
	}

	return nil
}

// compress audio
func compressAudio(srcPath, dstPath string) error {
	log.Infof("compress audio from %s to %s", srcPath, dstPath)
	if err := copyFile(srcPath, dstPath); err != nil {
		return errors.Wrapf(err, "failed to copy file: %s", srcPath)
	}

	return nil
}

// compress video
func compressVideo(srcPath, dstPath string) error {
	// if file size < 50M then copy
	fileInfo, err := os.Stat(srcPath)
	if err != nil {
		return errors.Wrapf(err, "failed to get file info: %s", srcPath)
	}
	if fileInfo.Size() < 50*1024*1024 { // 50 MB
		log.Infof("file size < 50M, copy file from %s to %s", srcPath, dstPath)
		return copyFile(srcPath, dstPath)
	}

	width, height, err := getVideoSize(srcPath)
	if err != nil {
		return errors.Wrapf(err, "failed to get video size: %s", srcPath)
	}
	log.Infof("video size: %dx%d", width, height)

	// cmd := exec.Command("ffmpeg", "-i", srcPath, "-c:v", "libx264", "-crf", "23", "-r", "25", "-c:a", "copy", fmt.Sprintf("'%s'", dstPath))
	cmd := exec.Command("sh", "-c", fmt.Sprintf("ffmpeg -i '%s' -y -c:v libx264 -crf 23 -r 25 -c:a copy '%s'", srcPath, dstPath))
	if width > 1920 || height > 1080 {
		cmd = exec.Command("sh", "-c", fmt.Sprintf("ffmpeg -i '%s' -y -c:v libx264 -crf 23 -r 25 -c:a copy -vf scale=1920:1080 '%s'", srcPath, dstPath))
	}
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	log.Infof("executing command: %s", cmd.String())

	if err := cmd.Run(); err != nil {
		log.Errorf("stderr: %s", stderr.String())
		return errors.Wrapf(err, "failed to execute compress command: %s", cmd.String())
	}

	return nil
}

func getVideoSize(filePath string) (int, int, error) {
	var stderr bytes.Buffer
	cmd := exec.Command("sh", "-c", fmt.Sprintf("ffprobe -v error -select_streams v:0 -show_entries stream=height,width -of default=noprint_wrappers=1:nokey=1 '%s'", filePath))
	cmd.Stderr = &stderr
	out, err := cmd.Output()
	if err != nil {
		log.Errorf("stderr: %s", stderr.String())
		return 0, 0, errors.Wrapf(err, "failed to execute command: %s", cmd.String())
	}
	lines := strings.Split(string(out), "\n")

	if len(lines) < 2 {
		return 0, 0, errors.New("failed to get video size")
	}

	// parse the height and width
	width, err := strconv.Atoi(strings.TrimSpace(lines[0]))
	if err != nil {
		return 0, 0, errors.Wrap(err, "failed to parse height")
	}

	height, err := strconv.Atoi(strings.TrimSpace(lines[1]))
	if err != nil {
		return 0, 0, errors.Wrap(err, "failed to parse width")
	}
	return width, height, nil
}