#!/usr/bin/env bash
# Strict mode for robust error handling
set -euo pipefail

# Default variables
INPUT_FILE=""
OUTPUT_FILE=""
CODEC="opus"
ENCODER_TYPE="native"
STREAM_INDEX="0"
BITRATE_PER_CHANNEL=64 # kbps per channel for Opus

usage() {
    cat <<EOF
Usage: $(basename "$0") -i <input_media> [OPTIONS]

A script to extract and encode a specific audio track from a movie file to FLAC or Opus.

Options:
  -i <file>    Input media file (video or audio)
  -o <file>    Output file name (optional; auto-generated if omitted)
  -c <codec>   Target codec: 'opus' or 'flac' (default: opus)
  -e <type>    Encoder type: 'native' (ffmpeg) or 'external' (flac/opusenc) (default: native)
  -s <index>   Audio stream index to select (default: 0, which is the first audio track)
  -h           Show this help message

Examples:
  $(basename "$0") -i movie.mkv -c opus -e native
  $(basename "$0") -i movie.mp4 -c flac -s 1 -o my_custom_audio.flac
EOF
    exit 1
}

# Parse command line arguments
while getopts "i:o:c:e:s:h" opt; do
    case "$opt" in
        i) INPUT_FILE="$OPTARG" ;;
        o) OUTPUT_FILE="$OPTARG" ;;
        c) CODEC=$(echo "$OPTARG" | tr '[:upper:]' '[:lower:]') ;;
        e) ENCODER_TYPE=$(echo "$OPTARG" | tr '[:upper:]' '[:lower:]') ;;
        s) STREAM_INDEX="$OPTARG" ;;
        h|*) usage ;;
    esac
done

# Validate inputs
if [[ -z "$INPUT_FILE" || ! -f "$INPUT_FILE" ]]; then
    echo "Error: Input file not found or not specified."
    usage
fi

if [[ "$CODEC" != "opus" && "$CODEC" != "flac" ]]; then
    echo "Error: Codec must be 'opus' or 'flac'."
    exit 1
fi

if [[ "$ENCODER_TYPE" != "native" && "$ENCODER_TYPE" != "external" ]]; then
    echo "Error: Encoder type must be 'native' or 'external'."
    exit 1
fi

# Validate stream index is a number
if ! [[ "$STREAM_INDEX" =~ ^[0-9]+$ ]]; then
    echo "Error: Stream index must be a non-negative integer."
    exit 1
fi

# Check dependencies
command -v ffmpeg >/dev/null 2>&1 || { echo "Error: ffmpeg is required."; exit 1; }
command -v ffprobe >/dev/null 2>&1 || { echo "Error: ffprobe is required."; exit 1; }

if [[ "$ENCODER_TYPE" == "external" ]]; then
    if [[ "$CODEC" == "flac" ]] && ! command -v flac >/dev/null 2>&1; then
        echo "Error: 'flac' tool is required for external flac encoding."
        exit 1
    fi
    if [[ "$CODEC" == "opus" ]] && ! command -v opusenc >/dev/null 2>&1; then
        echo "Error: 'opusenc' tool is required for external opus encoding."
        exit 1
    fi
fi

# Determine number of audio channels of the selected audio stream
CHANNELS=$(ffprobe -v error -select_streams "a:${STREAM_INDEX}" -show_entries stream=channels -of default=noprint_wrappers=1:nokey=1 "$INPUT_FILE")
if [[ -z "$CHANNELS" ]]; then
    echo "Error: Could not detect audio channels for stream index a:${STREAM_INDEX} in $INPUT_FILE. (Does that track exist?)"
    exit 1
fi

echo "Detected $CHANNELS channel(s) in audio stream ${STREAM_INDEX}."

# Calculate Opus bitrate (ignored for lossless FLAC)
TOTAL_BITRATE=$(( CHANNELS * BITRATE_PER_CHANNEL ))

# Determine output filename if not provided by user
if [[ -z "$OUTPUT_FILE" ]]; then
    OUTPUT_FILE="${INPUT_FILE%.*}_track${STREAM_INDEX}.${CODEC}"
fi

# Handle channel mapping for native Opus (required for >2 channels)
# Family 0 is for mono/stereo. Family 1 is for surround (3 to 8 channels).
MAPPING_FAMILY=0
if [[ "$CHANNELS" -gt 2 ]]; then
    MAPPING_FAMILY=1
fi

# Set FFmpeg logging profiles
# Native: Hide banner, only show errors, but keep the progress bar (-stats)
FFMPEG_NATIVE_LOG="-hide_banner -loglevel error -stats"
# External: Completely silent so it doesn't garble the external tool's output
FFMPEG_EXT_LOG="-hide_banner -loglevel quiet"

echo "Starting encode: Target=$CODEC | Encoder=$ENCODER_TYPE | Stream=a:${STREAM_INDEX} | Output=$OUTPUT_FILE"

# Execute encoding based on user choices
if [[ "$CODEC" == "opus" ]]; then
    echo "Target Bitrate: ${TOTAL_BITRATE}k"
    
    if [[ "$ENCODER_TYPE" == "native" ]]; then
        ffmpeg $FFMPEG_NATIVE_LOG -y -i "$INPUT_FILE" -map "0:a:${STREAM_INDEX}" \
            -c:a libopus -b:a "${TOTAL_BITRATE}k" -vbr on -mapping_family "$MAPPING_FAMILY" \
            "$OUTPUT_FILE"
            
    elif [[ "$ENCODER_TYPE" == "external" ]]; then
        # Pipe 24-bit PCM WAV to opusenc.
        ffmpeg $FFMPEG_EXT_LOG -y -i "$INPUT_FILE" -map "0:a:${STREAM_INDEX}" -f wav -c:a pcm_s24le - | \
            opusenc --bitrate "$TOTAL_BITRATE" - "$OUTPUT_FILE"
    fi

elif [[ "$CODEC" == "flac" ]]; then
    if [[ "$ENCODER_TYPE" == "native" ]]; then
        # Native FLAC encoding
        ffmpeg $FFMPEG_NATIVE_LOG -y -i "$INPUT_FILE" -map "0:a:${STREAM_INDEX}" \
            -c:a flac -compression_level 8 \
            "$OUTPUT_FILE"
            
    elif [[ "$ENCODER_TYPE" == "external" ]]; then
        # Pipe 24-bit PCM WAV to external FLAC encoder.
        ffmpeg $FFMPEG_EXT_LOG -y -i "$INPUT_FILE" -map "0:a:${STREAM_INDEX}" -f wav -c:a pcm_s24le - | \
            flac -8 - -o "$OUTPUT_FILE"
    fi
fi

echo -e "\nEncoding complete: $OUTPUT_FILE"