#!/bin/bash

# Script to compress WebM files using VP9 codec
# Usage: ./webm_compress.sh [directory_path] [compression_level] [audio_mode] [fps]

# Function to display usage
usage() {
    echo "Usage: $0 [directory_path] [compression_level] [audio_mode] [fps]"
    echo "  directory_path: Path to directory containing WebM files (default: current directory)"
    echo "  compression_level: Compression level 3-5 (default: 3)"
    echo "              - 3: Lossless, best compression (default)"
    echo "              - 4: Lossy, high quality (CRF 23)"
    echo "              - 5: Lossy, maximum compression (CRF 35)"
    echo "  audio_mode: Audio channel mode (default: stereo)"
    echo "              - stereo: Keep stereo audio (2 channels)"
    echo "              - mono: Convert to mono audio (1 channel, smaller files)"
    echo "  fps: Target frame rate (default: 30)"
    echo ""
    echo "Examples:"
    echo "  $0                            # Current directory, lossless, stereo, 30fps"
    echo "  $0 /path/to/files             # Custom directory, defaults for other params"
    echo "  $0 /path/to/files 4           # High quality lossy compression"
    echo "  $0 /path/to/files 4 mono      # High quality + mono audio"
    echo "  $0 /path/to/files 4 mono 15   # High quality + mono + 15fps"
    echo "  $0 /path/to/files 5 stereo 24 # Max compression + stereo + 24fps"
    echo ""
    echo "This script will:"
    echo "  - Find all .webm files in the specified directory"
    echo "  - Compress using VP9 codec with specified quality level"
    echo "  - Preserve alpha transparency when present"
    echo "  - Adjust frame rate and audio channels as specified"
    echo "  - Save results in a 'webm_compressed' subdirectory"
    exit 1
}

# Function to format file size in human-readable format
format_size() {
    local size=$1
    if command -v numfmt &> /dev/null; then
        numfmt --to=iec "$size"
    else
        # Fallback function for systems without numfmt
        if [ "$size" -lt 1024 ]; then
            echo "${size}B"
        elif [ "$size" -lt 1048576 ]; then
            echo "$(( size / 1024 ))K"
        elif [ "$size" -lt 1073741824 ]; then
            echo "$(( size / 1048576 ))M"
        else
            echo "$(( size / 1073741824 ))G"
        fi
    fi
}

# 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 libvpx decoders are available (needed for alpha channel support)
LIBVPX_CHECK=$(ffmpeg -decoders 2>/dev/null | grep libvpx)
if [ -z "$LIBVPX_CHECK" ]; then
    echo "Warning: libvpx decoders not found. Alpha channels may not be preserved properly."
    echo "Your ffmpeg installation may not support VP8/VP9 alpha transparency."
    echo ""
fi

# Check if bc is available (needed for floating point calculations)
if ! command -v bc &> /dev/null; then
    echo "Error: bc (basic calculator) is not installed or not in PATH"
    echo "Please install bc to use this script"
    exit 1
fi

# Set default values
DIRECTORY="."
COMPRESSION=3
AUDIO_MODE="stereo"
FPS=30

