av1dae/internal/encoder/encoder.go

package encoder

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"av1dae/internal/logger"
	"av1dae/internal/status"
	"av1dae/pkg/types"
)

var idetSummaryRegex = regexp.MustCompile(`Multi frame detection:\s+TFF:\s*(\d+)\s+BFF:\s*(\d+)\s+Progressive:\s*(\d+)\s+Undetermined:\s*(\d+)`)

type Encoder struct {
	ffmpegPath  string
	ffprobePath string
	opusencPath string
	log         *logger.Logger
	tracker     *status.Tracker

	procMu sync.Mutex
	cur    *os.Process // the in-flight video encode, for pause/resume
}

// Pause freezes the active video encode in place via SIGSTOP. libsvtav1 runs in
// the ffmpeg process (no forked children), so one signal suspends all its
// threads. No-op if nothing is encoding. The process keeps its memory and
// partial output and resumes exactly where it left off.
func (e *Encoder) Pause() error {
	e.procMu.Lock()
	defer e.procMu.Unlock()
	if e.cur == nil {
		return nil
	}
	if err := e.cur.Signal(syscall.SIGSTOP); err != nil {
		return err
	}
	if e.tracker != nil {
		e.tracker.SetPaused(true)
	}
	return nil
}

// Resume thaws a paused encode via SIGCONT. No-op if nothing is encoding.
func (e *Encoder) Resume() error {
	e.procMu.Lock()
	defer e.procMu.Unlock()
	if e.cur == nil {
		return nil
	}
	if err := e.cur.Signal(syscall.SIGCONT); err != nil {
		return err
	}
	if e.tracker != nil {
		e.tracker.SetPaused(false)
	}
	return nil
}

func (e *Encoder) setProc(p *os.Process) {
	e.procMu.Lock()
	e.cur = p
	e.procMu.Unlock()
}

func (e *Encoder) clearProc() {
	e.procMu.Lock()
	e.cur = nil
	e.procMu.Unlock()
	if e.tracker != nil {
		e.tracker.SetPaused(false)
	}
}

type StreamInfo struct {
	Width              int    `json:"width"`
	Height             int    `json:"height"`
	SampleAspectRatio  string `json:"sample_aspect_ratio"`
	DisplayAspectRatio string `json:"display_aspect_ratio"`
	CodecName          string `json:"codec_name"`
	CodecType          string `json:"codec_type"`
}

type ProbeResult struct {
	Streams []StreamInfo `json:"streams"`
}

type StreamMetadata struct {
	Index     int    `json:"index"`
	CodecType string `json:"codec_type"`
	CodecName string `json:"codec_name"`
	Channels  int    `json:"channels"`
	Language  string `json:"tags"`
	Title     string `json:"title"`
}

func New(log *logger.Logger, tracker *status.Tracker) *Encoder {
	return &Encoder{
		ffmpegPath:  "ffmpeg",
		ffprobePath: "ffprobe",
		opusencPath: "opusenc",
		log:         log,
		tracker:     tracker,
	}
}

// runCmd executes the command and emits a debug entry containing the full
// command line and its combined output. file is the source file the command
// is acting on (may be empty).
func (e *Encoder) runCmd(ctx context.Context, label, file, name string, args []string) ([]byte, error) {
	cmd := exec.CommandContext(ctx, name, args...)
	out, err := cmd.CombinedOutput()
	extra, _ := json.Marshal(struct {
		Cmd    string `json:"cmd"`
		Output string `json:"output"`
	}{
		Cmd:    name + " " + strings.Join(args, " "),
		Output: string(out),
	})
	e.log.Debug(label, file, string(extra))
	return out, err
}

