vibedit/ffmpeg.go

package main

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
)

var exeSuffix = func() string {
	if runtime.GOOS == "windows" {
		return ".exe"
	}
	return ""
}()

// Codec describes an encoder the export dialog can offer.
type Codec struct {
	Name   string `json:"name"`   // ffmpeg encoder name
	Label  string `json:"label"`  // shown to the user
	Kind   string `json:"kind"`   // video | audio
	Family string `json:"family"` // h264, hevc, av1, vp9, aac, opus, mp3, flac
	HW     bool   `json:"hw,omitempty"`
	DefQ   int    `json:"defQ,omitempty"` // default quality value (CRF-like, lower = better)
	MaxQ   int    `json:"maxQ,omitempty"`
}

var codecCatalog = []Codec{
	{Name: "libx264", Label: "H.264 (CPU, most compatible)", Kind: "video", Family: "h264", DefQ: 20, MaxQ: 51},
	{Name: "h264_nvenc", Label: "H.264 (NVIDIA GPU)", Kind: "video", Family: "h264", HW: true, DefQ: 23, MaxQ: 51},
	{Name: "h264_qsv", Label: "H.264 (Intel GPU)", Kind: "video", Family: "h264", HW: true, DefQ: 23, MaxQ: 51},
	{Name: "h264_amf", Label: "H.264 (AMD GPU)", Kind: "video", Family: "h264", HW: true, DefQ: 23, MaxQ: 51},
	{Name: "libx265", Label: "H.265 / HEVC (CPU, smaller files)", Kind: "video", Family: "hevc", DefQ: 24, MaxQ: 51},
	{Name: "hevc_nvenc", Label: "H.265 / HEVC (NVIDIA GPU)", Kind: "video", Family: "hevc", HW: true, DefQ: 25, MaxQ: 51},
	{Name: "hevc_qsv", Label: "H.265 / HEVC (Intel GPU)", Kind: "video", Family: "hevc", HW: true, DefQ: 25, MaxQ: 51},
	{Name: "hevc_amf", Label: "H.265 / HEVC (AMD GPU)", Kind: "video", Family: "hevc", HW: true, DefQ: 25, MaxQ: 51},
	{Name: "libsvtav1", Label: "AV1 (SVT-AV1, CPU)", Kind: "video", Family: "av1", DefQ: 32, MaxQ: 63},
	{Name: "libaom-av1", Label: "AV1 (libaom, CPU, slow)", Kind: "video", Family: "av1", DefQ: 32, MaxQ: 63},
	{Name: "av1_nvenc", Label: "AV1 (NVIDIA GPU)", Kind: "video", Family: "av1", HW: true, DefQ: 30, MaxQ: 51},
	{Name: "av1_qsv", Label: "AV1 (Intel GPU)", Kind: "video", Family: "av1", HW: true, DefQ: 30, MaxQ: 51},
	{Name: "av1_amf", Label: "AV1 (AMD GPU)", Kind: "video", Family: "av1", HW: true, DefQ: 30, MaxQ: 51},
	{Name: "libvpx-vp9", Label: "VP9 (CPU)", Kind: "video", Family: "vp9", DefQ: 32, MaxQ: 63},
	{Name: "aac", Label: "AAC", Kind: "audio", Family: "aac"},
	{Name: "libopus", Label: "Opus", Kind: "audio", Family: "opus"},
	{Name: "libmp3lame", Label: "MP3", Kind: "audio", Family: "mp3"},
	{Name: "flac", Label: "FLAC (lossless)", Kind: "audio", Family: "flac"},
}

func codecByName(name string) (Codec, bool) {
	for _, c := range codecCatalog {
		if c.Name == name {
			return c, true
		}
	}
	return Codec{}, false
}

// Container lists the codec families a re-encoded output may use.
type Container struct {
	Ext   string   `json:"ext"`
	Label string   `json:"label"`
	Video []string `json:"video"`
	Audio []string `json:"audio"`
}