# Parse arguments
if [ $# -eq 0 ]; then
    DIRECTORY="."
elif [ $# -eq 1 ]; then
    DIRECTORY="$1"
elif [ $# -eq 2 ]; then
    DIRECTORY="$1"
    COMPRESSION="$2"
elif [ $# -eq 3 ]; then
    DIRECTORY="$1"
    COMPRESSION="$2"
    AUDIO_MODE="$3"
elif [ $# -eq 4 ]; then
    DIRECTORY="$1"
    COMPRESSION="$2"
    AUDIO_MODE="$3"
    FPS="$4"
else
    echo "Error: Too many arguments"
    usage
fi

# Validate compression parameter
if ! [[ "$COMPRESSION" =~ ^[3-5]$ ]]; then
    echo "Error: compression level must be between 3-5"
    exit 1
fi

# Validate audio mode parameter
if [ "$AUDIO_MODE" != "stereo" ] && [ "$AUDIO_MODE" != "mono" ]; then
    echo "Error: audio_mode must be 'stereo' or 'mono'"
    exit 1
fi

# Validate fps parameter
if ! [[ "$FPS" =~ ^[0-9]+(\.[0-9]+)?$ ]] || (( $(echo "$FPS <= 0" | bc -l) )); then
    echo "Error: fps must be a positive number"
    exit 1
fi

# Set compression values based on level
case $COMPRESSION in
    3) VP9_PARAMS="-c:v libvpx-vp9 -lossless 1 -cpu-used 0"; COMP_DESC="Lossless, best compression" ;;
    4) VP9_PARAMS="-c:v libvpx-vp9 -crf 23 -b:v 0 -cpu-used 1"; COMP_DESC="Lossy, high quality (CRF 23)" ;;
    5) VP9_PARAMS="-c:v libvpx-vp9 -crf 35 -b:v 0 -cpu-used 2"; COMP_DESC="Lossy, maximum compression (CRF 35)" ;;
esac

# Set audio parameters based on mode
if [ "$AUDIO_MODE" = "mono" ]; then
    AUDIO_PARAMS="-c:a libopus -ac 1 -b:a 64k"
    AUDIO_DESC="mono (64k)"
else
    AUDIO_PARAMS="-c:a libopus -ac 2 -b:a 128k"
    AUDIO_DESC="stereo (128k)"
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 "Processing WebM files in: $DIRECTORY"
echo "Compression: Level $COMPRESSION ($COMP_DESC)"
echo "Audio: $AUDIO_DESC"
echo "Target FPS: $FPS"

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

# 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)"
echo "Note: Existing compressed files will be overwritten"
echo ""

# Counter for processed files
PROCESSED=0
FAILED=0
OVERWRITTEN=0
NEW_FILES=0
TOTAL_ORIGINAL_SIZE=0
TOTAL_COMPRESSED_SIZE=0

# Process each WebM file
for webm_file in "${WEBM_FILES[@]}"; do
    # Get the base filename without path and extension
    base_name=$(basename "$webm_file" .webm)

    # Output compressed WebM file path
    compressed_output="$COMPRESSED_DIR/${base_name}.webm"

    echo "Processing: $(basename "$webm_file")"

    # Get original file size
    ORIGINAL_SIZE=$(stat -f%z "$webm_file" 2>/dev/null || stat -c%s "$webm_file" 2>/dev/null)
    TOTAL_ORIGINAL_SIZE=$((TOTAL_ORIGINAL_SIZE + ORIGINAL_SIZE))

    # Check if output file already exists
    FILE_EXISTS=false
    if [ -f "$compressed_output" ]; then
        echo "  → Overwriting existing file: ${base_name}.webm"
        FILE_EXISTS=true
    fi

    # Detect codec type (VP8 or VP9) and other properties
    CODEC_INFO=$(ffprobe -v quiet -select_streams v:0 -show_entries stream=codec_name -of csv=p=0 "$webm_file" 2>/dev/null)
    PIXEL_FORMAT=$(ffprobe -v quiet -select_streams v:0 -show_entries stream=pix_fmt -of csv=p=0 "$webm_file" 2>/dev/null)
    HAS_ALPHA=$(echo "$PIXEL_FORMAT" | grep -iE "(yuva|rgba|gbra|argb|abgr|bgra|a444|a422|a420)")

    # Check for alpha_mode metadata (more reliable for WebM)
    ALPHA_MODE=$(ffprobe -v quiet -select_streams v:0 -show_entries stream_tags=alpha_mode -of csv=p=0 "$webm_file" 2>/dev/null)

    echo "  → Original codec: $CODEC_INFO (format: $PIXEL_FORMAT)"

    # Set the appropriate decoder for alpha preservation
    case $CODEC_INFO in
        "vp8")
            DECODER_PARAMS="-c:v libvpx"
            echo "  → Using libvpx decoder for VP8"
            ;;
        "vp9")
            DECODER_PARAMS="-c:v libvpx-vp9"
            echo "  → Using libvpx-vp9 decoder for VP9"
            ;;
        *)
            DECODER_PARAMS=""
            echo "  → Using default decoder for $CODEC_INFO"
            ;;
    esac

    if [ -n "$HAS_ALPHA" ] || [ "$ALPHA_MODE" = "1" ]; then
        echo "  → Alpha channel detected, preserving transparency"
        # For alpha content, use special VP9 settings and let ffmpeg choose pixel format
        case $COMPRESSION in
            3)
                # Lossless doesn't work well with alpha in VP9, use very high quality instead
                ALPHA_VP9_PARAMS="-c:v libvpx-vp9 -crf 0 -b:v 0 -cpu-used 0"
                COMP_DESC_ALPHA="Near-lossless (CRF 0) for alpha preservation"
                ;;
            4)
                ALPHA_VP9_PARAMS="-c:v libvpx-vp9 -crf 23 -b:v 0 -cpu-used 1"
                COMP_DESC_ALPHA="Lossy, high quality (CRF 23)"
                ;;
            5)
                ALPHA_VP9_PARAMS="-c:v libvpx-vp9 -crf 35 -b:v 0 -cpu-used 2"
                COMP_DESC_ALPHA="Lossy, maximum compression (CRF 35)"
                ;;
        esac
        # Don't force pixel format for alpha - let ffmpeg auto-detect
        ALPHA_PARAMS=""
        VIDEO_FILTER="fps=$FPS"
        FINAL_VP9_PARAMS="$ALPHA_VP9_PARAMS"
        echo "  → Using alpha-compatible settings: $COMP_DESC_ALPHA"
        HAS_ALPHA_CONTENT=true
    else
        echo "  → No alpha channel detected"
        ALPHA_PARAMS="-pix_fmt yuv420p"
        VIDEO_FILTER="fps=$FPS"
        FINAL_VP9_PARAMS="$VP9_PARAMS"
        HAS_ALPHA_CONTENT=false
        echo "  → Using standard VP9 compression"
    fi

    echo "  → Compressing with VP9..."

    # Compress the WebM file - use decoder params BEFORE input for alpha preservation
    if ffmpeg $DECODER_PARAMS -i "$webm_file" \
              -vf "$VIDEO_FILTER" \
              $FINAL_VP9_PARAMS \
              $ALPHA_PARAMS \
              $AUDIO_PARAMS \
              -threads 0 \
              -row-mt 1 \
              -y \
              "$compressed_output" \
              -v quiet -stats 2>/dev/null; then

        # Get compressed file size
        COMPRESSED_SIZE=$(stat -f%z "$compressed_output" 2>/dev/null || stat -c%s "$compressed_output" 2>/dev/null)
        TOTAL_COMPRESSED_SIZE=$((TOTAL_COMPRESSED_SIZE + COMPRESSED_SIZE))

        # Verify alpha preservation if original had alpha
        if [ "$HAS_ALPHA_CONTENT" = true ]; then
            OUTPUT_PIXEL_FORMAT=$(ffprobe -v quiet -select_streams v:0 -show_entries stream=pix_fmt -of csv=p=0 "$compressed_output" 2>/dev/null)
            OUTPUT_HAS_ALPHA=$(echo "$OUTPUT_PIXEL_FORMAT" | grep -iE "(yuva|rgba|gbra|argb|abgr|bgra|a444|a422|a420)")
            OUTPUT_ALPHA_MODE=$(ffprobe -v quiet -select_streams v:0 -show_entries stream_tags=alpha_mode -of csv=p=0 "$compressed_output" 2>/dev/null)

            if [ -n "$OUTPUT_HAS_ALPHA" ] || [ "$OUTPUT_ALPHA_MODE" = "1" ]; then
                echo "  ✓ Alpha channel preserved (format: $OUTPUT_PIXEL_FORMAT, alpha_mode: $OUTPUT_ALPHA_MODE)"
            else
                echo "  ⚠ Warning: Alpha channel may have been lost (output format: $OUTPUT_PIXEL_FORMAT)"
            fi
        fi

        # Calculate compression ratio
        RATIO=$(echo "scale=1; $COMPRESSED_SIZE * 100 / $ORIGINAL_SIZE" | bc)
        SAVINGS=$(echo "scale=1; 100 - $RATIO" | bc)

        if [ "$FILE_EXISTS" = true ]; then
            echo "  ✓ Successfully overwritten: ${base_name}.webm"
            ((OVERWRITTEN++))
        else
            echo "  ✓ Successfully created: ${base_name}.webm"
            ((NEW_FILES++))
        fi
        echo "  → Size: $(format_size $ORIGINAL_SIZE) → $(format_size $COMPRESSED_SIZE) (${RATIO}% of original, ${SAVINGS}% savings)"
        ((PROCESSED++))
    else
        echo "  ✗ Failed to compress: $(basename "$webm_file")"
        ((FAILED++))
    fi
    echo ""
