jiggablend/internal/runner/tasks/render.go

package tasks

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"jiggablend/internal/runner/blender"
	"jiggablend/internal/runner/sandbox"
	"jiggablend/internal/runner/workspace"
	"jiggablend/pkg/scripts"
	"jiggablend/pkg/types"
)

// RenderProcessor handles render tasks.
type RenderProcessor struct{}

// NewRenderProcessor creates a new render processor.
func NewRenderProcessor() *RenderProcessor {
	return &RenderProcessor{}
}

// gpuErrorSubstrings are log line substrings that indicate a GPU backend error (matched case-insensitively); any match triggers full GPU lockout.
var gpuErrorSubstrings = []string{
	"illegal address in hip", // HIP (AMD) e.g. "Illegal address in HIP" or "Illegal address in hip"
	"memory access fault",    // AMDGPU page fault during HIP/HIPRT (not necessarily OOM)
	"page not present",       // AMDGPU GCVM page fault detail
	"hiperror",               // hipError* codes
	"hip error",
	"cuda error",
	"cuerror",
	"optix error",
	"oneapi error",
	"opencl error",
}

// checkGPUErrorLine checks a log line for GPU error indicators and triggers runner GPU lockout if found.
// Once lockout is already active for this attempt (or was active at context creation), further
// matching lines are ignored so we do not spam lockout callbacks/logs on multi-line fault dumps.
func (p *RenderProcessor) checkGPUErrorLine(ctx *Context, line string) {
	if ctx.GPULockedOut || ctx.GPULockoutArmedThisAttempt {
		return
	}
	lower := strings.ToLower(line)
	for _, sub := range gpuErrorSubstrings {
		if strings.Contains(lower, sub) {
			if ctx.OnGPUError != nil {
				ctx.OnGPUError()
			}
			return
		}
	}
}

