av1dae/internal/encoder/encoder.go

package encoder

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

	"videnc-vibe/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
}

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"`
	Language  string `json:"tags"`
	Title     string `json:"title"`
}

func (e *Encoder) GetStreamLanguages(ctx context.Context, path string) ([]StreamMetadata, error) {
	cmd := exec.CommandContext(ctx, e.ffprobePath, "-v", "error", "-show_streams", "-print_format", "json", path)
	output, err := cmd.CombinedOutput()
	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"`
			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,
			Language:  lang,
			Title:     title,
		})
	}

	fmt.Printf("DEBUG stream languages: %+v\n", streams)
	return streams, nil
}

func New() *Encoder {
	return &Encoder{
		ffmpegPath:  "ffmpeg",
		ffprobePath: "ffprobe",
		opusencPath: "opusenc",
	}
}

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) GetMediaInfo(ctx context.Context, path string) (width, height int, interlaced bool, err error) {
	cmd := exec.CommandContext(ctx, e.ffprobePath, "-v", "error", "-show_streams", "-print_format", "json", path)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return 0, 0, false, fmt.Errorf("ffprobe error: %w", err)
	}

	fmt.Printf("DEBUG ffprobe output: %s\n", string(output))

	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
			fmt.Printf("DEBUG detected video: %dx%d, codec=%s, SAR=%s\n", width, height, stream.CodecName, stream.SampleAspectRatio)
			break
		}
	}
	if !foundVideo {
		return 0, 0, false, fmt.Errorf("no video stream found")
	}

	cmd = exec.CommandContext(ctx, e.ffmpegPath,
		"-hide_banner",
		"-nostats",
		"-i", path,
		"-vf", "idet",
		"-frames:v", "400",
		"-an", "-sn",
		"-f", "null", "-",
	)
	idetOut, _ := cmd.CombinedOutput()
	interlaced = detectInterlaced(string(idetOut))

	fmt.Printf("DEBUG interlaced: %v\n", interlaced)

	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) {
	cmd := exec.CommandContext(ctx, e.ffprobePath, "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=width,height,sample_aspect_ratio", "-print_format", "json", path)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return "", 0, fmt.Errorf("ffprobe error: %w", err)
	}

	fmt.Printf("DEBUG ffprobe for zscale: %s\n", string(output))

	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
	fmt.Printf("DEBUG stream: width=%d, height=%d, sar=%s\n", width, height, sar)

	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
	fmt.Printf("DEBUG calculated new width: %d (sar=%s)\n", newWidth, sar)

	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 {
	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))
		cmd := exec.CommandContext(ctx, e.ffmpegPath, "-i", input,
			"-map", fmt.Sprintf("0:a:%d", srcAudioIndex),
			"-vn", "-c:a", "pcm_s16le", "-ar", "48000",
			"-f", "wav", wavPath)
		if out, err := cmd.CombinedOutput(); 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)
		cmd := exec.CommandContext(ctx, e.opusencPath, "--bitrate", "128k", wav, out)
		if logOut, err := cmd.CombinedOutput(); err != nil {
			return nil, fmt.Errorf("opusenc: %s %w", logOut, err)
		}
		opusFiles = append(opusFiles, out)
	}
	return opusFiles, nil
}

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

	svtParams := fmt.Sprintf("film-grain=10:film-grain-denoise=1:scd=1:qm-min=4:qm-max=15:keyint=10s")

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

	fmt.Printf("DEBUG zscale: %s (newWidth=%d)\n", zscaleStr, newWidth)

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

	fmt.Printf("DEBUG final vf: %s\n", strings.Join(filters, ","))

	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))
		fmt.Printf("DEBUG setting audio language: stream %d -> language=%s\n", i, 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))
			fmt.Printf("DEBUG setting subtitle language: stream %d -> language=%s\n", idx, 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"}, args...)
	cmd := exec.CommandContext(ctx, e.ffmpegPath, args...)
	fmt.Printf("DEBUG FFmpeg command: ffmpeg %s\n", strings.Join(args, " "))
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("ffmpeg encode: %s %w", out, err)
	}

	return nil
}