var containers = []Container{
	{Ext: "mp4", Label: "MP4", Video: []string{"h264", "hevc", "av1"}, Audio: []string{"aac", "mp3", "opus"}},
	{Ext: "mkv", Label: "MKV", Video: []string{"h264", "hevc", "av1", "vp9"}, Audio: []string{"aac", "opus", "mp3", "flac"}},
	{Ext: "mov", Label: "MOV", Video: []string{"h264", "hevc"}, Audio: []string{"aac", "mp3"}},
	{Ext: "webm", Label: "WebM", Video: []string{"vp9", "av1"}, Audio: []string{"opus"}},
}

// copyExts are the output containers allowed for lossless (stream copy) export.
var copyExts = []string{"mp4", "mkv", "mov", "m4v", "webm", "ts", "mts", "m2ts", "avi", "flv", "wmv", "mpg", "mpeg", "3gp", "mxf", "ogv"}

// mediaExts are shown by the built-in file browser.
var mediaExts = []string{"mp4", "m4v", "mov", "mkv", "webm", "avi", "wmv", "asf", "flv", "mpg", "mpeg", "vob", "ts", "mts", "m2ts", "3gp", "mxf", "ogv", "dv"}

func containerByExt(ext string) (Container, bool) {
	for _, c := range containers {
		if c.Ext == ext {
			return c, true
		}
	}
	return Container{}, false
}

type Tools struct {
	FFmpeg  string
	FFprobe string
	Version string

	mu       sync.Mutex
	encoders map[string]bool // usable encoders (hardware ones are test-encoded)
	ready    bool
}

func findTools(ffmpegFlag, ffprobeFlag string) *Tools {
	t := &Tools{FFmpeg: ffmpegFlag, FFprobe: ffprobeFlag}
	if t.FFmpeg == "" {
		t.FFmpeg = findTool("ffmpeg")
	}
	if t.FFprobe == "" {
		t.FFprobe = findTool("ffprobe")
	}
	if t.FFmpeg != "" {
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		if out, err := command(ctx, t.FFmpeg, "-hide_banner", "-version").Output(); err == nil {
			line, _, _ := strings.Cut(string(out), "\n")
			t.Version = strings.TrimSpace(line)
		} else {
			t.FFmpeg = ""
		}
	}
	return t
}

// findTool looks next to the executable first, then on PATH.
func findTool(name string) string {
	file := name + exeSuffix
	if exe, err := os.Executable(); err == nil {
		if real, err := filepath.EvalSymlinks(exe); err == nil {
			exe = real
		}
		dir := filepath.Dir(exe)
		for _, d := range []string{dir, filepath.Join(dir, "ffmpeg"), filepath.Join(dir, "bin"), filepath.Join(dir, "ffmpeg", "bin")} {
			p := filepath.Join(d, file)
			if st, err := os.Stat(p); err == nil && !st.IsDir() {
				return p
			}
		}
	}
	if p, err := exec.LookPath(file); err == nil {
		if abs, err := filepath.Abs(p); err == nil {
			return abs
		}
		return p
	}
	return ""
}

func (t *Tools) OK() bool { return t.FFmpeg != "" && t.FFprobe != "" }

func (t *Tools) EncodersReady() bool {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.ready
}

func (t *Tools) HasEncoder(name string) bool {
	t.mu.Lock()
	defer t.mu.Unlock()
	return t.encoders[name]
}

// CodecList returns the catalog entries this ffmpeg build can use.
func (t *Tools) CodecList() []Codec {
	t.mu.Lock()
	defer t.mu.Unlock()
	var out []Codec
	for _, c := range codecCatalog {
		if t.encoders[c.Name] {
			out = append(out, c)
		}
	}
	return out
}