// Process executes a render task.
func (p *RenderProcessor) Process(ctx *Context) error {
	if err := ctx.CheckCancelled(); err != nil {
		return err
	}

	if ctx.FrameEnd > ctx.Frame {
		ctx.Info(fmt.Sprintf("Starting task: job %d, frames %d-%d, format: %s",
			ctx.JobID, ctx.Frame, ctx.FrameEnd, ctx.GetOutputFormat()))
		log.Printf("Processing task %d: job %d, frames %d-%d", ctx.TaskID, ctx.JobID, ctx.Frame, ctx.FrameEnd)
	} else {
		ctx.Info(fmt.Sprintf("Starting task: job %d, frame %d, format: %s",
			ctx.JobID, ctx.Frame, ctx.GetOutputFormat()))
		log.Printf("Processing task %d: job %d, frame %d", ctx.TaskID, ctx.JobID, ctx.Frame)
	}

	// Find .blend file
	blendFile, err := workspace.FindFirstBlendFile(ctx.WorkDir)
	if err != nil {
		return fmt.Errorf("failed to find blend file: %w", err)
	}

	// Runners must use manager-provided Blender versions; never fall back to system blender.
	version := ctx.GetBlenderVersion()
	if version == "" {
		return fmt.Errorf("job metadata missing blender_version: runner cannot use system blender")
	}

	ctx.Info(fmt.Sprintf("Job requires Blender %s", version))
	binaryPath, err := ctx.Blender.GetBinaryPath(version)
	if err != nil {
		return fmt.Errorf("failed to get Blender %s from manager: %w", version, err)
	}

	blenderBinary, err := blender.ResolveBinaryPath(binaryPath)
	if err != nil {
		return fmt.Errorf("failed to resolve Blender %s binary: %w", version, err)
	}
	ctx.Info(fmt.Sprintf("Using Blender binary: %s", blenderBinary))

	// Create output directory
	outputDir := filepath.Join(ctx.WorkDir, "output")
	if err := os.MkdirAll(outputDir, 0755); err != nil {
		return fmt.Errorf("failed to create output directory: %w", err)
	}

	// Create home directory for Blender inside workspace
	blenderHome := filepath.Join(ctx.WorkDir, "home")
	if err := os.MkdirAll(blenderHome, 0755); err != nil {
		return fmt.Errorf("failed to create Blender home directory: %w", err)
	}

	// We always render EXR (linear) for VFX accuracy; job output_format is the deliverable (EXR sequence or video).
	renderFormat := "EXR"

	if ctx.ShouldForceCPU() {
		if ctx.ForceCPURendering {
			ctx.Info("Runner compatibility flag is enabled: forcing CPU rendering for this job")
		} else if ctx.GPUDetectionFailed {
			ctx.Info("GPU backend detection failed at startup—we could not determine available GPU backends, so rendering will use CPU to avoid compatibility issues")
		} else {
			ctx.Info("GPU lockout active: using CPU rendering only")
		}
	} else if ctx.ShouldDisableRT() {
		ctx.Info("GPU ray tracing acceleration disabled for this job (--disable-rt)")
	}
	if ctx.ShouldBatchGPUSamples() {
		total := ctx.GetCyclesSamples()
		ctx.Info(fmt.Sprintf(
			"gfx115x HIP sample batching: rendering %d samples in passes of %d (ROCm driver limit, not system RAM)",
			total, ctx.HipGPUSampleBatch,
		))
	}

	// Render
	if ctx.FrameEnd > ctx.Frame {
		ctx.Info(fmt.Sprintf("Starting Blender render for frames %d-%d...", ctx.Frame, ctx.FrameEnd))
	} else {
		ctx.Info(fmt.Sprintf("Starting Blender render for frame %d...", ctx.Frame))
	}
	if err := p.renderFrames(ctx, blenderBinary, blendFile, outputDir, renderFormat, blenderHome); err != nil {
		if errors.Is(err, ErrJobCancelled) {
			ctx.Warn("Render stopped because job was cancelled")
			return err
		}
		ctx.Error(fmt.Sprintf("Blender render failed: %v", err))
		return err
	}

	// Verify output (range or single frame)
	if err := p.verifyOutputRange(ctx, outputDir, renderFormat); err != nil {
		ctx.Error(fmt.Sprintf("Output verification failed: %v", err))
		return err
	}
	if ctx.FrameEnd > ctx.Frame {
		ctx.Info(fmt.Sprintf("Blender render completed for frames %d-%d", ctx.Frame, ctx.FrameEnd))
	} else {
		ctx.Info(fmt.Sprintf("Blender render completed for frame %d", ctx.Frame))
	}

	return nil
}

type renderPassOptions struct {
	samplesOverride *int
	seedOverride    *int
}

func (p *RenderProcessor) renderFrames(ctx *Context, blenderBinary, blendFile, outputDir, renderFormat, blenderHome string) error {
	if !ctx.ShouldBatchGPUSamples() {
		if err := p.createRenderScript(ctx, renderFormat, nil); err != nil {
			return err
		}
		return p.runBlender(ctx, blenderBinary, blendFile, outputDir, renderFormat, blenderHome)
	}

	totalSamples := ctx.GetCyclesSamples()
	batchSize := ctx.HipGPUSampleBatch
	baseSeed := ctx.GetCyclesSeed()
	batchDir := filepath.Join(ctx.WorkDir, "sample_batches")
	if err := os.MkdirAll(batchDir, 0755); err != nil {
		return fmt.Errorf("failed to create sample batch directory: %w", err)
	}

	var batchOutputs []string
	remaining := totalSamples
	batchNum := 0
	for remaining > 0 {
		if err := ctx.CheckCancelled(); err != nil {
			return err
		}
		passSamples := batchSize
		if remaining < passSamples {
			passSamples = remaining
		}
		seed := baseSeed + batchNum
		batchNum++
		ctx.Info(fmt.Sprintf("GPU sample batch %d: %d samples (seed %d, %d remaining)", batchNum, passSamples, seed, remaining-passSamples))

		passOutput := filepath.Join(batchDir, fmt.Sprintf("batch_%03d", batchNum))
		if err := os.MkdirAll(passOutput, 0755); err != nil {
			return err
		}
		opts := &renderPassOptions{
			samplesOverride: &passSamples,
			seedOverride:    &seed,
		}
		if err := p.createRenderScript(ctx, renderFormat, opts); err != nil {
			return err
		}
		if err := p.runBlender(ctx, blenderBinary, blendFile, passOutput, renderFormat, blenderHome); err != nil {
			return err
		}
		if err := p.verifyOutputRange(ctx, passOutput, renderFormat); err != nil {
			return err
		}
		batchOutputs = append(batchOutputs, p.firstFrameOutputPath(ctx, passOutput, renderFormat))
		remaining -= passSamples
	}

	if err := os.MkdirAll(outputDir, 0755); err != nil {
		return err
	}
	finalOutput := p.firstFrameOutputPath(ctx, outputDir, renderFormat)
	ctx.Info(fmt.Sprintf("Merging %d GPU sample batches into final EXR...", len(batchOutputs)))
	if err := mergeEXRFiles(finalOutput, batchOutputs); err != nil {
		return err
	}
	ctx.Info("GPU sample batch merge completed")
	return nil
}