done

# Calculate overall compression statistics
if [ $TOTAL_ORIGINAL_SIZE -gt 0 ]; then
    OVERALL_RATIO=$(echo "scale=1; $TOTAL_COMPRESSED_SIZE * 100 / $TOTAL_ORIGINAL_SIZE" | bc)
    OVERALL_SAVINGS=$(echo "scale=1; 100 - $OVERALL_RATIO" | bc)
fi

# Summary
echo "=== Compression Complete ==="
echo "Total files processed: $PROCESSED"
if [ $NEW_FILES -gt 0 ]; then
    echo "New files created: $NEW_FILES"
fi
if [ $OVERWRITTEN -gt 0 ]; then
    echo "Existing files overwritten: $OVERWRITTEN"
fi
if [ $FAILED -gt 0 ]; then
    echo "Failed to process: $FAILED files"
fi
echo "Output directory: $COMPRESSED_DIR"
echo "Settings used - Compression: Level $COMPRESSION ($COMP_DESC), Audio: $AUDIO_DESC, FPS: $FPS"
echo "Note: Alpha-containing files may use different compression settings for transparency preservation"

if [ $PROCESSED -gt 0 ]; then
    echo ""
    echo "=== Compression Statistics ==="
    echo "Total original size: $(format_size $TOTAL_ORIGINAL_SIZE)"
    echo "Total compressed size: $(format_size $TOTAL_COMPRESSED_SIZE)"
    echo "Overall compression: ${OVERALL_RATIO}% of original (${OVERALL_SAVINGS}% space saved)"
fi
echo ""

# List generated files
if [ $PROCESSED -gt 0 ]; then
    echo "Compressed WebM files:"
    ls -la "$COMPRESSED_DIR"/*.webm 2>/dev/null | while read -r line; do
        echo "  $line"
    done
fi