av1-go/main.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"
)

var (
	RequiredTools = []string{
		"ffmpeg", "ffprobe", "mkvmerge",
		"opusenc", "mediainfo", "av1an", "HandBrakeCLI", "ffmsindex",
	}
	DirCompleted = "completed"
	DirOriginal  = "original"
	DirConvLogs  = "conv_logs"
	RemuxCodecs  = map[string]bool{"aac": true, "opus": true}
)

func main() {
	noDownmix := flag.Bool("no-downmix", false, "Preserve original audio channel layout")
	autocrop := flag.Bool("autocrop", false, "Automatically detect and crop black bars")
	speed := flag.String("speed", "", "Set encoding speed (slower, slow, medium, fast, faster)")
	quality := flag.String("quality", "", "Set encoding quality (lowest, low, medium, high, higher)")
	grain := flag.Int("grain", -1, "Set film-grain value")
	flag.Parse()

	CheckTools(RequiredTools)

	if *speed != "" {
		SvtAv1Params["speed"] = *speed
	}
	if *quality != "" {
		SvtAv1Params["quality"] = *quality
	}
	if *grain != -1 {
		SvtAv1Params["film-grain"] = *grain
	}

	os.MkdirAll(DirCompleted, 0755)
	os.MkdirAll(DirOriginal, 0755)
	os.MkdirAll(DirConvLogs, 0755)

	for {
		files, err := filepath.Glob("*.mkv")
		if err != nil {
			fmt.Println("Error listing files:", err)
			return
		}

		// Filter files
		var toProcess []string
		for _, f := range files {
			if strings.HasSuffix(f, ".ut.mkv") || strings.HasPrefix(f, "temp-") || strings.HasPrefix(f, "output-") || strings.HasSuffix(f, ".cfr_temp.mkv") {
				continue
			}
			toProcess = append(toProcess, f)
		}
		sort.Strings(toProcess)

		if len(toProcess) == 0 {
			fmt.Println("No more .mkv files found to process. Exiting.")
			break
		}

		file := toProcess[0]
		if !IsFfmpegDecodable(file) {
			fmt.Printf("ERROR: ffmpeg cannot decode '%s'. Skipping.\n", file)
			os.Rename(file, filepath.Join(DirOriginal, filepath.Base(file)))
			continue
		}

		ProcessFile(file, *noDownmix, *autocrop)
	}
}