func (p *RenderProcessor) firstFrameOutputPath(ctx *Context, outputDir, renderFormat string) string {
	ext := strings.ToLower(renderFormat)
	if ctx.FrameEnd > ctx.Frame {
		return filepath.Join(outputDir, fmt.Sprintf("frame_%04d.%s", ctx.Frame, ext))
	}
	return filepath.Join(outputDir, fmt.Sprintf("frame_%04d.%s", ctx.Frame, ext))
}

func (p *RenderProcessor) createRenderScript(ctx *Context, renderFormat string, opts *renderPassOptions) error {
	formatFilePath := filepath.Join(ctx.WorkDir, "output_format.txt")
	renderSettingsFilePath := filepath.Join(ctx.WorkDir, "render_settings.json")

	// Build unhide code conditionally
	unhideCode := ""
	if ctx.ShouldUnhideObjects() {
		unhideCode = scripts.UnhideObjects
	}

	// Load template and replace placeholders.
	// Job-bundled blender_addons/ always auto-install when present (users ship addons with scenes).
	// enable_execution only controls Blender --enable-autoexec (scripts inside .blend).
	scriptContent := scripts.RenderBlenderTemplate
	scriptContent = strings.ReplaceAll(scriptContent, "{{UNHIDE_CODE}}", unhideCode)
	scriptContent = strings.ReplaceAll(scriptContent, "{{FORMAT_FILE_PATH}}", fmt.Sprintf("%q", formatFilePath))
	scriptContent = strings.ReplaceAll(scriptContent, "{{RENDER_SETTINGS_FILE}}", fmt.Sprintf("%q", renderSettingsFilePath))

	scriptPath := filepath.Join(ctx.WorkDir, "enable_gpu.py")
	if err := os.WriteFile(scriptPath, []byte(scriptContent), 0644); err != nil {
		errMsg := fmt.Sprintf("failed to create GPU enable script: %v", err)
		ctx.Error(errMsg)
		return errors.New(errMsg)
	}

	// Write EXR to format file so Blender script sets OPEN_EXR (job output_format is for downstream deliverable only).
	ctx.Info("Writing output format 'EXR' to format file")
	if err := os.WriteFile(formatFilePath, []byte("EXR"), 0644); err != nil {
		errMsg := fmt.Sprintf("failed to create format file: %v", err)
		ctx.Error(errMsg)
		return errors.New(errMsg)
	}

	// Write render settings: merge job metadata with runner force_cpu (GPU lockout)
	var settingsMap map[string]interface{}
	if ctx.Metadata != nil && ctx.Metadata.RenderSettings.EngineSettings != nil {
		raw, err := json.Marshal(ctx.Metadata.RenderSettings)
		if err == nil {
			_ = json.Unmarshal(raw, &settingsMap)
		}
	}
	if settingsMap == nil {
		settingsMap = make(map[string]interface{})
	}
	settingsMap["force_cpu"] = ctx.ShouldForceCPU()
	settingsMap["disable_rt"] = ctx.ShouldDisableRT()
	if ctx.ShouldBatchGPUSamples() {
		// gfx115x: large tiles trigger AMDGPU page faults during HIP renders.
		settingsMap["tile_size_override"] = 256
		settingsMap["use_auto_tile_override"] = true
	}
	if opts != nil {
		if opts.samplesOverride != nil {
			settingsMap["samples_override"] = *opts.samplesOverride
		}
		if opts.seedOverride != nil {
			settingsMap["seed_override"] = *opts.seedOverride
		}
	}
	settingsJSON, err := json.Marshal(settingsMap)
	if err == nil {
		if err := os.WriteFile(renderSettingsFilePath, settingsJSON, 0644); err != nil {
			ctx.Warn(fmt.Sprintf("Failed to write render settings file: %v", err))
		}
	}

	return nil
}

