#!/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
TEMP_OPUS=""
TEMP_FLAC=""

# Ensure temporary files are cleaned up if the script exits or fails
cleanup() {
    if [[ -n "$TEMP_OPUS" && -f "$TEMP_OPUS" ]]; then rm -f "$TEMP_OPUS"; fi
    if [[ -n "$TEMP_FLAC" && -f "$TEMP_FLAC" ]]; then rm -f "$TEMP_FLAC"; fi
}
trap cleanup EXIT

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, Opus, or both.

Options:
  -i <file>    Input media file (video or audio)
  -o <file>    Output file name (optional; auto-generated if omitted)
  -c <codec>   Target codec: 'opus', 'flac', or 'dual' (default: opus)
               ('dual' encodes to both and muxes into an .mka file)
  -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 dual -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" && "$CODEC" != "dual" ]]; then
    echo "Error: Codec must be 'opus', 'flac', or 'dual'."
    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" || "$CODEC" == "dual" ]] && ! command -v flac >/dev/null 2>&1; then
        echo "Error: 'flac' tool is required for external flac/dual encoding."
        exit 1
    fi
    if [[ "$CODEC" == "opus" || "$CODEC" == "dual" ]] && ! command -v opusenc >/dev/null 2>&1; then
        echo "Error: 'opusenc' tool is required for external opus/dual 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
    if [[ "$CODEC" == "dual" ]]; then
        OUTPUT_FILE="${INPUT_FILE%.*}_track${STREAM_INDEX}.mka"
    else
        OUTPUT_FILE="${INPUT_FILE%.*}_track${STREAM_INDEX}.${CODEC}"
    fi
fi

# Check if output file already exists and abort if it does
if [[ -e "$OUTPUT_FILE" ]]; then
    echo "Error: Output file '$OUTPUT_FILE' already exists. Aborting."
    exit 1
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

# FILTER FIX: Force libopus to see standard 7.1, 5.1, stereo, or mono layouts.
# This silently fixes the "Invalid channel layout 5.1(side)" error.
CHANNEL_FORMAT_FILTER="-filter:a aformat=channel_layouts=7.1|5.1|stereo|mono"

# Set FFmpeg logging profiles
FFMPEG_NATIVE_LOG="-hide_banner -loglevel error -stats"
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}" \
            $CHANNEL_FORMAT_FILTER \
            -c:a libopus -b:a "${TOTAL_BITRATE}k" -vbr on -mapping_family "$MAPPING_FAMILY" \
            "$OUTPUT_FILE"

    elif [[ "$ENCODER_TYPE" == "external" ]]; then
        ffmpeg $FFMPEG_EXT_LOG -y -i "$INPUT_FILE" -map "0:a:${STREAM_INDEX}" \
            $CHANNEL_FORMAT_FILTER -f wav -c:a pcm_s24le - | \
            opusenc --bitrate "$TOTAL_BITRATE" - "$OUTPUT_FILE"
    fi

elif [[ "$CODEC" == "flac" ]]; then
    if [[ "$ENCODER_TYPE" == "native" ]]; then
        ffmpeg $FFMPEG_NATIVE_LOG -y -i "$INPUT_FILE" -map "0:a:${STREAM_INDEX}" \
            $CHANNEL_FORMAT_FILTER \
            -c:a flac -compression_level 8 \
            "$OUTPUT_FILE"

    elif [[ "$ENCODER_TYPE" == "external" ]]; then
        ffmpeg $FFMPEG_EXT_LOG -y -i "$INPUT_FILE" -map "0:a:${STREAM_INDEX}" \
            $CHANNEL_FORMAT_FILTER -f wav -c:a pcm_s24le - | \
            flac --best --ignore-chunk-sizes - -o "$OUTPUT_FILE"
    fi

elif [[ "$CODEC" == "dual" ]]; then
    echo "Target Opus Bitrate: ${TOTAL_BITRATE}k"

    if [[ "$ENCODER_TYPE" == "native" ]]; then
        # Map the stream twice. Note we use specific stream filters (-filter:a:0 and -filter:a:1) here
        # because applying a global -af filter breaks multiple mappings in FFmpeg.
        ffmpeg $FFMPEG_NATIVE_LOG -y -i "$INPUT_FILE" \
            -map "0:a:${STREAM_INDEX}" -filter:a:0 "aformat=channel_layouts=7.1|5.1|stereo|mono" -c:a:0 libopus -b:a:0 "${TOTAL_BITRATE}k" -vbr on -mapping_family "$MAPPING_FAMILY" -metadata:s:a:0 title="Opus (${TOTAL_BITRATE}k)" \
            -map "0:a:${STREAM_INDEX}" -filter:a:1 "aformat=channel_layouts=7.1|5.1|stereo|mono" -c:a:1 flac -compression_level 8 -metadata:s:a:1 title="FLAC (Lossless)" \
            "$OUTPUT_FILE"

    elif [[ "$ENCODER_TYPE" == "external" ]]; then
        TEMP_OPUS=$(mktemp --suffix=.opus)
        TEMP_FLAC=$(mktemp --suffix=.flac)

        echo "[1/3] Encoding Opus externally..."
        ffmpeg $FFMPEG_EXT_LOG -y -i "$INPUT_FILE" -map "0:a:${STREAM_INDEX}" \
            $CHANNEL_FORMAT_FILTER -f wav -c:a pcm_s24le - | \
            opusenc --quiet --bitrate "$TOTAL_BITRATE" - "$TEMP_OPUS"

        echo "[2/3] Encoding FLAC externally..."
        ffmpeg $FFMPEG_EXT_LOG -y -i "$INPUT_FILE" -map "0:a:${STREAM_INDEX}" \
            $CHANNEL_FORMAT_FILTER -f wav -c:a pcm_s24le - | \
            flac --silent --best --ignore-chunk-sizes - -o "$TEMP_FLAC"

        echo "[3/3] Multiplexing tracks into $OUTPUT_FILE..."
        ffmpeg $FFMPEG_NATIVE_LOG -y -i "$TEMP_OPUS" -i "$TEMP_FLAC" \
            -map 0:a -map 1:a -c copy \
            -metadata:s:a:0 title="Opus (${TOTAL_BITRATE}k)" \
            -metadata:s:a:1 title="FLAC (Lossless)" \
            "$OUTPUT_FILE"
    fi
fi

echo -e "\nEncoding complete: $OUTPUT_FILE"