flatrender/services/node-agent/internal/runner/runner.go

// Package runner executes After Effects render jobs and streams progress back
// via the provided callbacks. When AE_PATH is empty, a mock render is used
// (useful for CI and dev environments without a licensed AE installation).
package runner

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"sync/atomic"
	"time"
)

// aerender prints one "PROGRESS:  <timecode> (<frame>): <n> Seconds" line per
// rendered frame, and (on most versions) an up-front line naming the frame range
// or total count. We parse the current frame for live progress + the total to
// turn it into a real percentage and ETA.
var (
	reFrameNum   = regexp.MustCompile(`\((\d+)\)\s*:`)            // "(59):"
	reTotalRange = regexp.MustCompile(`(?i)\bto\s+(\d+)\b`)       // "0 to 299"
	reTotalCount = regexp.MustCompile(`(?i)\b(\d+)\s+frames?\b`)  // "300 frames"
)

// ProgressFn is called periodically during rendering with (percent 0-100, message).
type ProgressFn func(ctx context.Context, percent int, message string) error

// PreviewFn is called each time a new preview frame is ready.
// The argument is a base64-encoded PNG. Errors are non-fatal.
type PreviewFn func(ctx context.Context, imageB64 string) error

// Job holds the parameters for a single render.
type Job struct {
	JobID          string
	SavedProjectID string
	Quality        string
	Resolution     string
	FrameRate      int
	HasMusic       bool
	HasVoiceover   bool
	// AEPFilePath is the local path to the downloaded .aep project file.
	// In a full implementation the agent downloads this from MinIO before calling Run.
	AEPFilePath string
	// CompName is the composition to render (-comp), e.g. "frfinal". When empty the
	// node renders the project's render queue (-rqindex 1) instead.
	CompName string
	// AfterFxPath is afterfx.exe (the AE app, scriptable) — used by the binder to write
	// input values into the project before rendering. Render itself uses aerender.exe.
	AfterFxPath string
	// Bindings are the user's edited input values to write into the .aep before render.
	Bindings []Binding
}

// Binding is one input value written into the AE project before render.
type Binding struct {
	Key   string // AE layer/footage name, e.g. frl_c1t1 / frl_c1m1
	Type  string // content element type (Text, Media, …)
	Value string // text content, or a media URL
}

// Run executes the render job, calling onProgress and onPreview as it advances.
// Returns the path to the output MP4 file on success.
func Run(ctx context.Context, aePath, workDir string, job *Job, onProgress ProgressFn, onPreview PreviewFn) (string, error) {
	outputDir := filepath.Join(workDir, "renders", job.JobID)
	if err := os.MkdirAll(outputDir, 0o755); err != nil {
		return "", fmt.Errorf("create output dir: %w", err)
	}
	outputPath := filepath.Join(outputDir, "output.mp4")

	// Mock render when AE isn't installed (aePath empty) OR when this job has no
	// template project to render (AEPFilePath empty — the template bundle wasn't
	// uploaded/promoted yet). Mock drives progress+preview to completion so the job
	// doesn't hard-fail; a real render requires both AE and a downloaded .aep.
	if aePath == "" || job.AEPFilePath == "" {
		if aePath != "" && job.AEPFilePath == "" {
			log.Printf("[job %s] no template .aep available — falling back to mock render", job.JobID)
		}
		return mockRender(ctx, job, outputPath, onProgress, onPreview)
	}

	// Render binder: write the user's edited input values into the project before
	// rendering so the MP4 reflects their text/media. Non-fatal — on failure we render
	// the template defaults rather than failing the job.
	if len(job.Bindings) > 0 && job.AfterFxPath != "" {
		_ = onProgress(ctx, 6, "Applying your edits…")
		bound, berr := RunBinder(ctx, job, workDir)
		if berr != nil {
			log.Printf("[job %s] binder failed (%v) — rendering template defaults", job.JobID, berr)
		} else {
			job.AEPFilePath = bound // render the bound project with the user's values
		}
	}

	return aeRender(ctx, aePath, job, outputPath, onProgress, onPreview)
}

// ── Mock render (no AE installed) ────────────────────────────────────────────

func mockRender(ctx context.Context, job *Job, outputPath string, onProgress ProgressFn, onPreview PreviewFn) (string, error) {
	log.Printf("[mock] starting render for job %s (%s %s %dfps)", job.JobID, job.Quality, job.Resolution, job.FrameRate)

	steps := []struct {
		pct int
		msg string
	}{
		{5, "Preparing project…"},
		{15, "Loading template…"},
		{30, "Rendering frames…"},
		{50, "Rendering frames… (50%)"},
		{70, "Rendering frames… (70%)"},
		{85, "Encoding MP4…"},
		{95, "Uploading output…"},
	}

	for _, s := range steps {
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case <-time.After(800 * time.Millisecond):
		}
		if err := onProgress(ctx, s.pct, s.msg); err != nil {
			log.Printf("[mock] progress callback error: %v", err)
		}
		// Generate and push a preview frame at each step
		if onPreview != nil {
			b64 := GeneratePreviewB64(s.pct, job.Quality, job.Resolution)
			if err := onPreview(ctx, b64); err != nil {
				log.Printf("[mock] preview callback error: %v", err)
			}
		}
		log.Printf("[mock] %d%% — %s", s.pct, s.msg)
	}

	// Write a placeholder file so the path is valid
	if err := os.WriteFile(outputPath, []byte("mock-render-output"), 0o644); err != nil {
		return "", fmt.Errorf("write mock output: %w", err)
	}
	log.Printf("[mock] render complete: %s", outputPath)
	return outputPath, nil
}

