#!/bin/bash
# Desc: Encodes video in parallel by processing time ranges directly from the input file.
# Usage: ./ffmpeg_parallel_encode.sh input_video_file segment_duration_in_seconds
# Version: 0.1.5

# Limit the number of parallel jobs if necessary
MAX_JOBS=8;   # Adjust this number based on your system's CPU capacity
MAX_MEM=512M; # Adjust this number based on your system's memory

set -e;  # Exit immediately if a command exits with a non-zero status
set -u;  # Treat unset variables as an error

yell() { echo "$0: $*" >&2; }  # Print script path and all args to stderr
die() { yell "$*"; exit 111; }  # Same as yell() but exits with status 111
must() { "$@" || die "cannot $*"; }  # Runs command, dies if it fails

cleanup() {
    if [[ -n "${temp_dir:-}" && -d "$temp_dir" ]]; then
        echo "Cleaning up temporary files...";
        rm -rf "$temp_dir";
    fi
    if [[ -n "${temp_list_file:-}" && -f "$temp_list_file" ]]; then
        rm -f "$temp_list_file";
    fi;
}; # Function to clean up temporary files

# Trap signals to ensure cleanup happens on script exit or interruption
trap cleanup EXIT INT TERM;

if [[ "$#" -ne 2 ]]; then
    yell "Usage: $0 input_video_file segment_duration_in_seconds";
    exit 1;
fi;

input_file="$1";
segment_time="$2";

# Check if input file exists
if [[ ! -f "$input_file" ]]; then
    die "Input file '$input_file' does not exist.";
fi;

base_name="$(basename "$input_file" | sed 's/\.[^.]*$//')";  # Strip extension
segment_prefix="${base_name}_segment";
encode_suffix="_encoded";
declare -a segment_filenames;

# Set umask to ensure files and directories are not accessible by others
umask 0077;

# Create a unique temporary directory for encoded files in the working directory
timestamp="$(date +%s)";
temp_dir="$(mktemp -d "${base_name}_${timestamp}_XXXXXX")";

# Safety check: Ensure temp_dir is not empty and is a directory
if [[ -z "${temp_dir}" || ! -d "${temp_dir}" ]]; then
    die "Temporary directory creation failed.";
fi;

yell "Temporary directory created at $temp_dir";

# Get the total duration of the input video in seconds
yell "Getting total duration of input video...";
duration_str="$(ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 "$input_file")";
if [[ -z "$duration_str" ]]; then
    die "Failed to get duration of input video.";
fi;

duration="$(printf "%.0f" "$duration_str")";  # Convert to integer seconds

yell "Total duration: $duration seconds";

# Calculate the number of segments
num_segments=$(( (duration + segment_time - 1) / segment_time ));  # Ceiling division

yell "Number of segments: $num_segments";

# Generate start times and segment filenames, write to a temporary file
yell "Preparing list of start times and filenames...";
temp_list_file="$(mktemp -p "$temp_dir")";

# Use a separator that is unlikely to appear in the data (ASCII Unit Separator)
separator=$'\x1F';

for (( i=0; i<num_segments; i++ )); do
    start_time=$(( i * segment_time ));
    segment_filename="${segment_prefix}_$(printf "%03d" "$i")${encode_suffix}.mp4";
    segment_filepath="${temp_dir}/${segment_filename}";
    printf '%s%s%s\n' "$start_time" "$separator" "$segment_filepath" >> "$temp_list_file";
    segment_filenames[i]="$segment_filename";  # Store for later use
done;

# Encode segments in parallel without writing unencoded segments to disk
yell "Encoding segments in parallel...";

# Use GNU Parallel with the custom separator
must parallel -j "$MAX_JOBS" --memsuspend "$MAX_MEM" --no-notice --bar --colsep "$separator" --arg-file "$temp_list_file" \
    ffmpeg -hide_banner -i "$input_file" -ss '{1}' -t "$segment_time" -c:v libx264 -preset fast -crf 23 -c:a copy '{2}';

# Step 3: Create a list file for concatenation
segments_file="${temp_dir}/segments.txt";
yell "Preparing for concatenation...";
for (( i=0; i<num_segments; i++ )); do
    segment_filename="${segment_filenames[i]}";
    echo "file '$segment_filename'" >> "$segments_file";
done;

# Safety check: Ensure segments_file is not empty
if [[ ! -s "$segments_file" ]]; then
    die "Segments file is empty.";
fi;

# Step 4: Concatenate the encoded segments into a single output file
output_file="${base_name}_output.mp4";
yell "Concatenating encoded segments into $output_file...";
(
    cd "$temp_dir" || die "Failed to change directory to $temp_dir";
    must ffmpeg -f concat -safe 0 -i "segments.txt" -c copy "$output_file";
    mv "$output_file" ..
)

yell "Process completed successfully in ${SECONDS} seconds. Output video is $output_file";

# Author: Steven Baltakatei Sandoval
# License: GPLv3+