jiggablend/internal/runner/tasks/encode.go

package tasks

import (
	"bufio"
	"errors"
	"fmt"
	"log"
	"math"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"sort"
	"strings"

	"jiggablend/internal/runner/encoding"
)

// EncodeProcessor handles encode tasks.
type EncodeProcessor struct{}

// NewEncodeProcessor creates a new encode processor.
func NewEncodeProcessor() *EncodeProcessor {
	return &EncodeProcessor{}
}

// Process executes an encode task.
func (p *EncodeProcessor) Process(ctx *Context) error {
	ctx.Info(fmt.Sprintf("Starting encode task: job %d", ctx.JobID))
	log.Printf("Processing encode task %d for job %d", ctx.TaskID, ctx.JobID)

	// Create temporary work directory
	workDir, err := ctx.Workspace.CreateVideoDir(ctx.JobID)
	if err != nil {
		return fmt.Errorf("failed to create work directory: %w", err)
	}
	defer func() {
		if err := ctx.Workspace.CleanupVideoDir(ctx.JobID); err != nil {
			log.Printf("Warning: Failed to cleanup encode work directory: %v", err)
		}
	}()

	// Get output format and frame rate
	outputFormat := ctx.GetOutputFormat()
	if outputFormat == "" {
		outputFormat = "EXR_264_MP4"
	}
	frameRate := ctx.GetFrameRate()

	ctx.Info(fmt.Sprintf("Encode: detected output format '%s'", outputFormat))
	ctx.Info(fmt.Sprintf("Encode: using frame rate %.2f fps", frameRate))

	// Get job files
	files, err := ctx.Manager.GetJobFiles(ctx.JobID)
	if err != nil {
		ctx.Error(fmt.Sprintf("Failed to get job files: %v", err))
		return fmt.Errorf("failed to get job files: %w", err)
	}

	ctx.Info(fmt.Sprintf("GetJobFiles returned %d total files for job %d", len(files), ctx.JobID))

	// Log all files for debugging
	for _, file := range files {
		ctx.Info(fmt.Sprintf("File: %s (type: %s, size: %d)", file.FileName, file.FileType, file.FileSize))
	}

	// Determine source format based on output format
	sourceFormat := "exr"
	fileExt := ".exr"

	// Find and deduplicate frame files (EXR or PNG)
	frameFileSet := make(map[string]bool)
	var frameFilesList []string
	for _, file := range files {
		if file.FileType == "output" && strings.HasSuffix(strings.ToLower(file.FileName), fileExt) {
			// Deduplicate by filename
			if !frameFileSet[file.FileName] {
				frameFileSet[file.FileName] = true
				frameFilesList = append(frameFilesList, file.FileName)
			}
		}
	}

	if len(frameFilesList) == 0 {
		// Log why no files matched (deduplicate for error reporting)
		outputFileSet := make(map[string]bool)
		frameFilesOtherTypeSet := make(map[string]bool)
		var outputFiles []string
		var frameFilesOtherType []string

		for _, file := range files {
			if file.FileType == "output" {
				if !outputFileSet[file.FileName] {
					outputFileSet[file.FileName] = true
					outputFiles = append(outputFiles, file.FileName)
				}
			}
			if strings.HasSuffix(strings.ToLower(file.FileName), fileExt) {
				key := fmt.Sprintf("%s (type: %s)", file.FileName, file.FileType)
				if !frameFilesOtherTypeSet[key] {
					frameFilesOtherTypeSet[key] = true
					frameFilesOtherType = append(frameFilesOtherType, key)
				}
			}
		}
		ctx.Error(fmt.Sprintf("no %s frame files found for encode: found %d total files, %d unique output files, %d unique %s files (with other types)", strings.ToUpper(fileExt[1:]), len(files), len(outputFiles), len(frameFilesOtherType), strings.ToUpper(fileExt[1:])))
		if len(outputFiles) > 0 {
			ctx.Error(fmt.Sprintf("Output files found: %v", outputFiles))
		}
		if len(frameFilesOtherType) > 0 {
			ctx.Error(fmt.Sprintf("%s files with wrong type: %v", strings.ToUpper(fileExt[1:]), frameFilesOtherType))
		}
		err := fmt.Errorf("no %s frame files found for encode", strings.ToUpper(fileExt[1:]))
		return err
	}

	ctx.Info(fmt.Sprintf("Found %d %s frames for encode", len(frameFilesList), strings.ToUpper(fileExt[1:])))

	// Download frames
	ctx.Info(fmt.Sprintf("Downloading %d %s frames for encode...", len(frameFilesList), strings.ToUpper(fileExt[1:])))

	var frameFiles []string
	for i, fileName := range frameFilesList {
		ctx.Info(fmt.Sprintf("Downloading frame %d/%d: %s", i+1, len(frameFilesList), fileName))
		framePath := filepath.Join(workDir, fileName)
		if err := ctx.Manager.DownloadFrame(ctx.JobID, fileName, framePath); err != nil {
			ctx.Error(fmt.Sprintf("Failed to download %s frame %s: %v", strings.ToUpper(fileExt[1:]), fileName, err))
			log.Printf("Failed to download %s frame for encode %s: %v", strings.ToUpper(fileExt[1:]), fileName, err)
			continue
		}
		ctx.Info(fmt.Sprintf("Successfully downloaded frame %d/%d: %s", i+1, len(frameFilesList), fileName))
		frameFiles = append(frameFiles, framePath)
	}

	if len(frameFiles) == 0 {
		err := fmt.Errorf("failed to download any %s frames for encode", strings.ToUpper(fileExt[1:]))
		ctx.Error(err.Error())
		return err
	}

	sort.Strings(frameFiles)
	ctx.Info(fmt.Sprintf("Downloaded %d frames", len(frameFiles)))

	// Check if EXR files have alpha channel and HDR content (only for EXR source format)
	hasAlpha := false
	hasHDR := false
	if sourceFormat == "exr" {
		// Check first frame for alpha channel and HDR using ffprobe
		firstFrame := frameFiles[0]
		hasAlpha = detectAlphaChannel(ctx, firstFrame)
		if hasAlpha {
			ctx.Info("Detected alpha channel in EXR files")
		} else {
			ctx.Info("No alpha channel detected in EXR files")
		}

		hasHDR = detectHDR(ctx, firstFrame)
		if hasHDR {
			ctx.Info("Detected HDR content in EXR files")
		} else {
			ctx.Info("No HDR content detected in EXR files (SDR range)")
		}
	}

	// Generate video
	// Use alpha if:
	// 1. User explicitly enabled it OR source has alpha channel AND
	// 2. Codec supports alpha (AV1 or VP9)
	preserveAlpha := ctx.ShouldPreserveAlpha()
	useAlpha := (preserveAlpha || hasAlpha) && (outputFormat == "EXR_AV1_MP4" || outputFormat == "EXR_VP9_WEBM")
	if (preserveAlpha || hasAlpha) && outputFormat == "EXR_264_MP4" {
		ctx.Warn("Alpha channel requested/detected but H.264 does not support alpha. Consider using EXR_AV1_MP4 or EXR_VP9_WEBM to preserve alpha.")
	}
	if preserveAlpha && !hasAlpha {
		ctx.Warn("Alpha preservation requested but no alpha channel detected in EXR files.")
	}
	if useAlpha {
		if preserveAlpha && hasAlpha {
			ctx.Info("Alpha preservation enabled: Using alpha channel encoding")
		} else if hasAlpha {
			ctx.Info("Alpha channel detected - automatically enabling alpha encoding")
		}
	}
	var outputExt string
	switch outputFormat {
	case "EXR_VP9_WEBM":
		outputExt = "webm"
		ctx.Info("Encoding WebM video with VP9 codec (with alpha channel and HDR support)...")
	case "EXR_AV1_MP4":
		outputExt = "mp4"
		ctx.Info("Encoding MP4 video with AV1 codec (with alpha channel)...")
	default:
		outputExt = "mp4"
		ctx.Info("Encoding MP4 video with H.264 codec...")
	}

	outputVideo := filepath.Join(workDir, fmt.Sprintf("output_%d.%s", ctx.JobID, outputExt))

	// Build input pattern
	firstFrame := frameFiles[0]
	baseName := filepath.Base(firstFrame)
	re := regexp.MustCompile(`_(\d+)\.`)
	var pattern string
	var startNumber int
	frameNumStr := re.FindStringSubmatch(baseName)
	if len(frameNumStr) > 1 {
		pattern = re.ReplaceAllString(baseName, "_%04d.")
		fmt.Sscanf(frameNumStr[1], "%d", &startNumber)
	} else {
		startNumber = extractFrameNumber(baseName)
		pattern = strings.Replace(baseName, fmt.Sprintf("%d", startNumber), "%04d", 1)
	}
	patternPath := filepath.Join(workDir, pattern)

	// Select encoder and build command (software encoding only)
	var encoder encoding.Encoder
	switch outputFormat {
	case "EXR_AV1_MP4":
		encoder = ctx.Encoder.SelectAV1()
	case "EXR_VP9_WEBM":
		encoder = ctx.Encoder.SelectVP9()
	default:
		encoder = ctx.Encoder.SelectH264()
	}

	ctx.Info(fmt.Sprintf("Using encoder: %s (%s)", encoder.Name(), encoder.Codec()))

	// All software encoders use 2-pass for optimal quality
	ctx.Info("Starting 2-pass encode for optimal quality...")

	// Pass 1
	ctx.Info("Pass 1/2: Analyzing content for optimal encode...")
	softEncoder := encoder.(*encoding.SoftwareEncoder)
	// Use HDR if: user explicitly enabled it OR HDR content was detected
	preserveHDR := (ctx.ShouldPreserveHDR() || hasHDR) && sourceFormat == "exr"
	if hasHDR && !ctx.ShouldPreserveHDR() {
		ctx.Info("HDR content detected - automatically enabling HDR preservation")
	}
	pass1Cmd := softEncoder.BuildPass1Command(&encoding.EncodeConfig{
		InputPattern: patternPath,
		OutputPath:   outputVideo,
		StartFrame:   startNumber,
		FrameRate:    frameRate,
		WorkDir:      workDir,
		UseAlpha:     useAlpha,
		TwoPass:      true,
		SourceFormat: sourceFormat,
		PreserveHDR:  preserveHDR,
	})
	if err := pass1Cmd.Run(); err != nil {
		ctx.Warn(fmt.Sprintf("Pass 1 completed (warnings expected): %v", err))
	}

	// Pass 2
	ctx.Info("Pass 2/2: Encoding with optimal quality...")

	preserveHDR = (ctx.ShouldPreserveHDR() || hasHDR) && sourceFormat == "exr"
	if preserveHDR {
		if hasHDR && !ctx.ShouldPreserveHDR() {
			ctx.Info("HDR preservation enabled (auto-detected): Using HLG transfer with bt709 primaries")
		} else {
			ctx.Info("HDR preservation enabled: Using HLG transfer with bt709 primaries")
		}
	}

	config := &encoding.EncodeConfig{
		InputPattern: patternPath,
		OutputPath:   outputVideo,
		StartFrame:   startNumber,
		FrameRate:    frameRate,
		WorkDir:      workDir,
		UseAlpha:     useAlpha,
		TwoPass:      true, // Software encoding always uses 2-pass for quality
		SourceFormat: sourceFormat,
		PreserveHDR:  preserveHDR,
	}

	cmd := encoder.BuildCommand(config)
	if cmd == nil {
		return errors.New("failed to build encode command")
	}

	// Set up pipes
	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("failed to create stdout pipe: %w", err)
	}

	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("failed to create stderr pipe: %w", err)
	}

	if err := cmd.Start(); err != nil {
		return fmt.Errorf("failed to start encode command: %w", err)
	}

	ctx.Processes.Track(ctx.TaskID, cmd)
	defer ctx.Processes.Untrack(ctx.TaskID)

	// Stream stdout
	stdoutDone := make(chan bool)
	go func() {
		defer close(stdoutDone)
		scanner := bufio.NewScanner(stdoutPipe)
		for scanner.Scan() {
			line := scanner.Text()
			if line != "" {
				ctx.Info(line)
			}
		}
	}()

	// Stream stderr
	stderrDone := make(chan bool)
	go func() {
		defer close(stderrDone)
		scanner := bufio.NewScanner(stderrPipe)
		for scanner.Scan() {
			line := scanner.Text()
			if line != "" {
				ctx.Warn(line)
			}
		}
	}()

	err = cmd.Wait()
	<-stdoutDone
	<-stderrDone

	if err != nil {
		var errMsg string
		if exitErr, ok := err.(*exec.ExitError); ok {
			if exitErr.ExitCode() == 137 {
				errMsg = "FFmpeg was killed due to excessive memory usage (OOM)"
			} else {
				errMsg = fmt.Sprintf("ffmpeg encoding failed: %v", err)
			}
		} else {
			errMsg = fmt.Sprintf("ffmpeg encoding failed: %v", err)
		}

		if sizeErr := checkFFmpegSizeError(errMsg); sizeErr != nil {
			ctx.Error(sizeErr.Error())
			return sizeErr
		}

		ctx.Error(errMsg)
		return errors.New(errMsg)
	}

	// Verify output
	if _, err := os.Stat(outputVideo); os.IsNotExist(err) {
		err := fmt.Errorf("video %s file not created: %s", outputExt, outputVideo)
		ctx.Error(err.Error())
		return err
	}

	// Clean up 2-pass log files
	os.Remove(filepath.Join(workDir, "ffmpeg2pass-0.log"))
	os.Remove(filepath.Join(workDir, "ffmpeg2pass-0.log.mbtree"))

	ctx.Info(fmt.Sprintf("%s video encoded successfully", strings.ToUpper(outputExt)))

	// Upload video
	ctx.Info(fmt.Sprintf("Uploading encoded %s video...", strings.ToUpper(outputExt)))

	uploadPath := fmt.Sprintf("/api/runner/jobs/%d/upload", ctx.JobID)
	if err := ctx.Manager.UploadFile(uploadPath, ctx.JobToken, outputVideo); err != nil {
		ctx.Error(fmt.Sprintf("Failed to upload %s: %v", strings.ToUpper(outputExt), err))
		return fmt.Errorf("failed to upload %s: %w", strings.ToUpper(outputExt), err)
	}

	ctx.Info(fmt.Sprintf("Successfully uploaded %s: %s", strings.ToUpper(outputExt), filepath.Base(outputVideo)))

	log.Printf("Successfully generated and uploaded %s for job %d: %s", strings.ToUpper(outputExt), ctx.JobID, filepath.Base(outputVideo))
	return nil
}

