#!/bin/bash

# Script to extract first frame from WebM files and convert to WebP with alpha transparency
# Usage: ./webm_to_webp.sh [directory_path] [background] [compression]

# Function to display usage
usage() {
    echo "Usage: $0 [directory_path] [background] [compression]"
    echo "  directory_path: Path to directory containing WebM files (default: current directory)"
    echo "  background: Background color for transparency (default: transparent)"
    echo "              - 'transparent' for alpha transparency"
    echo "              - Hex colors: '#000000', '#ffffff', '#ff0000', etc."
    echo "              - Named colors: 'black', 'white', 'red', 'blue', etc."
    echo "  compression: Compression level 1-5 (default: 3)"
    echo "              - 1: Lossless, fast (largest files)"
    echo "              - 2: Lossless, balanced"
    echo "              - 3: Lossless, best compression (default)"
    echo "              - 4: Lossy, high quality (quality 90)"
    echo "              - 5: Lossy, maximum compression (quality 75)"
    echo ""
    echo "Examples:"
    echo "  $0                              # Current directory, transparent, lossless best"
    echo "  $0 /path/to/files               # Custom directory, defaults"
    echo "  $0 /path/to/files black         # Black background, lossless best"
    echo "  $0 /path/to/files '#ff0000' 5   # Red background, maximum compression"
    echo ""
    echo "This script will:"
    echo "  - Find all .webm files in the specified directory"
    echo "  - Extract the frame with maximum coverage from each WebM file"
    echo "  - Convert to WebP format maintaining alpha transparency"
    echo "  - Save results in a 'webp' 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

# Set default values
DIRECTORY="."
BACKGROUND="transparent"
COMPRESSION=3

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

# Validate background parameter
if [ "$BACKGROUND" != "transparent" ]; then
    # Check if it's a valid hex color or named color
    if ! [[ "$BACKGROUND" =~ ^#[0-9a-fA-F]{6}$ ]] && ! [[ "$BACKGROUND" =~ ^[a-zA-Z]+$ ]]; then
        echo "Error: background must be 'transparent', a hex color (#000000), or a named color (black, white, etc.)"
        exit 1
    fi
fi

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

# Set compression values based on level
case $COMPRESSION in
    1) WEBP_LOSSLESS=1; WEBP_COMP_LEVEL=0; WEBP_QUALITY=100; WEBP_METHOD=0; COMP_DESC="Lossless, fast" ;;
    2) WEBP_LOSSLESS=1; WEBP_COMP_LEVEL=3; WEBP_QUALITY=100; WEBP_METHOD=0; COMP_DESC="Lossless, balanced" ;;
    3) WEBP_LOSSLESS=1; WEBP_COMP_LEVEL=6; WEBP_QUALITY=100; WEBP_METHOD=0; COMP_DESC="Lossless, best compression" ;;
    4) WEBP_LOSSLESS=0; WEBP_COMP_LEVEL=4; WEBP_QUALITY=90; WEBP_METHOD=4; COMP_DESC="Lossy, high quality (q90)" ;;
    5) WEBP_LOSSLESS=0; WEBP_COMP_LEVEL=4; WEBP_QUALITY=75; WEBP_METHOD=6; COMP_DESC="Lossy, maximum compression (q75)" ;;
esac

# 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 "Background: $BACKGROUND"
echo "Compression: Level $COMPRESSION ($COMP_DESC)"

# Create webp output directory
WEBP_DIR="$DIRECTORY/webp"
if [ ! -d "$WEBP_DIR" ]; then
    mkdir -p "$WEBP_DIR"
    echo "Created output directory: $WEBP_DIR"
else
    echo "Using existing output directory: $WEBP_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 WebP files will be overwritten to ensure proper transparency"
echo ""