// runCmdProgress runs ffmpeg with `-progress pipe:1`, streaming progress
// samples to the tracker (instead of buffering all output like runCmd). stdout
// carries only the key=value progress stream; stderr carries real errors and is
// returned for the caller's error message. Used solely for the video encode —
// the one step long enough to be worth watching live.
func (e *Encoder) runCmdProgress(ctx context.Context, label, file, name string, args []string) ([]byte, error) {
	cmd := exec.CommandContext(ctx, name, args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	var stderr bytes.Buffer
	cmd.Stderr = &stderr

	if err := cmd.Start(); err != nil {
		return nil, err
	}
	// Publish the process so Pause/Resume can signal it; clear on exit.
	e.setProc(cmd.Process)
	defer e.clearProc()

	// Reads stdout to EOF (when ffmpeg exits), so Wait below is safe afterwards.
	var lastLog time.Time
	_ = status.ScanProgress(stdout, func(s status.ProgressSample) {
		if e.tracker == nil {
			return
		}
		e.tracker.Update(s.OutTimeSec, s.FPS, s.Speed)
		if time.Since(lastLog) >= 5*time.Second {
			lastLog = time.Now()
			snap := e.tracker.Snapshot()
			e.log.Progress(fmt.Sprintf("encoding %s · %.1f%% · %.1ffps · %.2fx · ETA %s",
				filepath.Base(file), snap.Percent, snap.FPS, snap.Speed, fmtETA(snap.ETASec)))
		}
	})

	err = cmd.Wait()

	extra, _ := json.Marshal(struct {
		Cmd    string `json:"cmd"`
		Output string `json:"output"`
	}{
		Cmd:    name + " " + strings.Join(args, " "),
		Output: stderr.String(),
	})
	e.log.Debug(label, file, string(extra))
	return stderr.Bytes(), err
}

// fmtETA renders a seconds count as a compact "11h03m" / "4m12s" / "9s" string.
func fmtETA(sec int) string {
	if sec <= 0 {
		return "--"
	}
	d := time.Duration(sec) * time.Second
	switch {
	case d >= time.Hour:
		return fmt.Sprintf("%dh%02dm", int(d.Hours()), int(d.Minutes())%60)
	case d >= time.Minute:
		return fmt.Sprintf("%dm%02ds", int(d.Minutes()), sec%60)
	default:
		return fmt.Sprintf("%ds", sec)
	}
}

// GetDuration returns the source container duration in seconds via ffprobe.
func (e *Encoder) GetDuration(ctx context.Context, path string) (float64, error) {
	args := []string{"-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", path}
	out, err := e.runCmd(ctx, "ffprobe duration", path, e.ffprobePath, args)
	if err != nil {
		return 0, fmt.Errorf("ffprobe duration: %w", err)
	}
	d, err := strconv.ParseFloat(strings.TrimSpace(string(out)), 64)
	if err != nil {
		return 0, fmt.Errorf("parsing duration %q: %w", strings.TrimSpace(string(out)), err)
	}
	return d, nil
}

func (e *Encoder) CheckDeps() error {
	if _, err := exec.LookPath(e.ffmpegPath); err != nil {
		return fmt.Errorf("ffmpeg not found")
	}
	if _, err := exec.LookPath(e.ffprobePath); err != nil {
		return fmt.Errorf("ffprobe not found")
	}
	if _, err := exec.LookPath(e.opusencPath); err != nil {
		return fmt.Errorf("opusenc not found")
	}
	return nil
}

func (e *Encoder) GetStreamLanguages(ctx context.Context, path string) ([]StreamMetadata, error) {
	args := []string{"-v", "error", "-show_streams", "-print_format", "json", path}
	output, err := e.runCmd(ctx, "ffprobe stream languages", path, e.ffprobePath, args)
	if err != nil {
		return nil, fmt.Errorf("ffprobe error: %w", err)
	}

	var result struct {
		Streams []struct {
			Index     int               `json:"index"`
			CodecType string            `json:"codec_type"`
			CodecName string            `json:"codec_name"`
			Channels  int               `json:"channels"`
			Tags      map[string]string `json:"tags"`
		} `json:"streams"`
	}
	if err := json.Unmarshal(output, &result); err != nil {
		return nil, fmt.Errorf("parsing ffprobe: %w", err)
	}

	var streams []StreamMetadata
	for _, s := range result.Streams {
		lang := ""
		title := ""
		if s.Tags != nil {
			lang = s.Tags["language"]
			title = s.Tags["title"]
		}
		streams = append(streams, StreamMetadata{
			Index:     s.Index,
			CodecType: s.CodecType,
			CodecName: s.CodecName,
			Channels:  s.Channels,
			Language:  lang,
			Title:     title,
		})
	}

	streamsJSON, _ := json.Marshal(streams)
	e.log.Debug("parsed stream languages", path, string(streamsJSON))
	return streams, nil
}

func (e *Encoder) GetMediaInfo(ctx context.Context, path string) (width, height int, interlaced bool, err error) {
	probeArgs := []string{"-v", "error", "-show_streams", "-print_format", "json", path}
	output, err := e.runCmd(ctx, "ffprobe media info", path, e.ffprobePath, probeArgs)
	if err != nil {
		return 0, 0, false, fmt.Errorf("ffprobe error: %w", err)
	}

	var result ProbeResult
	if err := json.Unmarshal(output, &result); err != nil {
		return 0, 0, false, fmt.Errorf("parsing ffprobe output: %w", err)
	}

	foundVideo := false
	for _, stream := range result.Streams {
		if stream.CodecType == "video" {
			width = stream.Width
			height = stream.Height
			foundVideo = true
			info, _ := json.Marshal(map[string]interface{}{
				"width": width, "height": height, "codec": stream.CodecName, "sar": stream.SampleAspectRatio,
			})
			e.log.Debug("detected video stream", path, string(info))
			break
		}
	}
	if !foundVideo {
		return 0, 0, false, fmt.Errorf("no video stream found")
	}

	idetArgs := []string{
		"-hide_banner",
		"-nostats",
		"-i", path,
		"-vf", "idet",
		"-frames:v", "400",
		"-an", "-sn",
		"-f", "null", "-",
	}
	idetOut, _ := e.runCmd(ctx, "ffmpeg idet", path, e.ffmpegPath, idetArgs)
	interlaced = detectInterlaced(string(idetOut))

	info, _ := json.Marshal(map[string]bool{"interlaced": interlaced})
	e.log.Debug("interlace detection", path, string(info))

	return width, height, interlaced, nil
}

// detectInterlaced parses ffmpeg idet's "Multi frame detection" summary line
// and returns true when TFF+BFF clearly outnumber Progressive frames among
// the decided frames. Undetermined frames are ignored. If the summary line
// is missing, returns false (default to progressive).
func detectInterlaced(idetOutput string) bool {
	m := idetSummaryRegex.FindStringSubmatch(idetOutput)
	if m == nil {
		return false
	}
	tff, _ := strconv.Atoi(m[1])
	bff, _ := strconv.Atoi(m[2])
	prog, _ := strconv.Atoi(m[3])
	interlaced := tff + bff
	decided := interlaced + prog
	if decided == 0 {
		return false
	}
	return interlaced > prog
}

func (e *Encoder) calculateZscaleWidth(ctx context.Context, path string, originalHeight int) (string, int, error) {
	args := []string{"-v", "error", "-select_streams", "v:0", "-show_entries", "stream=width,height,sample_aspect_ratio", "-print_format", "json", path}
	output, err := e.runCmd(ctx, "ffprobe zscale info", path, e.ffprobePath, args)
	if err != nil {
		return "", 0, fmt.Errorf("ffprobe error: %w", err)
	}

	var result ProbeResult
	if err := json.Unmarshal(output, &result); err != nil {
		return "", 0, fmt.Errorf("parsing ffprobe: %w", err)
	}

	if len(result.Streams) == 0 {
		return "", 0, fmt.Errorf("no video streams found")
	}

	stream := result.Streams[0]
	width := stream.Width
	height := stream.Height
	sar := stream.SampleAspectRatio

	newWidth := width

	if sar == "1:1" || sar == "N/A" || sar == "" {
		return "", newWidth, nil
	}

	parts := strings.Split(sar, ":")
	if len(parts) != 2 {
		return "", newWidth, nil
	}
	num, err1 := strconv.Atoi(parts[0])
	den, err2 := strconv.Atoi(parts[1])
	if err1 != nil || err2 != nil || num == 0 || den == 0 {
		return "", newWidth, nil
	}

	// Round to nearest, then snap down to even (mod-2 widths preferred by AV1/H.264).
	newWidth = ((width*num + den/2) / den) &^ 1

	info, _ := json.Marshal(map[string]interface{}{
		"width": width, "height": height, "sar": sar, "new_width": newWidth,
	})
	e.log.Debug("zscale calculation", path, string(info))

	if newWidth == width {
		return "", newWidth, nil
	}

	zscale := fmt.Sprintf("zscale=w=%d:h=%d:filter=spline36", newWidth, height)
	return zscale, newWidth, nil
}

func (e *Encoder) Transcode(ctx context.Context, input, workDir string, job *types.Job, metadata *types.Metadata, interlaced bool, streamLangs []StreamMetadata) error {
	if e.tracker != nil {
		e.tracker.SetPhase(status.PhaseAudio)
	}
	audioWavs, err := e.extractAudio(ctx, input, workDir, streamLangs)
	if err != nil {
		return fmt.Errorf("extracting audio: %w", err)
	}

	opusFiles, err := e.encodeOpus(ctx, audioWavs, workDir)
	if err != nil {
		return fmt.Errorf("encoding opus: %w", err)
	}

	if err := e.encodeVideo(ctx, input, workDir, opusFiles, job, metadata, interlaced, streamLangs); err != nil {
		return fmt.Errorf("encoding video: %w", err)
	}

	return nil
}

// audioStreamsInSourceOrder returns the audio entries of streamLangs sorted by
// their source-side Index. The resulting slice position is the 0-based audio
// stream index used by ffmpeg selectors like `0:a:<n>`.
func audioStreamsInSourceOrder(streamLangs []StreamMetadata) []StreamMetadata {
	var audio []StreamMetadata
	for _, s := range streamLangs {
		if s.CodecType == "audio" {
			audio = append(audio, s)
		}
	}
	sort.Slice(audio, func(i, j int) bool { return audio[i].Index < audio[j].Index })
	return audio
}

func (e *Encoder) extractAudio(ctx context.Context, input, workDir string, streamLangs []StreamMetadata) ([]string, error) {
	audio := audioStreamsInSourceOrder(streamLangs)
	if len(audio) == 0 {
		return nil, fmt.Errorf("no audio streams found")
	}

	wavs := make([]string, 0, len(audio))
	for srcAudioIndex := range audio {
		wavPath := filepath.Join(workDir, fmt.Sprintf("audio.%d.wav", srcAudioIndex))
		args := []string{
			"-i", input,
			"-map", fmt.Sprintf("0:a:%d", srcAudioIndex),
			"-vn", "-c:a", "pcm_s16le", "-ar", "48000",
			"-f", "wav", wavPath,
		}
		out, err := e.runCmd(ctx, "ffmpeg extract audio", input, e.ffmpegPath, args)
		if err != nil {
			return wavs, fmt.Errorf("ffmpeg extract a:%d: %s %w", srcAudioIndex, out, err)
		}
		wavs = append(wavs, wavPath)
	}
	return wavs, nil
}

func (e *Encoder) encodeOpus(ctx context.Context, wavs []string, workDir string) ([]string, error) {
	var opusFiles []string
	for _, wav := range wavs {
		base := filepath.Base(strings.Replace(wav, ".wav", ".opus", 1))
		out := filepath.Join(workDir, base)
		args := []string{"--bitrate", "128k", wav, out}
		logOut, err := e.runCmd(ctx, "opusenc", wav, e.opusencPath, args)
		if err != nil {
			return nil, fmt.Errorf("opusenc: %s %w", logOut, err)
		}
		opusFiles = append(opusFiles, out)
	}
	return opusFiles, nil
}

// svtav1Params builds the -svtav1-params value, appending lp=N (logical
// processors / encoder thread count) only when lp > 0. lp=0 lets SVT-AV1
// auto-detect, preserving prior behavior.
func svtav1Params(lp int) string {
	p := "film-grain=10:film-grain-denoise=1:scd=1:qm-min=4:qm-max=15:keyint=10s"
	if lp > 0 {
		p += fmt.Sprintf(":lp=%d", lp)
	}
	return p
}

func (e *Encoder) encodeVideo(ctx context.Context, input, workDir string, opusFiles []string, job *types.Job, metadata *types.Metadata, interlaced bool, streamLangs []StreamMetadata) error {
	outFile := filepath.Join(workDir, "output.mkv")

	// Switch the tracker to the encode phase and feed it the source duration so
	// progress samples can be turned into a percentage. A failed duration probe
	// just means no percent — it must not abort the encode.
	if e.tracker != nil {
		e.tracker.SetPhase(status.PhaseEncoding)
		if dur, derr := e.GetDuration(ctx, input); derr == nil {
			e.tracker.SetTotal(dur)
		} else {
			e.log.Debug("duration probe failed", input, derr.Error())
		}
	}

	svtParams := svtav1Params(job.LP)

	zscaleStr, newWidth, err := e.calculateZscaleWidth(ctx, input, 0)
	if err != nil {
		return fmt.Errorf("calculating zscale: %w", err)
	}

	var filters []string
	if interlaced {
		filters = append(filters, "bwdif=mode=0:par=-1:-1")
	}
	if zscaleStr != "" {
		filters = append(filters, zscaleStr)
	}

	filterInfo, _ := json.Marshal(map[string]interface{}{
		"zscale": zscaleStr, "new_width": newWidth, "interlaced": interlaced, "vf": strings.Join(filters, ","),
	})
	e.log.Debug("video filter chain", input, string(filterInfo))

	args := []string{
		"-y",
		"-i", input,
	}

	for _, opus := range opusFiles {
		args = append(args, "-i", opus)
	}

	if len(filters) > 0 {
		args = append(args, "-vf", strings.Join(filters, ","))
	}
	args = append(args, "-c:v", "libsvtav1")
	args = append(args, "-crf", strconv.Itoa(job.CRF))
	args = append(args, "-preset", strconv.Itoa(job.Preset))
	args = append(args, "-svtav1-params", svtParams)
	args = append(args, "-pix_fmt", "yuv420p10le")

	args = append(args, "-map", "0:v")
	args = append(args, "-map", "0:s?")

	for i := range opusFiles {
		args = append(args, "-map", fmt.Sprintf("%d:a", i+1))
	}

	args = append(args, "-c:a", "copy")
	args = append(args, "-c:s", "copy")
	args = append(args, "-map_metadata", "-1")

	// Walk opusFiles by output audio index (0..N-1) and look up the source
	// audio language at the same source-audio-index position. This keeps the
	// `-metadata:s:a:i` index aligned with the output stream order, regardless
	// of whether some source streams lacked a language tag.
	audioStreams := audioStreamsInSourceOrder(streamLangs)
	for i := range opusFiles {
		if i >= len(audioStreams) {
			break
		}
		lang := audioStreams[i].Language
		if lang == "" {
			continue
		}
		args = append(args, fmt.Sprintf("-metadata:s:a:%d", i), fmt.Sprintf("language=%s", lang))
	}

	for _, stream := range streamLangs {
		if stream.CodecType == "subtitle" && stream.Language != "" {
			idx := 0
			for _, s := range streamLangs {
				if s.CodecType == "subtitle" && s.Index < stream.Index {
					idx++
				}
			}
			args = append(args, fmt.Sprintf("-metadata:s:s:%d", idx), fmt.Sprintf("language=%s", stream.Language))
		}
	}

	args = append(args, "-metadata", fmt.Sprintf("TITLE=%s", metadata.Title))
	args = append(args, "-metadata", fmt.Sprintf("DATE_RELEASED=%s", metadata.DateReleased))
	args = append(args, "-metadata", fmt.Sprintf("IMDBID=%s", metadata.IMDBID))
	args = append(args, "-metadata", fmt.Sprintf("ORIGINAL_MEDIA_TYPE=%s", metadata.OriginalMedia))

	if metadata.IsSeries {
		args = append(args, "-metadata", fmt.Sprintf("COLLECTION=%s", metadata.Collection))
		args = append(args, "-metadata", fmt.Sprintf("SEASON=%s", metadata.Season))
		args = append(args, "-metadata", fmt.Sprintf("EPISODE=%s", metadata.Episode))
		if metadata.TVMazeID != "" {
			args = append(args, "-metadata", fmt.Sprintf("TVMAZE_ID=%s", metadata.TVMazeID))
		}
	}

	args = append(args, outFile)

	args = append([]string{"-hide_banner", "-v", "error", "-progress", "pipe:1", "-nostats"}, args...)
	out, err := e.runCmdProgress(ctx, "ffmpeg encode", input, e.ffmpegPath, args)
	if err != nil {
		return fmt.Errorf("ffmpeg encode: %s %w", out, err)
	}

	return nil
}