// detectAlphaChannel checks if an EXR file has an alpha channel using ffprobe
func detectAlphaChannel(ctx *Context, filePath string) bool {
	// Use ffprobe to check pixel format and stream properties
	// EXR files with alpha will have formats like gbrapf32le (RGBA) vs gbrpf32le (RGB)
	cmd := exec.Command("ffprobe",
		"-v", "error",
		"-select_streams", "v:0",
		"-show_entries", "stream=pix_fmt:stream=codec_name",
		"-of", "default=noprint_wrappers=1",
		filePath,
	)

	output, err := cmd.Output()
	if err != nil {
		// If ffprobe fails, assume no alpha (conservative approach)
		ctx.Warn(fmt.Sprintf("Failed to detect alpha channel in %s: %v", filepath.Base(filePath), err))
		return false
	}

	outputStr := string(output)
	// Check pixel format - EXR with alpha typically has 'a' in the format name (e.g., gbrapf32le)
	// Also check for formats that explicitly indicate alpha
	hasAlpha := strings.Contains(outputStr, "pix_fmt=gbrap") ||
		strings.Contains(outputStr, "pix_fmt=rgba") ||
		strings.Contains(outputStr, "pix_fmt=yuva") ||
		strings.Contains(outputStr, "pix_fmt=abgr")

	if hasAlpha {
		ctx.Info(fmt.Sprintf("Detected alpha channel in EXR file: %s", filepath.Base(filePath)))
	}

	return hasAlpha
}

