#!/bin/bash

# Exit immediately if a command exits with a non-zero status.
set -e

# Enable debugging (optional)
# Uncomment the following line to enable debug mode
# set -x

# Function to display informational messages in blue
echo_info() {
    echo -e "\033[1;34m[INFO]\033[0m $1"
}

# Function to display error messages in red
echo_error() {
    echo -e "\033[1;31m[ERROR]\033[0m $1"
}

# Function to display usage instructions
usage() {
    echo "Usage: $0 <segments_directory>"
    echo "Example: $0 606-short_segments"
    exit 1
}

# Check if exactly one argument is provided
if [ "$#" -ne 1 ]; then
    echo_error "Incorrect number of arguments."

    usage
fi

# Assign the first argument to SEGMENTS_DIR

SEGMENTS_DIR="$1"

# Resolve the absolute path of SEGMENTS_DIR
if command -v realpath &> /dev/null; then

    SEGMENTS_DIR=$(realpath "$SEGMENTS_DIR")
else
    # If realpath is not available, use a fallback
    SEGMENTS_DIR="$(cd "$SEGMENTS_DIR" && pwd)"
fi

# Check if the provided argument is a directory
if [ ! -d "$SEGMENTS_DIR" ]; then
    echo_error "The provided path '$SEGMENTS_DIR' is not a directory or does not exist."
    exit 1

fi


echo_info "Segments directory set to: $SEGMENTS_DIR"

# Check if ffmpeg is installed
if ! command -v ffmpeg &> /dev/null; then
    echo_error "ffmpeg could not be found. Please install ffmpeg and ensure it's in your PATH."

    exit 1
fi

# Iterate over each segment directory within SEGMENTS_DIR
for segment_dir in "$SEGMENTS_DIR"/segment_*; do
    if [ -d "$segment_dir" ]; then
        segment_name=$(basename "$segment_dir")
        echo_info "Processing directory: $segment_name"

        # Define paths
        frames_dir="$segment_dir/output_frames"  # Using output_frames for video frames

        # Check if output_frames directory exists
        if [ ! -d "$frames_dir" ]; then
            echo_error "  Directory 'output_frames' does not exist in '$segment_name'. Skipping."
            continue
        fi

        # Find the audio source file (output_*.mp4)
        # Assuming only one output_*.mp4 per segment; adjust if multiple exist
        audio_file=$(find "$segment_dir" -maxdepth 1 -type f -iname "output_*.mp4" | head -n 1)


        if [ -z "$audio_file" ]; then
            echo_error "  No 'output_*.mp4' file found in '$segment_name'. Skipping."
            continue
        fi

        echo_info "  Found audio source: $(basename "$audio_file")"

        # Define output video path

        output_video="$segment_dir/output-matted.mp4"


        # Check if output video already exists
        if [ -f "$output_video" ]; then
            echo_info "  '$output_video' already exists. Overwriting..."
        fi

        # Execute ffmpeg command
        echo_info "  Starting video creation with ffmpeg..."

        # Determine the number of digits in frame filenames
        # Example: 0001.jpg has 4 digits
        first_frame=$(ls "$frames_dir"/*.jpg | head -n 1)
        if [[ "$first_frame" =~ ([0-9]+)\.jpg$ ]]; then
            num_digits=${#BASH_REMATCH[1]}
        else
            echo_error "  Unable to determine frame numbering in '$frames_dir'. Skipping."
            continue
        fi

        # Create a pattern for ffmpeg based on the number of digits
        frame_pattern=$(printf "%%0%dd.jpg" "$num_digits")

        # Run ffmpeg to create the video with audio
        ffmpeg -y \
            -framerate 59.94 \
            -i "$frames_dir/$frame_pattern" \
            -i "$audio_file" \
            -c:v hevc_nvenc \
            -preset slow \
            -b:v 50M \
            -s 7680x3840 \
            -r 59.94 \
            -c:a copy \
            -map 0:v:0 \
            -map 1:a:0 \
            "$output_video" \
            && echo_info "  Successfully created 'output-matted.mp4' in '$segment_name'"

    else
        echo_error "'$segment_dir' is not a directory. Skipping."
    fi

done

echo_info "All segments processed."