// wrapBlenderCmd builds the Blender *exec.Cmd, optionally through the sandbox wrapper.
func wrapBlenderCmd(ctx *Context, blenderBinary string, args, env []string, blenderHome string) (*exec.Cmd, error) {
	sb := ctx.Sandbox
	if sb == nil {
		var err error
		sb, err = sandbox.New(sandbox.Options{Backend: sandbox.BackendNone})
		if err != nil {
			return nil, err
		}
	}
	if sb.Name() != sandbox.BackendNone {
		ctx.Info(fmt.Sprintf("Running Blender under sandbox backend %q", sb.Name()))
	}
	return sb.Wrap(sandbox.Spec{
		BlenderBinary: blenderBinary,
		Args:          args,
		WorkDir:       ctx.WorkDir,
		HomeDir:       blenderHome,
		Env:           env,
		HasAMD:        ctx.HasAMD,
		HasNVIDIA:     ctx.HasNVIDIA,
		HasIntel:      ctx.HasIntel,
		ForceCPU:      ctx.ShouldForceCPU(),
	})
}

func (p *RenderProcessor) runBlender(ctx *Context, blenderBinary, blendFile, outputDir, renderFormat, blenderHome string) error {
	scriptPath := filepath.Join(ctx.WorkDir, "enable_gpu.py")
	blendFileAbs, err := filepath.Abs(blendFile)
	if err != nil {
		return fmt.Errorf("failed to resolve blend file path: %w", err)
	}
	scriptPathAbs, err := filepath.Abs(scriptPath)
	if err != nil {
		return fmt.Errorf("failed to resolve blender script path: %w", err)
	}

	args := []string{"-b", blendFileAbs, "--python", scriptPathAbs}
	if ctx.ShouldEnableExecution() {
		args = append(args, "--enable-autoexec")
	}

	// Output pattern
	outputPattern := filepath.Join(outputDir, fmt.Sprintf("frame_####.%s", strings.ToLower(renderFormat)))
	outputAbsPattern, _ := filepath.Abs(outputPattern)
	args = append(args, "-o", outputAbsPattern)

	// Render single frame or range: -f N for one frame, -s start -e end -a for range
	if ctx.FrameEnd > ctx.Frame {
		args = append(args, "-s", fmt.Sprintf("%d", ctx.Frame), "-e", fmt.Sprintf("%d", ctx.FrameEnd), "-a")
	} else {
		args = append(args, "-f", fmt.Sprintf("%d", ctx.Frame))
	}

	// Set up environment: LD_LIBRARY_PATH for tarball Blender, then custom HOME
	env := os.Environ()
	env = blender.TarballEnv(blenderBinary, env)
	newEnv := make([]string, 0, len(env)+1)
	for _, e := range env {
		if !strings.HasPrefix(e, "HOME=") {
			newEnv = append(newEnv, e)
		}
	}
	newEnv = append(newEnv, fmt.Sprintf("HOME=%s", blenderHome))

	cmd, err := wrapBlenderCmd(ctx, blenderBinary, args, newEnv, blenderHome)
	if err != nil {
		return fmt.Errorf("failed to build sandboxed blender command: %w", err)
	}

	// 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 blender: %w", err)
	}
	stopMonitor := ctx.StartCancellationMonitor(cmd, "render")
	defer stopMonitor()

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

	// Stream stdout and watch for GPU error lines (lock out all GPU on any backend error)
	stdoutDone := make(chan bool)
	go func() {
		defer close(stdoutDone)
		scanner := bufio.NewScanner(stdoutPipe)
		for scanner.Scan() {
			line := scanner.Text()
			if line != "" {
				p.checkGPUErrorLine(ctx, line)
				shouldFilter, logLevel := blender.FilterLog(line)
				if !shouldFilter {
					ctx.Log(logLevel, line)
				}
			}
		}
		if err := scanner.Err(); err != nil {
			log.Printf("Error reading stdout: %v", err)
		}
	}()

	// Stream stderr and watch for GPU error lines
	stderrDone := make(chan bool)
	go func() {
		defer close(stderrDone)
		scanner := bufio.NewScanner(stderrPipe)
		for scanner.Scan() {
			line := scanner.Text()
			if line != "" {
				p.checkGPUErrorLine(ctx, line)
				shouldFilter, logLevel := blender.FilterLog(line)
				if !shouldFilter {
					if logLevel == types.LogLevelInfo {
						logLevel = types.LogLevelWarn
					}
					ctx.Log(logLevel, line)
				}
			}
		}
		if err := scanner.Err(); err != nil {
			log.Printf("Error reading stderr: %v", err)
		}
	}()

	// Wait for completion
	err = cmd.Wait()
	<-stdoutDone
	<-stderrDone

	if err != nil {
		if cancelled, checkErr := ctx.IsJobCancelled(); checkErr == nil && cancelled {
			return ErrJobCancelled
		}
		if exitErr, ok := err.(*exec.ExitError); ok {
			if exitErr.ExitCode() == 137 {
				return errors.New("Blender was killed due to excessive memory usage (OOM)")
			}
		}
		return fmt.Errorf("blender failed: %w", err)
	}

	return nil
}