// detectHDR checks if an EXR file contains HDR content using ffprobe
func detectHDR(ctx *Context, filePath string) bool {
	// First, check if the pixel format supports HDR (32-bit float)
	cmd := exec.Command("ffprobe",
		"-v", "error",
		"-select_streams", "v:0",
		"-show_entries", "stream=pix_fmt",
		"-of", "default=noprint_wrappers=1:nokey=1",
		filePath,
	)

	output, err := cmd.Output()
	if err != nil {
		// If ffprobe fails, assume no HDR (conservative approach)
		ctx.Warn(fmt.Sprintf("Failed to detect HDR in %s: %v", filepath.Base(filePath), err))
		return false
	}

	pixFmt := strings.TrimSpace(string(output))
	// EXR files with 32-bit float format (gbrpf32le, gbrapf32le) can contain HDR
	// Check if it's a 32-bit float format
	isFloat32 := strings.Contains(pixFmt, "f32") || strings.Contains(pixFmt, "f32le")

	if !isFloat32 {
		// Not a float format, definitely not HDR
		return false
	}

	// For 32-bit float EXR, sample pixels to check if values exceed SDR range (> 1.0)
	// Use ffmpeg to extract pixel statistics - check max pixel values
	// This is more efficient than sampling individual pixels
	cmd = exec.Command("ffmpeg",
		"-v", "error",
		"-i", filePath,
		"-vf", "signalstats",
		"-f", "null",
		"-",
	)

	output, err = cmd.CombinedOutput()
	if err != nil {
		// If stats extraction fails, try sampling a few pixels directly
		return detectHDRBySampling(ctx, filePath)
	}

	// Check output for max pixel values
	outputStr := string(output)
	// Look for max values in the signalstats output
	// If we find values > 1.0, it's HDR
	if strings.Contains(outputStr, "MAX") {
		// Try to extract max values from signalstats output
		// Format is typically like: YMAX:1.234 UMAX:0.567 VMAX:0.890
		// For EXR (RGB), we need to check R, G, B channels
		// Since signalstats works on YUV, we'll use a different approach
		return detectHDRBySampling(ctx, filePath)
	}

	// Fallback to pixel sampling
	return detectHDRBySampling(ctx, filePath)
}