# Counter for processed files
PROCESSED=0
FAILED=0
OVERWRITTEN=0
NEW_FILES=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 WebP file path
    webp_output="$WEBP_DIR/${base_name}.webp"

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

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

    # Detect codec type (VP8 or VP9) to use correct decoder
    CODEC_INFO=$(ffprobe -v quiet -select_streams v:0 -show_entries stream=codec_name -of csv=p=0 "$webm_file" 2>/dev/null)

    # Set appropriate decoder based on codec
    if [ "$CODEC_INFO" = "vp9" ]; then
        DECODER="libvpx-vp9"
        echo "  → Using VP9 decoder (optimal for alpha preservation)"
    else
        DECODER="libvpx"  # Default to VP8 decoder
        echo "  → Using VP8 decoder"
    fi

    echo "  → Finding frame with maximum coverage..."

    # Get total frame count and duration
    FRAME_COUNT=$(ffprobe -v quiet -select_streams v:0 -count_frames -show_entries stream=nb_read_frames -of csv=p=0 "$webm_file" 2>/dev/null)

    if [ -z "$FRAME_COUNT" ] || [ "$FRAME_COUNT" -eq 0 ]; then
        echo "  ✗ Could not determine frame count"
        ((FAILED++))
        continue
    fi

    # Sample 10 frames evenly distributed throughout the video
    SAMPLE_COUNT=10
    if [ "$FRAME_COUNT" -lt "$SAMPLE_COUNT" ]; then
        SAMPLE_COUNT=$FRAME_COUNT
    fi

    MAX_SIZE=0
    BEST_FRAME=0
    TEMP_DIR=$(mktemp -d)

    echo "  → Analyzing $SAMPLE_COUNT sample frames from $FRAME_COUNT total frames..."

    # Sample frames evenly throughout the video
    for ((i=0; i<SAMPLE_COUNT; i++)); do
        # Calculate frame position (evenly distributed)
        FRAME_POS=$((i * FRAME_COUNT / SAMPLE_COUNT))

        # Extract single frame as WebP to temp location
        TEMP_FRAME="$TEMP_DIR/frame_${FRAME_POS}.webp"

        ffmpeg -c:v "$DECODER" \
               -i "$webm_file" \
               -vf "select=eq(n\\,$FRAME_POS)" \
               -vframes 1 \
               -c:v libwebp \
               -lossless 1 \
               -y \
               "$TEMP_FRAME" \
               -v quiet 2>/dev/null

        if [ -f "$TEMP_FRAME" ]; then
            # Use file size as coverage indicator (larger = more content)
            FRAME_SIZE=$(stat -f%z "$TEMP_FRAME" 2>/dev/null || stat -c%s "$TEMP_FRAME" 2>/dev/null)

            if [ -n "$FRAME_SIZE" ] && [ "$FRAME_SIZE" -gt "$MAX_SIZE" ]; then
                MAX_SIZE=$FRAME_SIZE
                BEST_FRAME=$FRAME_POS
            fi
        fi
    done

    # Clean up temp directory
    rm -rf "$TEMP_DIR"

    echo "  → Best frame: #$BEST_FRAME (size: $MAX_SIZE bytes)"

    # Set up filter chain based on background setting
    if [ "$BACKGROUND" = "transparent" ]; then
        FILTER_CHAIN="select=eq(n\\,$BEST_FRAME)"
        echo "  → Extracting frame with alpha transparency preserved"
    else
        FILTER_CHAIN="select=eq(n\\,$BEST_FRAME),split[a][b];[a]geq=r=0:g=0:b=0:a=255,drawbox=color=$BACKGROUND@1.0:replace=1:t=fill[bg];[bg][b]overlay=format=auto"
        echo "  → Extracting frame with background color: $BACKGROUND"
    fi

    # Build WebP encoding parameters based on compression level
    if [ "$WEBP_LOSSLESS" -eq 1 ]; then
        WEBP_PARAMS="-c:v libwebp -lossless 1 -compression_level $WEBP_COMP_LEVEL"
    else
        WEBP_PARAMS="-c:v libwebp -lossless 0 -quality $WEBP_QUALITY -method $WEBP_METHOD -compression_level $WEBP_COMP_LEVEL"
    fi

    # Extract the frame with maximum coverage
    if ffmpeg -c:v "$DECODER" \
              -i "$webm_file" \
              -vf "$FILTER_CHAIN" \
              -vframes 1 \
              $WEBP_PARAMS \
              -an \
              -y \
              "$webp_output" \
              -v quiet -stats; then
        if [ "$FILE_EXISTS" = true ]; then
            echo "  ✓ Successfully overwritten: ${base_name}.webp"
            ((OVERWRITTEN++))
        else
            echo "  ✓ Successfully created: ${base_name}.webp"
            ((NEW_FILES++))
        fi
        ((PROCESSED++))
    else
        echo "  ✗ Failed to process: $(basename "$webm_file")"
        ((FAILED++))
    fi
    echo ""
done

# Summary
echo "=== Processing 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: $WEBP_DIR"
echo "Settings used - Background: $BACKGROUND, Compression: Level $COMPRESSION ($COMP_DESC)"
echo ""

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