// verifyOutputRange checks that output files exist for the task's frame range (first and last at minimum).
func (p *RenderProcessor) verifyOutputRange(ctx *Context, outputDir, renderFormat string) error {
	entries, err := os.ReadDir(outputDir)
	if err != nil {
		return fmt.Errorf("failed to read output directory: %w", err)
	}

	ctx.Info("Checking output directory for files...")
	ext := strings.ToLower(renderFormat)

	// Check first and last frame in range (minimum required for range; single frame = one check)
	framesToCheck := []int{ctx.Frame}
	if ctx.FrameEnd > ctx.Frame {
		framesToCheck = append(framesToCheck, ctx.FrameEnd)
	}
	for _, frame := range framesToCheck {
		found := false
		// Try frame_0001.ext, frame_1.ext, 0001.ext
		for _, name := range []string{
			fmt.Sprintf("frame_%04d.%s", frame, ext),
			fmt.Sprintf("frame_%d.%s", frame, ext),
			fmt.Sprintf("%04d.%s", frame, ext),
		} {
			if _, err := os.Stat(filepath.Join(outputDir, name)); err == nil {
				found = true
				ctx.Info(fmt.Sprintf("Found output file: %s", name))
				break
			}
		}
		if !found {
			// Search entries for this frame number
			frameStr := fmt.Sprintf("%d", frame)
			frameStrPadded := fmt.Sprintf("%04d", frame)
			for _, entry := range entries {
				if entry.IsDir() {
					continue
				}
				fileName := entry.Name()
				if strings.Contains(fileName, "%04d") || strings.Contains(fileName, "%d") {
					continue
				}
				if (strings.Contains(fileName, frameStrPadded) ||
					strings.Contains(fileName, frameStr)) && strings.HasSuffix(strings.ToLower(fileName), ext) {
					found = true
					ctx.Info(fmt.Sprintf("Found output file: %s", fileName))
					break
				}
			}
		}
		if !found {
			fileList := []string{}
			for _, e := range entries {
				if !e.IsDir() {
					fileList = append(fileList, e.Name())
				}
			}
			return fmt.Errorf("output file for frame %d not found; files in output directory: %v", frame, fileList)
		}
	}
	return nil
}