#!/bin/bash

# Script to trim empty frames from start and end of WebM files
# Usage: ./trim_webm_empty_frames.sh [directory_path] [sample_rate]

# Function to display usage
usage() {
    echo "Usage: $0 [directory_path] [sample_rate]"
    echo "  directory_path: Path to directory containing WebM files to trim (default: current directory)"
    echo "  sample_rate: How many frames to sample per second for analysis (default: 2)"
    echo "                Higher values = more accurate but slower (1-10 recommended)"
    echo ""
    echo "Examples:"
    echo "  $0                          # Current directory, sample 2 frames/sec"
    echo "  $0 /path/to/files           # Custom directory, default sampling"
    echo "  $0 /path/to/files 4         # Custom directory, sample 4 frames/sec"
    echo ""
    echo "This script will:"
    echo "  - Analyze WebM files to find first and last non-empty frames"
    echo "  - Trim away all empty frames from start and end"
    echo "  - Preserve original quality and alpha transparency"
    echo "  - Save trimmed files in a 'trim' subdirectory"
    exit 1
}

# Check if ffmpeg is installed
if ! command -v ffmpeg &> /dev/null; then
    echo "Error: ffmpeg is not installed or not in PATH"
    echo "Please install ffmpeg to use this script"
    exit 1
fi

# Check if bc is available for floating point calculations
if ! command -v bc >/dev/null 2>&1; then
    echo "Warning: 'bc' calculator not found. Using basic arithmetic (less precise)."
    echo "For better precision, install 'bc': brew install bc (macOS) or apt-get install bc (Linux)"
    echo ""
fi

# Set default values
DIRECTORY="."
SAMPLE_RATE=2