// detectEncoders lists the encoders compiled into ffmpeg and test-encodes a
// few frames with each hardware encoder, since being compiled in does not mean
// the GPU is present.
func (t *Tools) detectEncoders() {
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()
	out, err := command(ctx, t.FFmpeg, "-hide_banner", "-encoders").Output()
	compiled := map[string]bool{}
	if err == nil {
		for _, line := range strings.Split(string(out), "\n") {
			f := strings.Fields(line)
			if len(f) >= 2 && len(f[0]) == 6 && (f[0][0] == 'V' || f[0][0] == 'A') {
				compiled[f[1]] = true
			}
		}
	}
	usable := map[string]bool{}
	var wg sync.WaitGroup
	var mu sync.Mutex
	for _, c := range codecCatalog {
		if !compiled[c.Name] {
			continue
		}
		if !c.HW {
			usable[c.Name] = true
			continue
		}
		wg.Add(1)
		go func(name string) {
			defer wg.Done()
			ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
			defer cancel()
			err := command(ctx, t.FFmpeg, "-hide_banner", "-loglevel", "error", "-nostdin",
				"-f", "lavfi", "-i", "color=c=black:s=640x360:r=25", "-frames:v", "5",
				"-pix_fmt", hwPixFmt(name), "-c:v", name, "-f", "null", "-").Run()
			if err == nil {
				mu.Lock()
				usable[name] = true
				mu.Unlock()
			}
		}(c.Name)
	}
	wg.Wait()
	t.mu.Lock()
	t.encoders = usable
	t.ready = true
	t.mu.Unlock()
}

func hwPixFmt(encoder string) string {
	if strings.HasSuffix(encoder, "_qsv") {
		return "nv12"
	}
	return "yuv420p"
}

// command builds an exec.Cmd that never pops up a console window.
func command(ctx context.Context, name string, args ...string) *exec.Cmd {
	cmd := exec.CommandContext(ctx, name, args...)
	cmd.WaitDelay = 3 * time.Second
	hideConsole(cmd)
	return cmd
}

// ffmpegBaseArgs are shared by every ffmpeg run that reports progress.
func ffmpegBaseArgs() []string {
	return []string{"-hide_banner", "-nostdin", "-y", "-loglevel", "error", "-progress", "pipe:1", "-nostats"}
}

// runFFmpeg runs ffmpeg and calls onProgress with the output position in
// seconds. It returns ctx.Err() when canceled.
func (t *Tools) runFFmpeg(ctx context.Context, args []string, onProgress func(sec float64, speed string)) error {
	cmd := command(ctx, t.FFmpeg, args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	stderr := &tailBuffer{max: 16 << 10}
	cmd.Stderr = stderr
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("could not start ffmpeg: %w", err)
	}
	sc := bufio.NewScanner(stdout)
	var sec float64
	var speed string
	for sc.Scan() {
		k, v, ok := strings.Cut(sc.Text(), "=")
		if !ok {
			continue
		}
		v = strings.TrimSpace(v)
		switch k {
		case "out_time_us":
			if us, err := strconv.ParseInt(v, 10, 64); err == nil && us >= 0 {
				sec = float64(us) / 1e6
			}
		case "speed":
			if v != "N/A" {
				speed = v
			}
		case "progress":
			if onProgress != nil {
				onProgress(sec, speed)
			}
		}
	}
	err = cmd.Wait()
	if ctx.Err() != nil {
		return ctx.Err()
	}
	if err != nil {
		msg := lastLines(strings.TrimSpace(stderr.String()), 12)
		if msg == "" {
			msg = err.Error()
		}
		return errors.New(msg)
	}
	return nil
}

// tailBuffer keeps the last max bytes written to it.
type tailBuffer struct {
	mu  sync.Mutex
	max int
	buf []byte
}

func (b *tailBuffer) Write(p []byte) (int, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.buf = append(b.buf, p...)
	if len(b.buf) > b.max {
		b.buf = b.buf[len(b.buf)-b.max:]
	}
	return len(p), nil
}

func (b *tailBuffer) String() string {
	b.mu.Lock()
	defer b.mu.Unlock()
	return string(b.buf)
}

func lastLines(s string, n int) string {
	lines := strings.Split(s, "\n")
	if len(lines) > n {
		lines = lines[len(lines)-n:]
	}
	return strings.Join(lines, "\n")
}

// ffTime formats seconds for ffmpeg's -ss/-t options.
func ffTime(sec float64) string {
	return strconv.FormatFloat(sec, 'f', 6, 64)
}