func ProcessFile(filePath string, noDownmix, autocrop bool) {
	baseName := strings.TrimSuffix(filepath.Base(filePath), filepath.Ext(filePath))
	logFileName := baseName + ".log"
	logFilePath := filepath.Join(DirConvLogs, logFileName)

	SetupLogging(logFilePath)
	LogInfo("STARTING LOG FOR: %s", filepath.Base(filePath))
	LogInfo("Processing started at: %s", time.Now().Format(time.RFC3339))
	absPath, _ := filepath.Abs(filePath)
	LogInfo("Full input file path: %s", absPath)
	LogInfo(strings.Repeat("-", 80))

	intermediateOutput := "output-" + filepath.Base(filePath)
	var audioTempDir string
	var handbrakeCleanup string
	processingError := false
	startTime := time.Now()

	defer func() {
		LogInfo("--- Starting Universal Cleanup ---")
		if audioTempDir != "" {
			os.RemoveAll(audioTempDir)
			LogInfo("    - Deleted audio temp dir: %s", audioTempDir)
		}

		if _, err := os.Stat(intermediateOutput); err == nil {
			if processingError {
				LogInfo("    - WARNING: Error occurred. Intermediate file preserved: %s", intermediateOutput)
			} else {
				os.Remove(intermediateOutput)
				LogInfo("    - Deleted intermediate output file: %s", intermediateOutput)
			}
		} else {
			if !processingError {
				LogInfo("    - Intermediate output file successfully moved.")
			}
		}

		LogInfo("\nTotal runtime: %s", time.Since(startTime))
		LogInfo("FINISHED LOG FOR: %s", filepath.Base(filePath))
		if LogFileHandle != nil {
			LogFileHandle.Close()
			LogFileHandle = nil
		}
	}()

	// Create audio temp dir
	var err error
	audioTempDir, err = ioutil.TempDir("", "anime_audio_")
	if err != nil {
		LogInfo("Error creating temp dir: %v", err)
		processingError = true
		return
	}
	LogInfo("Audio temporary directory created at: %s", audioTempDir)

	// Analyze file
	LogInfo("Analyzing file: %s", absPath)

	ffprobeJson, err := RunCmd("ffprobe", []string{"-v", "quiet", "-print_format", "json", "-show_streams", "-show_format", absPath}, true)
	if err != nil {
		LogInfo("ffprobe failed: %v", err)
		processingError = true
		return
	}
	var ffprobeInfo FFprobeOutput
	json.Unmarshal([]byte(ffprobeJson), &ffprobeInfo)

	mkvmergeJson, err := RunCmd("mkvmerge", []string{"-J", absPath}, true)
	if err != nil {
		LogInfo("mkvmerge failed: %v", err)
		processingError = true
		return
	}
	var mkvInfo MkvMergeOutput
	json.Unmarshal([]byte(mkvmergeJson), &mkvInfo)

	mediainfoJson, err := RunCmd("mediainfo", []string{"--Output=JSON", "-f", absPath}, true)
	if err != nil {
		LogInfo("mediainfo failed: %v", err)
		processingError = true
		return
	}
	var mediaInfo MediaInfoOutput
	json.Unmarshal([]byte(mediainfoJson), &mediaInfo)

	// VFR Detection
	isVfr := false
	targetCfrFps := ""

	// Find video track in mediainfo
	var videoTrack *MediaInfoTrack
	for _, t := range mediaInfo.Media.Track {
		if t.Type == "Video" {
			// Make a copy to avoid pointer to loop variable issues if we were storing it,
			// but here we just assign to a pointer variable.
			val := t
			videoTrack = &val
			break
		}
	}

	if videoTrack != nil {
		mode := strings.ToUpper(videoTrack.FrameRateMode)
		if mode == "VFR" || mode == "VARIABLE" {
			isVfr = true
			LogInfo("    - Detected VFR based on MediaInfo FrameRate_Mode: %s", videoTrack.FrameRateMode)

			origStr := videoTrack.FrameRateOriginalString
			re := regexp.MustCompile(`\((\d+/\d+)\)`)
			match := re.FindStringSubmatch(origStr)
			if len(match) > 1 {
				targetCfrFps = match[1]
			} else {
				targetCfrFps = videoTrack.FrameRateOriginal
			}

			if targetCfrFps == "" {
				targetCfrFps = videoTrack.FrameRate
				if targetCfrFps != "" {
					LogInfo("    - Using MediaInfo FrameRate (%s) as fallback.", targetCfrFps)
				}
			}

			if targetCfrFps != "" {
				LogInfo("    - Target CFR for HandBrake: %s", targetCfrFps)
				if strings.Contains(targetCfrFps, "/") {
					parts := strings.Split(targetCfrFps, "/")
					if len(parts) == 2 {
						num, _ := strconv.ParseFloat(parts[0], 64)
						den, _ := strconv.ParseFloat(parts[1], 64)
						if den != 0 {
							targetCfrFps = fmt.Sprintf("%.3f", num/den)
							LogInfo("    - Converted fractional FPS to decimal: %s", targetCfrFps)
						}
					}
				}
			} else {
				LogInfo("    - Warning: VFR detected but no target CFR found.")
				isVfr = false
			}
		} else {
			LogInfo("    - Video appears to be CFR.")
		}
	}

	autocropFilter := ""
	if autocrop {
		LogInfo("--- Running autocrop detection ---")
		autocropFilter = DetectAutocropFilter(absPath)
		if autocropFilter != "" {
			LogInfo("    - Autocrop filter detected: %s", autocropFilter)
		} else {
			LogInfo("    - No crop needed or detected.")
		}
	}

	encodedVideoFile, hbCleanup, err := ConvertVideo(baseName, absPath, isVfr, targetCfrFps, autocropFilter)
	if err != nil {
		LogInfo("Video conversion failed: %v", err)
		processingError = true
		return
	}
	handbrakeCleanup = hbCleanup

	LogInfo("--- Starting Audio Processing ---")
	var processedAudio []ProcessedAudioFile
	var remuxAudioTracks []string

	// Map mkvmerge tracks
	mkvAudioTracks := make(map[int]MkvMergeTrack)
	for _, t := range mkvInfo.Tracks {
		if t.Type == "audio" {
			mkvAudioTracks[t.ID] = t
		}
	}

	// Map mediainfo tracks by StreamOrder
	mediaAudioTracks := make(map[int]MediaInfoTrack)
	for _, t := range mediaInfo.Media.Track {
		if t.Type == "Audio" {
			so, _ := strconv.Atoi(t.StreamOrder)
			mediaAudioTracks[so] = t
		}
	}

	for _, s := range ffprobeInfo.Streams {
		if s.CodecType != "audio" {
			continue
		}

		// Find mkv track
		var mkvTrack MkvMergeTrack
		found := false
		for _, t := range mkvInfo.Tracks {
			if t.Type == "audio" && t.Properties.StreamID == s.Index {
				mkvTrack = t
				found = true
				break
			}
		}
		if !found {
			// Fallback by index if possible, but stream index might not match track index directly
			// Simplified: just skip or warn? Python script had fallback.
		}

		trackID := mkvTrack.ID
		trackTitle := mkvTrack.Properties.TrackName
		lang := s.Tags["language"]
		if lang == "" {
			lang = "und"
		}

		// Delay
		delay := 0
		if mt, ok := mediaAudioTracks[s.Index]; ok {
			if mt.VideoDelay != "" {
				dVal, _ := strconv.ParseFloat(mt.VideoDelay, 64)
				// MediaInfo usually gives ms if int, or s if float?
				// Python script logic: if < 1, *1000.
				if dVal < 1 && dVal > -1 && dVal != 0 {
					delay = int(dVal * 1000)
				} else {
					delay = int(dVal)
				}
			}
		}

		LogInfo("Processing Audio Stream #%d (TID: %d, Codec: %s, Channels: %d)", s.Index, trackID, s.CodecName, s.Channels)

		if RemuxCodecs[s.CodecName] {
			remuxAudioTracks = append(remuxAudioTracks, strconv.Itoa(trackID))
		} else {
			opusFile, err := ConvertAudioTrack(s.Index, s.Channels, lang, audioTempDir, absPath, !noDownmix)
			if err != nil {
				LogInfo("Audio conversion failed: %v", err)
				processingError = true
				return
			}
			processedAudio = append(processedAudio, ProcessedAudioFile{
				Path:     opusFile,
				Language: lang,
				Title:    trackTitle,
				Delay:    delay,
			})
		}
	}
	LogInfo("--- Finished Audio Processing ---")

	LogInfo("Assembling final file with mkvmerge...")
	mkvArgs := []string{"-o", intermediateOutput, encodedVideoFile}

	for _, aud := range processedAudio {
		if aud.Delay != 0 {
			mkvArgs = append(mkvArgs, "--sync", fmt.Sprintf("0:%d", aud.Delay))
		}
		mkvArgs = append(mkvArgs, "--language", fmt.Sprintf("0:%s", aud.Language))
		mkvArgs = append(mkvArgs, "--track-name", fmt.Sprintf("0:%s", aud.Title))
		mkvArgs = append(mkvArgs, aud.Path)
	}

	mkvArgs = append(mkvArgs, "--no-video")
	if len(remuxAudioTracks) > 0 {
		mkvArgs = append(mkvArgs, "--audio-tracks", strings.Join(remuxAudioTracks, ","))
	} else {
		mkvArgs = append(mkvArgs, "--no-audio")
	}
	mkvArgs = append(mkvArgs, absPath)

	_, err = RunCmd("mkvmerge", mkvArgs, false)
	if err != nil {
		LogInfo("mkvmerge assembly failed: %v", err)
		processingError = true
		return
	}

	LogInfo("Moving files to final destinations...")
	err = os.Rename(filePath, filepath.Join(DirOriginal, filepath.Base(filePath)))
	if err != nil {
		LogInfo("Failed to move original file: %v", err)
	}
	err = os.Rename(intermediateOutput, filepath.Join(DirCompleted, filepath.Base(filePath)))
	if err != nil {
		LogInfo("Failed to move completed file: %v", err)
	}

	LogInfo("Cleaning up persistent video temporary files...")
	temps := []string{
		baseName + ".vpy",
		baseName + ".ut.mkv",
		encodedVideoFile,
		baseName + ".ut.mkv.lwi",
		baseName + ".ut.mkv.ffindex",
	}
	if handbrakeCleanup != "" {
		temps = append(temps, handbrakeCleanup)
	}

	for _, t := range temps {
		if _, err := os.Stat(t); err == nil {
			LogInfo("    Deleting: %s", t)
			os.Remove(t)
		}
	}
}