// detectHDRBySampling samples pixels from multiple regions to detect HDR content
func detectHDRBySampling(ctx *Context, filePath string) bool {
	// Sample multiple 10x10 regions from different parts of the image
	// This gives us better coverage than a single sample
	sampleRegions := []string{
		"crop=10:10:iw/4:ih/4",     // Top-left quadrant
		"crop=10:10:iw*3/4:ih/4",   // Top-right quadrant
		"crop=10:10:iw/4:ih*3/4",   // Bottom-left quadrant
		"crop=10:10:iw*3/4:ih*3/4", // Bottom-right quadrant
		"crop=10:10:iw/2:ih/2",     // Center
	}

	for _, region := range sampleRegions {
		cmd := exec.Command("ffmpeg",
			"-v", "error",
			"-i", filePath,
			"-vf", fmt.Sprintf("%s,scale=1:1", region),
			"-f", "rawvideo",
			"-pix_fmt", "gbrpf32le",
			"-",
		)

		output, err := cmd.Output()
		if err != nil {
			continue // Skip this region if sampling fails
		}

		// Parse the float32 values (4 bytes per float, 3 channels RGB)
		if len(output) >= 12 { // At least 3 floats (RGB) = 12 bytes
			for i := 0; i < len(output)-11; i += 12 {
				// Read RGB values (little-endian float32)
				r := float32FromBytes(output[i : i+4])
				g := float32FromBytes(output[i+4 : i+8])
				b := float32FromBytes(output[i+8 : i+12])

				// Check if any channel exceeds 1.0 (SDR range)
				if r > 1.0 || g > 1.0 || b > 1.0 {
					maxVal := max(r, max(g, b))
					ctx.Info(fmt.Sprintf("Detected HDR content in EXR file: %s (max value: %.2f)", filepath.Base(filePath), maxVal))
					return true
				}
			}
		}
	}

	// If we sampled multiple regions and none exceed 1.0, it's likely SDR content
	// But since it's 32-bit float format, user can still manually enable HDR if needed
	return false
}