// findRenderedOutput locates the file aerender actually produced. The requested
// path is e.g. <dir>/output.mp4, but the output module may have written
// output.avi / output.mov / output.mp4. Prefer an exact match, then .mp4, then
// the largest output.* file in the directory.
func findRenderedOutput(requested string) string {
	if st, err := os.Stat(requested); err == nil && st.Size() > 0 {
		return requested
	}
	dir := filepath.Dir(requested)
	base := strings.TrimSuffix(filepath.Base(requested), filepath.Ext(requested)) // "output"
	matches, _ := filepath.Glob(filepath.Join(dir, base+".*"))
	var best string
	var bestSize int64 = -1
	for _, m := range matches {
		st, err := os.Stat(m)
		if err != nil || st.IsDir() {
			continue
		}
		// Prefer .mp4 immediately.
		if strings.EqualFold(filepath.Ext(m), ".mp4") && st.Size() > 0 {
			return m
		}
		if st.Size() > bestSize {
			best, bestSize = m, st.Size()
		}
	}
	return best
}

// ffmpegPath locates an ffmpeg binary: $FFMPEG_PATH, then `ffmpeg(.exe)` next to
// the agent executable, then PATH. Returns "" when none is found.
func ffmpegPath() string {
	if p := os.Getenv("FFMPEG_PATH"); p != "" {
		if _, err := os.Stat(p); err == nil {
			return p
		}
	}
	name := "ffmpeg"
	if runtime.GOOS == "windows" {
		name = "ffmpeg.exe"
	}
	if exe, err := os.Executable(); err == nil {
		cand := filepath.Join(filepath.Dir(exe), name)
		if _, err := os.Stat(cand); err == nil {
			return cand
		}
	}
	if p, err := exec.LookPath(name); err == nil {
		return p
	}
	return ""
}

// transcodeToMP4 converts a lossless AE render (AVI/MOV) to a web-playable H.264
// MP4 using ffmpeg. Returns the .mp4 path. Errors if ffmpeg is unavailable.
// resolutionHeight maps a quality-tier label to its output height (mirrors render-svc).
func resolutionHeight(resolution string) int {
	switch strings.ToLower(strings.TrimSpace(resolution)) {
	case "360p":
		return 360
	case "540p":
		return 540
	case "720p":
		return 720
	case "1080p", "fullhd":
		return 1080
	case "4k", "2160p":
		return 2160
	default:
		return 0 // unknown → no scaling
	}
}

func transcodeToMP4(ctx context.Context, src, requested string, height int) (string, error) {
	ff := ffmpegPath()
	if ff == "" {
		return "", fmt.Errorf("ffmpeg not found (set FFMPEG_PATH or place ffmpeg.exe next to the agent)")
	}
	dst := strings.TrimSuffix(requested, filepath.Ext(requested)) + ".mp4"
	args := []string{
		"-y", "-i", src,
		"-c:v", "libx264", "-preset", "medium", "-crf", "20", "-pix_fmt", "yuv420p",
	}
	// Downscale (or up) to the selected quality tier. -2 keeps width even & aspect.
	if height > 0 {
		args = append(args, "-vf", fmt.Sprintf("scale=-2:%d", height))
	}
	args = append(args,
		"-c:a", "aac", "-b:a", "192k",
		"-movflags", "+faststart",
		dst,
	)
	log.Printf("[ffmpeg] %s %v", ff, args)
	cmd := exec.CommandContext(ctx, ff, args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return "", fmt.Errorf("ffmpeg: %w", err)
	}
	if st, err := os.Stat(dst); err != nil || st.Size() == 0 {
		return "", fmt.Errorf("ffmpeg produced no output")
	}
	return dst, nil
}

// ── Real AE render via aerender.exe ──────────────────────────────────────────