# Parse arguments
if [ $# -eq 0 ]; then
    DIRECTORY="."
elif [ $# -eq 1 ]; then
    if [ "$1" = "-h" ] || [ "$1" = "--help" ]; then
        usage
    fi
    DIRECTORY="$1"
elif [ $# -eq 2 ]; then
    DIRECTORY="$1"
    SAMPLE_RATE="$2"
else
    echo "Error: Too many arguments"
    usage
fi

# Validate sample rate
if ! [[ "$SAMPLE_RATE" =~ ^[1-9][0-9]*$ ]] || [ "$SAMPLE_RATE" -gt 10 ]; then
    echo "Error: sample_rate must be between 1-10"
    exit 1
fi

# Check if directory exists
if [ ! -d "$DIRECTORY" ]; then
    echo "Error: Directory '$DIRECTORY' does not exist"
    exit 1
fi

# Convert to absolute path
DIRECTORY=$(realpath "$DIRECTORY")
echo "Trimming WebM files in: $DIRECTORY"
echo "Sample rate: $SAMPLE_RATE frames per second"

# Create trim output directory
TRIM_DIR="$DIRECTORY/trim"
if [ ! -d "$TRIM_DIR" ]; then
    mkdir -p "$TRIM_DIR"
    echo "Created output directory: $TRIM_DIR"
else
    echo "Using existing output directory: $TRIM_DIR"
fi

# Create temporary directory for frame analysis
temp_dir=$(mktemp -d)
trap "rm -rf $temp_dir" EXIT

# Function to check if a frame is empty (returns 0 for empty, 1 for content, 2 for unclear)
check_frame_empty() {
    local frame_file="$1"

    # Check if frame file exists
    if [ ! -f "$frame_file" ]; then
        return 2  # unclear/error
    fi

    # Get file size - very small files are likely empty/transparent
    file_size=$(stat -f%z "$frame_file" 2>/dev/null || stat -c%s "$frame_file" 2>/dev/null)

    # If file is very small (less than 2KB), likely empty
    if [ -n "$file_size" ] && [ "$file_size" -lt 2048 ]; then
        return 0  # definitely empty
    fi

    # If file is reasonably large, assume it has content
    if [ "$file_size" -gt 50000 ]; then
        return 1  # definitely has content
    fi

    # For medium-sized files, try alpha analysis
    alpha_info=$(ffmpeg -i "$frame_file" -vf "extractplanes=a" -f null - 2>&1 | grep "mean\|avg" 2>/dev/null)

    if [ -n "$alpha_info" ]; then
        alpha_mean=$(echo "$alpha_info" | grep -o "mean:[0-9.]*" | cut -d: -f2 | head -1)

        if [ -n "$alpha_mean" ]; then
            # Check if mean is very low (close to 0 = transparent)
            is_empty=$(awk "BEGIN {print ($alpha_mean < 2.0) ? 1 : 0}")

            if [ "$is_empty" = "1" ]; then
                return 0  # definitely empty
            else
                return 1  # definitely has content
            fi
        fi
    fi

    # Default: assume medium files have some content (conservative approach)
    return 1
}

# Function to find first non-empty frame
find_first_content_frame() {
    local video_file="$1"
    local duration="$2"
    local decoder="$3"

    # Check frames at 0, 0.5, 1.0, 1.5 seconds etc
    for i in 0 0.5 1.0 1.5 2.0 2.5 3.0 3.5 4.0 4.5 5.0; do
        local test_frame="$temp_dir/start_${i}.png"

        ffmpeg -c:v "$decoder" -ss "$i" -i "$video_file" -frames:v 1 -y "$test_frame" -v quiet 2>/dev/null

        if [ -f "$test_frame" ]; then
            check_frame_empty "$test_frame"
            if [ $? -eq 1 ]; then
                echo "$i"
                rm -f "$test_frame"
                return 0
            fi
            rm -f "$test_frame"
        fi
    done

    echo "0"
}

# Function to find last non-empty frame
find_last_content_frame() {
    local video_file="$1"
    local duration="$2"
    local decoder="$3"

    # Work backwards from end
    local dur_int=${duration%.*}

    for i in 0.1 0.5 1.0 1.5 2.0 2.5 3.0 3.5 4.0 4.5 5.0; do
        local check_time=$(echo "$duration - $i" | bc 2>/dev/null || echo "$dur_int")
        local test_frame="$temp_dir/end_${i}.png"

        ffmpeg -c:v "$decoder" -ss "$check_time" -i "$video_file" -frames:v 1 -y "$test_frame" -v quiet 2>/dev/null

        if [ -f "$test_frame" ]; then
            check_frame_empty "$test_frame"
            if [ $? -eq 1 ]; then
                echo "$check_time"
                rm -f "$test_frame"
                return 0
            fi
            rm -f "$test_frame"
        fi
    done

    echo "$duration"
}

# Find all WebM files in the directory
WEBM_FILES=($(find "$DIRECTORY" -maxdepth 1 -name "*.webm" -type f))

# Check if any WebM files were found
if [ ${#WEBM_FILES[@]} -eq 0 ]; then
    echo "No WebM files found in '$DIRECTORY'"
    exit 0
fi

echo "Found ${#WEBM_FILES[@]} WebM file(s) to process"
echo ""

# Counters
processed=0
skipped=0
failed=0

# Process each WebM file
for webm_file in "${WEBM_FILES[@]}"; do
    count=$((count + 1))
    base_name=$(basename "$webm_file" .webm)

    echo "[$((processed + skipped + failed + 1))/${#WEBM_FILES[@]}] Processing: $base_name.webm"

    # Check if output file already exists
    output_file="$TRIM_DIR/${base_name}.webm"
    if [ -f "$output_file" ]; then
        echo "    → Skipping: Output file already exists"
        ((skipped++))
        echo ""
        continue
    fi

    # Test if file is readable by ffprobe
    if ! ffprobe -v quiet "$webm_file" >/dev/null 2>&1; then
        echo "    ✗ Error: File cannot be read by ffprobe"
        ((failed++))
        echo ""
        continue
    fi

    # Detect codec type and get video info
    CODEC_INFO=$(ffprobe -v quiet -select_streams v:0 -show_entries stream=codec_name -of csv=p=0 "$webm_file" 2>/dev/null)
    duration=$(ffprobe -v quiet -show_entries format=duration -of csv=p=0 "$webm_file" 2>/dev/null)

    if [ -z "$duration" ] || [ "$duration" = "N/A" ]; then
        echo "    ✗ Error: Could not get video duration"
        ((failed++))
        echo ""
        continue
    fi

    # Set appropriate decoder
    if [ "$CODEC_INFO" = "vp9" ]; then
        DECODER="libvpx-vp9"
        echo "    → VP9 codec detected"
    else
        DECODER="libvpx"
        echo "    → VP8 codec detected"
    fi

    echo "    → Duration: ${duration}s"

    # Find content boundaries
    echo "    → Analyzing frames..."
    start_time=$(find_first_content_frame "$webm_file" "$duration" "$DECODER")
    end_time=$(find_last_content_frame "$webm_file" "$duration" "$DECODER")

    # Validate times
    if [ -z "$start_time" ] || [ -z "$end_time" ]; then
        echo "    ✗ Error: Could not determine content boundaries"
        ((failed++))
        echo ""
        continue
    fi

    echo "    → Content from ${start_time}s to ${end_time}s"

    # Calculate new duration
    new_duration=$(echo "$end_time - $start_time" | bc 2>/dev/null || echo "$end_time")

    # Simple trimming check
    if [ "$start_time" = "0" ] && [ "$end_time" = "$duration" ]; then
        echo "    → No trimming needed, copying original..."
        cp "$webm_file" "$output_file"
        if [ $? -eq 0 ]; then
            echo "    ✓ File copied successfully"
            ((processed++))
        else
            echo "    ✗ Failed to copy file"
            ((failed++))
        fi
    else
        echo "    → Trimming (new duration: ${new_duration}s)"

        if ffmpeg -ss "$start_time" -i "$webm_file" -t "$new_duration" -c copy -y "$output_file" -v quiet 2>/dev/null; then
            echo "    ✓ Successfully trimmed"
            ((processed++))
        else
            echo "    ✗ Failed to trim"
            ((failed++))
        fi
    fi

    echo ""
done

# Summary
echo "==================================="
echo "Trimming Complete!"
echo "Total files processed: $processed"
echo "Files skipped (already exist): $skipped"
echo "Files failed: $failed"
echo "Output directory: $TRIM_DIR"

if [ $processed -gt 0 ]; then
    echo ""
    echo "Trimmed files:"
    ls -la "$TRIM_DIR"/*.webm 2>/dev/null | while read -r line; do
        echo "  $line"
    done
fi