// float32FromBytes converts 4 bytes (little-endian) to float32
func float32FromBytes(bytes []byte) float32 {
	if len(bytes) < 4 {
		return 0
	}
	bits := uint32(bytes[0]) | uint32(bytes[1])<<8 | uint32(bytes[2])<<16 | uint32(bytes[3])<<24
	return math.Float32frombits(bits)
}

// max returns the maximum of two float32 values
func max(a, b float32) float32 {
	if a > b {
		return a
	}
	return b
}

func extractFrameNumber(filename string) int {
	parts := strings.Split(filepath.Base(filename), "_")
	if len(parts) < 2 {
		return 0
	}
	framePart := strings.Split(parts[1], ".")[0]
	var frameNum int
	fmt.Sscanf(framePart, "%d", &frameNum)
	return frameNum
}

func checkFFmpegSizeError(output string) error {
	outputLower := strings.ToLower(output)

	if strings.Contains(outputLower, "hardware does not support encoding at size") {
		constraintsMatch := regexp.MustCompile(`constraints:\s*width\s+(\d+)-(\d+)\s+height\s+(\d+)-(\d+)`).FindStringSubmatch(output)
		if len(constraintsMatch) == 5 {
			return fmt.Errorf("video frame size is outside hardware encoder limits. Hardware requires: width %s-%s, height %s-%s",
				constraintsMatch[1], constraintsMatch[2], constraintsMatch[3], constraintsMatch[4])
		}
		return fmt.Errorf("video frame size is outside hardware encoder limits")
	}

	if strings.Contains(outputLower, "picture size") && strings.Contains(outputLower, "is invalid") {
		sizeMatch := regexp.MustCompile(`picture size\s+(\d+)x(\d+)`).FindStringSubmatch(output)
		if len(sizeMatch) == 3 {
			return fmt.Errorf("invalid video frame size: %sx%s", sizeMatch[1], sizeMatch[2])
		}
		return fmt.Errorf("invalid video frame size")
	}

	if strings.Contains(outputLower, "error while opening encoder") &&
		(strings.Contains(outputLower, "width") || strings.Contains(outputLower, "height") || strings.Contains(outputLower, "size")) {
		sizeMatch := regexp.MustCompile(`at size\s+(\d+)x(\d+)`).FindStringSubmatch(output)
		if len(sizeMatch) == 3 {
			return fmt.Errorf("hardware encoder cannot encode frame size %sx%s", sizeMatch[1], sizeMatch[2])
		}
		return fmt.Errorf("hardware encoder error: frame size may be invalid")
	}

	if strings.Contains(outputLower, "invalid") &&
		(strings.Contains(outputLower, "width") || strings.Contains(outputLower, "height") || strings.Contains(outputLower, "dimension")) {
		return fmt.Errorf("invalid frame dimensions detected")
	}

	return nil
}