func aeRender(ctx context.Context, aePath string, job *Job, outputPath string, onProgress ProgressFn, onPreview PreviewFn) (string, error) {
	if job.AEPFilePath == "" {
		return "", fmt.Errorf("AEPFilePath is required for real AE render")
	}

	// aerender flags:
	//   -project <path.aep>
	//   -comp <name>  (or -rqindex 1 when no comp name is known)
	//   -output  <output.mp4>
	// Modern AE can't reliably write H.264 directly from aerender, so we let it
	// render with the project's output module (typically Lossless AVI/MOV) and
	// transcode to MP4 with ffmpeg afterwards (see transcodeToMP4).
	// Without -comp/-rqindex, aerender ignores -output and renders nothing.
	args := []string{"-project", job.AEPFilePath}
	if job.CompName != "" {
		args = append(args, "-comp", job.CompName)
	} else {
		args = append(args, "-rqindex", "1")
	}
	args = append(args, "-output", outputPath)

	log.Printf("[ae] running: %s %v", aePath, args)
	cmd := exec.CommandContext(ctx, aePath, args...)
	// Run from the project's folder so a .zip bundle's relative footage/font paths
	// resolve correctly (the .aep sits alongside its assets after extraction).
	cmd.Dir = filepath.Dir(job.AEPFilePath)
	stdout, pipeErr := cmd.StdoutPipe()
	if pipeErr != nil {
		return "", fmt.Errorf("stdout pipe: %w", pipeErr)
	}
	cmd.Stderr = os.Stderr

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

	// Scan aerender stdout in the background: echo it to our log AND extract the
	// current frame + total frame count for real progress.
	var curFrame, totalFrames int64
	go func() {
		sc := bufio.NewScanner(stdout)
		sc.Buffer(make([]byte, 64*1024), 1024*1024)
		for sc.Scan() {
			line := sc.Text()
			_, _ = io.WriteString(os.Stdout, line+"\n")
			if m := reFrameNum.FindStringSubmatch(line); m != nil {
				if n, e := strconv.ParseInt(m[1], 10, 64); e == nil {
					atomic.StoreInt64(&curFrame, n)
				}
			}
			if atomic.LoadInt64(&totalFrames) == 0 {
				if m := reTotalRange.FindStringSubmatch(line); m != nil {
					if n, e := strconv.ParseInt(m[1], 10, 64); e == nil && n > 1 {
						atomic.StoreInt64(&totalFrames, n)
					}
				} else if m := reTotalCount.FindStringSubmatch(line); m != nil {
					if n, e := strconv.ParseInt(m[1], 10, 64); e == nil && n > 1 {
						atomic.StoreInt64(&totalFrames, n)
					}
				}
			}
		}
	}()

	done := make(chan error, 1)
	go func() { done <- cmd.Wait() }()

	_ = onProgress(ctx, 5, "After Effects starting…")
	start := time.Now()
	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()
	lastPreview := time.Time{}

	for {
		select {
		case err := <-done:
			if err != nil {
				return "", fmt.Errorf("aerender exit: %w", err)
			}
			actual := findRenderedOutput(outputPath)
			if actual == "" {
				return "", fmt.Errorf("aerender finished but no output file found in %s", filepath.Dir(outputPath))
			}
			if strings.EqualFold(filepath.Ext(actual), ".mp4") {
				_ = onProgress(ctx, 98, "Encoding complete")
				return actual, nil
			}
			_ = onProgress(ctx, 92, "Transcoding to MP4…")
			mp4, terr := transcodeToMP4(ctx, actual, outputPath, resolutionHeight(job.Resolution))
			if terr != nil {
				log.Printf("[ae] transcode failed (%v) — uploading raw %s", terr, filepath.Ext(actual))
				return actual, nil
			}
			_ = onProgress(ctx, 98, "Encoding complete")
			_ = os.Remove(actual) // drop the multi-GB intermediate
			return mp4, nil
		case <-ticker.C:
			cur := atomic.LoadInt64(&curFrame)
			tot := atomic.LoadInt64(&totalFrames)
			pct, msg := aeProgress(cur, tot, time.Since(start))
			_ = onProgress(ctx, pct, msg)
			// Preview ~every 8s so the box shows something soon after start.
			if onPreview != nil && time.Since(lastPreview) >= 8*time.Second {
				lastPreview = time.Now()
				if err := onPreview(ctx, GeneratePreviewB64(pct, job.Quality, job.Resolution)); err != nil {
					log.Printf("[ae] preview push error: %v", err)
				}
			}
		case <-ctx.Done():
			_ = cmd.Process.Kill()
			return "", ctx.Err()
		}
	}
}

// aeProgress turns the current/total frame counts (and elapsed time) into a
// render percentage (5–90, leaving headroom for transcode/upload) plus a human
// message. When the total is known the percentage is real; otherwise it eases
// toward 88% over time so the bar keeps moving without ever sticking or lying
// about being done.
func aeProgress(cur, total int64, elapsed time.Duration) (int, string) {
	if total > 1 && cur >= 0 {
		frac := float64(cur) / float64(total)
		if frac > 1 {
			frac = 1
		}
		pct := 5 + int(frac*85) // 5..90
		return pct, fmt.Sprintf("در حال رندر… فریم %d از %d", cur, total)
	}
	// Unknown total: asymptotic ease toward 88% (~half-way by ~90s).
	secs := elapsed.Seconds()
	pct := 5 + int(83*(secs/(secs+90)))
	if pct > 88 {
		pct = 88
	}
	if cur > 0 {
		return pct, fmt.Sprintf("در حال رندر… فریم %d", cur)
	}
	return pct, "در حال رندر…"
}