Dual Audio Release Creating

#!/bin/bash
# This script checks a file for whether it has both german and english audio.
# If one's missing. it makes radarr grab a new release to extract the audio from
# and inserts it into the original file.
# It also synchronizes the inserted audio track with the original audio track
# using their audio energy profiles
#
# Usage: ./createSuperRelease_radarr.sh /path/to/file.mkv

set -euo pipefail
main(){
  ###############################################################################################################
  #################################### Modify following Variables to your liking ################################
  ###############################################################################################################

  RADARR_API_KEY="secretkey" # your radarr API key
  RADARR_URL="http://127.0.0.1:7878" # your local radarr URL
  RADARR_DB_PATH="/opt/radarr/radarr.db" # absolute host path to radarr database

  SABNZBD_URL="http://127.0.0.1:8080" # your local SABnzbd URL
  SABNZBD_API_KEY="secretkey" # your SABnzbd API key

  WORKING_DIR="/temp/transcode_cache/" # work directory for merging audio streams and creating new end file
  PERMS_GROUP="root:root" # file permission user:group for resulting file
  PERMS_BITS="755" # file permission bits user|group|other

  # uncomment this section and execute it to retrieve the quality profile names and IDs from your radarr database
  # sqlite3 -noheader -batch $RADARR_DB_PATH \
  #   "SELECT Id,Name FROM QualityProfiles"
  # exit 0
  QualityProfileId_ger=21 # one profile is needed for first language and another for second language, you gotta create these QualityProfiles yourself and retrieve the internal IDs
  QualityProfileId_eng=20

  ###############################################################################################################
  #################################### Modify above Variables to your liking ####################################
  ###############################################################################################################

  FILE_PATH="$1"

  set +H
  START_TIME=$(date -u +"%Y-%m-%dT%H:%M:%SZ")
  echo "$START_TIME IS START_TIME"

  ENGLISH_ID=1 # internal radarr Language IDs, do not change unless you know what you're doing
  GERMAN_ID=4

  if [[ ! -f "$FILE_PATH" ]]; then
    echo "Input File doesn't exist"
    exit 1
  fi


  ### Retrieve file path and trim to Short Path ###
  echo "Processing: $FILE_PATH"
  REL_PATH="$(basename "$FILE_PATH")"
  echo "Handling file: $REL_PATH"

  ### Retrieve unique movie file ID and Release Group ###
  IFS='|' read -r MOVIE_FILE_ID RELEASE_GROUP MOVIE_ID < <(
      sqlite3 -noheader -batch $RADARR_DB_PATH \
      "SELECT Id,ReleaseGroup,MovieId FROM MovieFiles WHERE RelativePath = '$REL_PATH'"
  )
  echo "Found MOVIE FILE ID: $MOVIE_FILE_ID AND Release Group: $RELEASE_GROUP and Movie ID: $MOVIE_ID"

  ### check whether file has english audio according to radarr's ffprobe scan ###
  HAS_ENGLISH=$(sqlite3 -noheader -batch $RADARR_DB_PATH "
  SELECT EXISTS (
    SELECT 1
    FROM MovieFiles,
        json_each(MovieFiles.Languages)
    WHERE MovieFiles.Id = $MOVIE_FILE_ID
      AND json_each.value = $ENGLISH_ID
  );
  ")

  ### check whether file has german audio according to radarr's ffprobe scan ###
  HAS_GERMAN=$(sqlite3 -noheader -batch $RADARR_DB_PATH "
  SELECT EXISTS (
    SELECT 1
    FROM MovieFiles,
        json_each(MovieFiles.Languages)
    WHERE MovieFiles.Id = $MOVIE_FILE_ID
      AND json_each.value = $GERMAN_ID
  );
  ")

  ### only process files that either have german or english audio stream, not both or neither ###
  if [[ "$HAS_ENGLISH" -eq 1 && "$HAS_GERMAN" -eq 1 ]]; then
      echo "German and English included. Exiting"
      exit 3
  elif [[ "$HAS_ENGLISH" != 1 && "$HAS_GERMAN" != 1 ]]; then
      echo "German and English missing. Exiting"
      exit 3
  elif [[ "$HAS_GERMAN" -eq 1 ]]; then
      echo "English missing"
      MISSING_LANG="$QualityProfileId_eng"
      MISSING_LANG_NAME="eng"
  elif [[ "$HAS_ENGLISH" -eq 1 ]]; then
      echo "German missing"
      MISSING_LANG="$QualityProfileId_ger"
      MISSING_LANG_NAME="ger"
  fi

  ### execute ffprobe to check stream properties ###
  echo "Executing ffprobe for check"
  ffprobe "$FILE_PATH"

  ### Retrieve current Quality Profile ID ###
  QualityProfileId_original=$(sqlite3 -noheader -batch $RADARR_DB_PATH \
    "SELECT QualityProfileId FROM Movies WHERE Id = $MOVIE_ID")

  ### Retrieve new Quality Profile Name ###
  QualityProfileName=$(sqlite3 -noheader -batch $RADARR_DB_PATH \
    "SELECT Name FROM QualityProfiles WHERE Id = $MISSING_LANG")

  ### Change Quality Profile to missing language ###
  echo "Temporarily changing Movie QualityProfile to $MISSING_LANG $QualityProfileName"
  sqlite3 -noheader -batch $RADARR_DB_PATH \
    "UPDATE Movies SET QualityProfileId = $MISSING_LANG WHERE Id = $MOVIE_ID"

  ### Wait a short moment so the database changes can propagate ###
  sleep 1

  ### Use radarr API to start searching for new release that has the missing language ###
  echo "Searching for new release matching missing language"
  RESPONSE=$(curl -s -X POST "$RADARR_URL/api/v3/command" \
    -H "X-Api-Key: $RADARR_API_KEY" \
    -H "Content-Type: application/json" \
    -d "{
          \"name\":\"MoviesSearch\",
          \"movieIds\":[$MOVIE_ID]
        }")
  COMMAND_ID=$(jq -r '.id' <<< "$RESPONSE")

  ### Wait a short moment for the radarr search to commence
  sleep 1

  ### Retrieve Name of original Quality Profile ###
  echo "Retrieving original QualityProfileName"
  QualityProfileName=$(sqlite3 -noheader -batch $RADARR_DB_PATH \
    "SELECT Name FROM QualityProfiles WHERE Id = $QualityProfileId_original")

  ### Change Movie Quality Profile back to original Quality Profile while radarr keeps searching for new release ###
  echo "Changing Movie QualityProfile back to $QualityProfileName"
  sqlite3 -noheader -batch $RADARR_DB_PATH \
    "UPDATE Movies SET QualityProfileId = $QualityProfileId_original WHERE Id = $MOVIE_ID"


  ### Use radarr API to periodically check whether movie search has found a new release yet ###
  while true; do
      RESPONSE=$(curl -s \
        -H "X-Api-Key: $RADARR_API_KEY" \
        "$RADARR_URL/api/v3/command/$COMMAND_ID")
      echo "$RESPONSE"

      STATUS=$(jq -r '.status' <<< "$RESPONSE")
      RESULT=$(jq -r '.result' <<< "$RESPONSE")
      REPORTS=$(jq -r '.message | capture("(?<count>[0-9]+) reports downloaded") | .count' <<< "$RESPONSE")

      echo "radarr Command Status: $STATUS"

      if [[ "$STATUS" == "completed" ]]; then
          if [[ "$RESULT" != "successful" || "$REPORTS" -eq 0 ]]; then
            echo "No Release found. Exiting"
            exit 4
          fi
          echo "Release found. Reports downloaded: $REPORTS"
          break
      fi
      echo "Release is still being searched"
      echo "Next check in 5 seconds"
      sleep 5
  done


  ### Use radarr API to periodically check the queue for when the download starts ###
  DOWNLOAD_ID="0"
  RETRY_COUNTER=0 # let's set a limit of 3 retries in case something goes wrong
  while true; do
      RESPONSE=$(curl -s \
        -H "X-Api-Key: $RADARR_API_KEY" \
        "$RADARR_URL/api/v3/queue?page=1&pageSize=100")

      echo "$RESPONSE" | jq

      RADARR_ENTRY=$(jq -c --argjson id "$MOVIE_ID" \
          '.records[] | select(.movieId == $id)' <<< "$RESPONSE")

      if [[ -z $RADARR_ENTRY ]]; then
        ((RETRY_COUNTER+=1))

        if [[ $RETRY_COUNTER -ge 3 ]]; then
          echo "No download found after 3 tries. Something's wrong!"
          echo "Exiting Script"
          exit 99
        fi

        sleep 5
        continue
      else
        echo "Found Download:"
        echo "$RADARR_ENTRY"
        QUEUE_ID=$(echo "$RADARR_ENTRY" | jq -r '.id')

        echo "Found Queue Id: $QUEUE_ID"
        echo "Removing download from radarr Queue so it doesn't get imported by accident. The download still continues on download client"

        curl -s -X DELETE \
            -H "X-Api-Key: $RADARR_API_KEY" \
            "$RADARR_URL/api/v3/queue/$QUEUE_ID?removeFromClient=false"
        DOWNLOAD_ID=$(echo "$RADARR_ENTRY" | jq -r '.downloadId')
        echo "Found Download Id: $DOWNLOAD_ID"

        break
      fi
  done

  ### Use Sabnzbd API to periodically check the queue for when the download is finished ###
  while true; do
      SABNZBD_ENTRY=$(curl -sG "$SABNZBD_URL/api" \
      --data-urlencode "mode=queue" \
      --data-urlencode "output=json" \
      --data-urlencode "apikey=$SABNZBD_API_KEY" \
      --data-urlencode "start=0" \
      --data-urlencode "limit=100" |
      jq --arg id "$DOWNLOAD_ID" '.queue.slots[] | select(.nzo_id == $id)')

      if [[ -z $SABNZBD_ENTRY ]]; then # no download found in Sabnzbd Queue after 3 retries or after the download has started
        echo "Download disappeared from queue. Is it completed?"

        SABNZBD_ENTRY=$(curl -sG "$SABNZBD_URL/api" \
          --data-urlencode "mode=history" \
          --data-urlencode "output=json" \
          --data-urlencode "apikey=$SABNZBD_API_KEY" \
          --data-urlencode "start=0" \
          --data-urlencode "limit=10" |
          jq --arg id "$DOWNLOAD_ID" '.history.slots[] | select(.nzo_id == $id)')

        if [[ ! -z "$SABNZBD_ENTRY" ]]; then # found download in Sabnzbd download history
          echo "Found History Entry"
          echo "$SABNZBD_ENTRY"

          STATUS=$(jq -r '.status' <<< "$SABNZBD_ENTRY")

          if [[ $STATUS == "Completed" ]]; then # download is completed
            echo "Download is completed"
            STORAGE_PATH=$(jq -r '.storage' <<< "$SABNZBD_ENTRY")
            if [[ -d "$STORAGE_PATH" ]]; then
                OUTPUT_PATH="$STORAGE_PATH"
            elif [[ -f "$STORAGE_PATH" ]]; then
                OUTPUT_PATH=$(dirname "$STORAGE_PATH")
            fi
            echo "Found Download Path $OUTPUT_PATH"
            break
          elif [[ $STATUS == "Failed" ]]; then # download failed
            echo "Download failed. Blacklisting release"

            HISTORY=$(curl -s \ # retieve radarr History for this specific release
              -H "X-Api-Key: $RADARR_API_KEY" \
              "$RADARR_URL/api/v3/history?downloadid=$DOWNLOAD_ID&pageSize=50")

            # retrieve history ID
            echo "Retrieving History ID of Download"
            echo "$HISTORY"
            FAILED_HISTORY_ID=$(echo "$HISTORY" | jq -r ' # only grab latest history event
              .records
              | map(select(.eventType == "grabbed"))
              | sort_by(.date)
              | reverse
              | .[0].id
            ')

            # mark release as failed in radarr
            echo "Marking release as failed: $FAILED_HISTORY_ID"
            curl -s -X POST \
              "$RADARR_URL/api/v3/history/failed/$FAILED_HISTORY_ID" \
              -H "X-Api-Key: $RADARR_API_KEY" \
              > /dev/null
            echo "Re-Running Script in 5 seconds to search for a new release"
            sleep 5
            exec "$0" "$@"
          else # download is neither failed nor completed, probably post processing
            ACTION_LINE=$(jq -r '.action_line' <<< "$SABNZBD_ENTRY")
            echo "Waiting for post processing: $ACTION_LINE"
            echo "Next check in 10 seconds"
            sleep 10
            continue
          fi
        else # found no download in Sabnzbd queue nor download history, so where is it? Maybe misconfigured download client
          echo "Found no history entry. Something's wrong!"
          echo "Exiting Script"
          exit 99
        fi
      fi

      echo "Found Sabnzbd Download" # regular download status check
      echo "$SABNZBD_ENTRY"

      DOWNLOAD_STATUS=$(echo "$SABNZBD_ENTRY" | jq -r '.status')

      if [[ true ]]; then # "$DOWNLOAD_STATUS" == "Downloading"
          echo "Sabnzbd download is ongoing: $DOWNLOAD_STATUS"
          TIMELEFT=$(echo "$SABNZBD_ENTRY" | jq -r '.timeleft')

          if [[ "$TIMELEFT" =~ ^[0-9]+:[0-9]{2}:[0-9]{2}$ ]]; then
              IFS=: read -r HOURS MINUTES SECONDS <<< "$TIMELEFT"
              SLEEP_TIME=$((10#$HOURS * 3600 + 10#$MINUTES * 60 + 10#$SECONDS))
          else
              SLEEP_TIME=10
          fi

          if (( SLEEP_TIME > 30 )); then
              SLEEP_TIME=$((SLEEP_TIME / 2))
          elif (( SLEEP_TIME < 10 )); then
              SLEEP_TIME=10
          fi
      fi
      echo "Next Check in $SLEEP_TIME seconds"
      sleep "$SLEEP_TIME"
  done
  echo "radarr download is done. Downloaded file is in $OUTPUT_PATH"

  ### Find video file in newly downloaded file path in case it has more than one video file ###
  echo "Searching for matching episode file"
  AUDIO_SOURCE=""
  mapfile -d '' VIDEO_FILES < <(find "$OUTPUT_PATH" -type f \( \
      -iname "*.mkv" -o \
      -iname "*.mp4" -o \
      -iname "*.avi" -o \
      -iname "*.mov" -o \
      -iname "*.webm" -o \
      -iname "*.m4v" -o \
      -iname "*.ts" \
      \) -print0)

  if [[ "${#VIDEO_FILES[@]}" -eq 0 ]]; then
      echo "No video file found."
      exit 11
  elif [[ "${#VIDEO_FILES[@]}" -eq 1 ]]; then
      AUDIO_SOURCE="${VIDEO_FILES[0]}"
      echo "Found matching file $AUDIO_SOURCE"
  else
      echo "Found multiple video files in Download Path. No clue which one's the right one."
      echo "Exiting script"
      exit 13
  fi


  ### Checking whether new file actually has wanted language ###
  echo "Checking downloaded file for whether it has wanted audio language"
  if [[ "$MISSING_LANG_NAME" == "ger" ]]; then
      AUDIO_LANG="de|deu|ger"
  elif [[ "$MISSING_LANG_NAME" == "eng" ]]; then
      AUDIO_LANG="en|eng"
  else
      echo "Invalid language: $LANG"
      exit 99
  fi

  if ffprobe -v error \
      -select_streams a \
      -show_entries stream_tags=language \
      -of csv=p=0 \
      "$AUDIO_SOURCE" | grep -Eqi "^$AUDIO_LANG$"; then
      echo "File has wanted audio language"
  else ### If File doesn't have the wanted audio language, check for the other language. Although we don\'t want it, we can blacklist this release if it doesn\'t have any viable language (here: neither german or english) ###
      if [[ "$AUDIO_LANG" == "en" ]]; then
        AUDIO_LANG="de|deu|ger"
      else
        AUDIO_LANG="en|eng"
      fi
      echo "File doesn't have wanted audio language"
      if ! ffprobe -v error \
          -select_streams a \
          -show_entries stream_tags=language \
          -of csv=p=0 \
          "$AUDIO_SOURCE" | grep -Eqi "^$AUDIO_LANG$";
      then
          echo "File also doesn't have other viable language. Blacklisting release..."

          HISTORY=$(curl -s \
            -H "X-Api-Key: $RADARR_API_KEY" \
            "$RADARR_URL/api/v3/history?downloadid=$DOWNLOAD_ID&pageSize=50")

          # retrieve history ID
          echo "Retrieving History ID of Download"
          echo "$HISTORY"
          FAILED_HISTORY_ID=$(echo "$HISTORY" | jq -r '
            .records
            | map(select(.eventType == "grabbed"))
            | sort_by(.date)
            | reverse
            | .[0].id
          ')

          # mark release as failed
          echo "Marking release as failed: $FAILED_HISTORY_ID"
          curl -s -X POST \
            "$RADARR_URL/api/v3/history/failed/$FAILED_HISTORY_ID" \
            -H "X-Api-Key: $RADARR_API_KEY" \
            > /dev/null

          echo "Removing downloaded file"
          rm -r "$OUTPUT_PATH"
      fi

      echo "Re-Running Script"
      sleep 5
      exec "$0" "$@"
      exit 99
  fi

  ### Move Found Video File to work Directory ###
  extension="${AUDIO_SOURCE##*.}"
  AUDIO_FILE_PATH="${WORKING_DIR}${MOVIE_ID}_audio.${extension}"

  echo "Moving file to work directory $WORKING_DIR"
  mv "$AUDIO_SOURCE" "$AUDIO_FILE_PATH"
  #cp "$AUDIO_SOURCE" "$AUDIO_FILE_PATH"

  ### Delete file from download directory ###
  echo "Deleting file from download directory"
  rm -r "$OUTPUT_PATH"

  ### Copy Original File to work Directory ###
  echo "Copying original file to work directory"
  cp "$FILE_PATH" "$WORKING_DIR"

  ### Retrieve absolute file path in work directory ###
  FILE_PATH_BASE="$(basename "$FILE_PATH")"
  BASE_FILE_PATH="${WORKING_DIR}$FILE_PATH_BASE"
  OUTPUT_FILE_PATH="${WORKING_DIR}output.mkv"


  ### Start executing python script to calculate audio stream offset and synchronize audio streams of original file and new file ###
  echo "Executing command: process_audio $BASE_FILE_PATH $AUDIO_FILE_PATH  $OUTPUT_FILE_PATH $MISSING_LANG_NAME"
  process_audio "$BASE_FILE_PATH" "$AUDIO_FILE_PATH" "$OUTPUT_FILE_PATH" "$MISSING_LANG_NAME"

  EXIT_CODE=$?

  if [[ $EXIT_CODE -ge 1 ]]; then
      echo "Merging Audio Streams failed in python script. Exiting Script."
      exit 20
  fi


  ### Delete work files to only keep file with merged audio streams ###
  echo "Audios have been merged. Deleting work files..."
  rm "$BASE_FILE_PATH"
  rm "$AUDIO_FILE_PATH"

  if [[ ! -f "$OUTPUT_FILE_PATH" ]]; then
    echo "Merging Audio Streams failed for some reason. Exiting Script."
    exit 21
  fi

  echo "Adjusting permissions of resulting file so radarr can access it"
  chown $PERMS_GROUP "$OUTPUT_FILE_PATH"
  chmod $PERMS_BITS "$OUTPUT_FILE_PATH"

  echo "Renaming resulting file to original file name"
  mv "$OUTPUT_FILE_PATH" "$BASE_FILE_PATH"


  echo "Deleting original file in order to replace it"
  curl -X DELETE \
    "$RADARR_URL/api/v3/moviefile/$MOVIE_FILE_ID" \
    -H "X-Api-Key: $RADARR_API_KEY"

  sleep 2

  echo "Calling radarr API to import resulting file"
  RESPONSE=$(curl -X POST \
    -H "X-Api-Key: $RADARR_API_KEY" \
    -H "Content-Type: application/json" \
    "$RADARR_URL/api/v3/command" \
    -d '{
      "name": "DownloadedMoviesScan",
      "path": "'"$BASE_FILE_PATH"'"
    }')
  COMMAND_ID=$(jq -r '.id' <<< "$RESPONSE")

  while true; do
      RESPONSE=$(curl -s \
        -H "X-Api-Key: $RADARR_API_KEY" \
        "$RADARR_URL/api/v3/command/$COMMAND_ID")
      echo "$RESPONSE"

      STATUS=$(jq -r '.status' <<< "$RESPONSE")
      RESULT=$(jq -r '.result' <<< "$RESPONSE")

      echo "radarr Command Status: $STATUS"

      if [[ "$STATUS" == "completed" ]]; then
        if [[ $RESULT != "successful" ]]; then
          echo "Import failed for some reason. Exiting Script."
          exit 30
        fi
        echo "File import successful. Continuing..."
        break
      fi

      echo "File hasn't been imported yet. Checking again in 5 seconds"
      sleep 5
  done


  sleep 1

  if [[ ! "$RELEASE_GROUP" == "" ]]; then # keep original release group although it's a new file now
    echo "Setting ReleaseGroup of new file to the same of the original file"
    echo "UPDATE MovieFiles SET ReleaseGroup='$RELEASE_GROUP' WHERE MovieId = '$MOVIE_ID'"
    sqlite3 -noheader -batch $RADARR_DB_PATH \
      "UPDATE MovieFiles SET ReleaseGroup='$RELEASE_GROUP' WHERE MovieId = '$MOVIE_ID'"
    sleep 1
  else
    echo "Skipping Release Group Adjustment since it was empty before Script Execution anyway"
  fi

  echo "Making radarr rescan movie to retrieve updated Release Group"
  curl -X POST \
    -H "X-Api-Key: $RADARR_API_KEY" \
    -H "Content-Type: application/json" \
    "$RADARR_URL/api/v3/command" \
    -d "{\"name\":\"RefreshMovie\",\"movieId\":$MOVIE_ID}"

  echo "Script execution is completed."
}
process_audio() {
python3 - "$1" "$2" "$3" "$4" <<'PY'
"""
sync_audio_track.py

Sucht in einer zweiten (Quell-)Datei die Audiospur mit einem bestimmten
Sprach-Tag, extrahiert sie nach AAC (44100 Hz), ermittelt per
Energieverlauf-Kreuzkorrelation die zeitliche Verschiebung gegenüber der
Original-Audiospur der ersten (Ziel-)Datei, fügt die neue Spur
zeitlich synchronisiert in eine Kopie der Zieldatei ein und setzt deren
Disposition auf "nicht default", sodass die Original-Tonspur weiterhin
Default bleibt.

Aufruf:
    python3 sync_audio_track.py ORIGINAL.mkv ZWEITE_DATEI.mkv OUTPUT.mkv de

Abhängigkeiten:
    - ffmpeg / ffprobe im PATH
    - Python-Paket: numpy

Ablauf im Detail:
    1. ffprobe auf ZWEITE_DATEI -> passenden Audiostream per Sprach-Tag finden.
    2. Diesen Stream nach AAC/44100 Hz extrahieren (temporäre Datei).
    3. Aus ORIGINAL die erste/Default-Audiospur ebenfalls als Rohsignal
       extrahieren (zum Vergleich, nicht als Output).
    4. Aus beiden Signalen eine grobe Energiehüllkurve (RMS je Zeitfenster)
       berechnen und per Kreuzkorrelation den zeitlichen Versatz (Offset)
       bestimmen.
    5. Je nach Offsetgröße/Richtung:
         - |offset| <= 5s               -> normaler itsoffset-Delay
         - offset > 5s und Delay > 10s  -> Delay einbauen, aber Ende
                                            auf Videolänge kappen
         - offset < -5s (Linksverschiebung
           um mehr als 5s)              -> Anfang der neuen Spur um
                                            |offset| kappen (synchron ab 0),
                                            die dadurch am ENDE fehlende
                                            Zeit (Dauer |offset|) wird mit
                                            dem entsprechenden Ausschnitt
                                            vom ENDE der Original-Audiospur
                                            aufgefüllt, sodass die Gesamt-
                                            länge der Videolänge entspricht
    6. Neue synchronisierte Spur als zusätzlichen Audiostream in eine Kopie
       von ORIGINAL muxen, disposition:a:<neu> auf 0 setzen, Sprach-Tag
       auf das gesuchte Sprach-Tag setzen.
"""

import argparse
import json
import math
import subprocess
import sys
import tempfile
import os
from pathlib import Path

import numpy as np

# ---------------------------------------------------------------------------
# Konfiguration
# ---------------------------------------------------------------------------

AAC_SAMPLE_RATE = 48000
ANALYSIS_SAMPLE_RATE = 8000       # suffices for energy profile cross correlation, quite quick
WINDOW_MS = 50                    # windows span for energy profile
NORMAL_THRESHOLD_S = 5.0          # threshold for when to fill gap at beginning due to shifting with audio from original audio stream
RIGHT_SHIFT_HARD_LIMIT_S = 10.0   # threshold for when to cut off audio stream at end when shifting it too far to the right
MAX_XCORR_LAG_S = 120.0           # max search span for cross correlation


def log(msg: str) -> None:
    print(f"[sync_audio_track] {msg}", file=sys.stderr, flush=True)


def run(cmd, **kwargs):
    log("Executing: " + " ".join(str(c) for c in cmd))
    return subprocess.run(cmd, check=True, **kwargs)


def run_capture(cmd):
    log("Executing: " + " ".join(str(c) for c in cmd))
    result = subprocess.run(cmd, check=True, capture_output=True, text=True)
    return result.stdout


# ---------------------------------------------------------------------------
# ffprobe-Hilfsfunktionen
# ---------------------------------------------------------------------------

def ffprobe_streams(path: str) -> list:
    out = run_capture([
        "ffprobe", "-v", "error",
        "-print_format", "json",
        "-show_streams",
        str(path),
    ])
    data = json.loads(out)
    return data.get("streams", [])


def ffprobe_format_duration(path: str) -> float:
    out = run_capture([
        "ffprobe", "-v", "error",
        "-print_format", "json",
        "-show_entries", "format=duration",
        str(path),
    ])
    data = json.loads(out)
    dur = data.get("format", {}).get("duration")
    if dur is None:
        raise RuntimeError(f"Could not evaluate duration of {path}")
    return float(dur)


def find_audio_stream_by_language(path: str, lang_tag: str) -> dict:
    """
    Sucht in `path` den Audiostream, dessen Sprach-Tag `lang_tag` entspricht.
    Vergleich case-insensitive. Wirft einen Fehler, wenn keiner gefunden wird.
    """
    streams = ffprobe_streams(path)
    audio_streams = [s for s in streams if s.get("codec_type") == "audio"]

    lang_tag_lower = lang_tag.strip().lower()
    matches = []
    for s in audio_streams:
        tags = s.get("tags", {}) or {}
        lang = (tags.get("language") or tags.get("LANGUAGE") or "").strip().lower()
        if lang == lang_tag_lower:
            matches.append(s)

    if not matches:
        available = [
            (s.get("index"), (s.get("tags", {}) or {}).get("language", "?"))
            for s in audio_streams
        ]
        raise RuntimeError(
            f"No audio stream with '{lang_tag}' found in {path} "
            f"Existing audio streams (index, language): {available}"
        )

    if len(matches) > 1:
        log(
            f"WARNING: Multiple audio streams with wanted language tag '{lang_tag}' found "
            f"(indices: {[m.get('index') for m in matches]}). Picking first one."
        )

    return matches[0]


def find_default_or_first_audio_stream(path: str) -> dict:
    """
    Liefert die Default-Audiospur der Datei, oder falls keine als Default
    markiert ist, die erste Audiospur.
    """
    streams = ffprobe_streams(path)
    audio_streams = [s for s in streams if s.get("codec_type") == "audio"]
    if not audio_streams:
        raise RuntimeError(f"No audio stream found in {path}")

    for s in audio_streams:
        disp = s.get("disposition", {}) or {}
        if disp.get("default") == 1:
            return s

    return audio_streams[0]


# ---------------------------------------------------------------------------
# Audio-Extraktion
# ---------------------------------------------------------------------------

def extract_audio_to_aac(src_path: str, stream_index: int, out_path: str) -> None:
    """
    Extrahiert den Audiostream mit dem absoluten ffmpeg-Streamindex
    `stream_index` aus `src_path` und kodiert ihn nach AAC, 44100 Hz.
    """
    run([
        "ffmpeg", "-y",
        "-i", str(src_path),
        "-map", f"0:{stream_index}",
        "-vn", "-sn", "-dn",
        "-c:a", "aac",
        "-ar", str(AAC_SAMPLE_RATE),
        "-b:a", "192k",
        str(out_path),
    ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def extract_audio_to_wav_mono(src_path: str, stream_index: int, out_path: str,
                               sample_rate: int = ANALYSIS_SAMPLE_RATE) -> None:
    """
    Extrahiert einen Audiostream als mono PCM-WAV in niedriger Samplerate,
    ausschließlich für die Analyse (Kreuzkorrelation) gedacht.
    """
    run([
        "ffmpeg", "-y",
        "-i", str(src_path),
        "-map", f"0:{stream_index}",
        "-vn", "-sn", "-dn",
        "-ac", "1",
        "-ar", str(sample_rate),
        "-c:a", "pcm_s16le",
        str(out_path),
    ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


# ---------------------------------------------------------------------------
# Energieverlauf & Kreuzkorrelation
# ---------------------------------------------------------------------------

def read_wav_mono_pcm16(path: str) -> np.ndarray:
    import wave
    with wave.open(str(path), "rb") as wf:
        n_frames = wf.getnframes()
        raw = wf.readframes(n_frames)
    samples = np.frombuffer(raw, dtype=np.int16).astype(np.float64)
    return samples


def compute_rms_envelope(samples: np.ndarray, sample_rate: int,
                          window_ms: int = WINDOW_MS) -> np.ndarray:
    """
    Berechnet eine RMS-Energiehüllkurve über nicht überlappende Fenster.
    Ein Wert pro Fenster; Fenstergröße in ms.
    """
    window_size = max(1, int(sample_rate * window_ms / 1000.0))
    n_windows = len(samples) // window_size
    if n_windows == 0:
        return np.array([])

    trimmed = samples[: n_windows * window_size]
    reshaped = trimmed.reshape(n_windows, window_size)
    rms = np.sqrt(np.mean(reshaped ** 2, axis=1) + 1e-12)
    return rms


def normalize_envelope(env: np.ndarray) -> np.ndarray:
    env = env - np.mean(env)
    norm = np.linalg.norm(env)
    if norm > 1e-9:
        env = env / norm
    return env


def compute_offset_seconds(env_a: np.ndarray, env_b: np.ndarray,
                            window_ms: int = WINDOW_MS,
                            max_lag_s: float = MAX_XCORR_LAG_S) -> float:
    """
    Bestimmt per Kreuzkorrelation den zeitlichen Versatz zwischen
    Referenz-Hüllkurve `env_a` (Original) und `env_b` (neue/eingefügte Spur).

    Rückgabewert (offset_seconds):
        > 0  -> env_b liegt zeitlich HINTER env_a
                (die neue Spur muss nach RECHTS verschoben/verzögert werden,
                 damit sie zur Originalspur passt)
        < 0  -> env_b liegt zeitlich VOR env_a
                (die neue Spur muss nach LINKS verschoben werden /
                 am Anfang beschnitten werden)
    """
    a = normalize_envelope(env_a)
    b = normalize_envelope(env_b)

    max_lag_windows = int(max_lag_s * 1000.0 / window_ms)

    # np.correlate mit mode='full' liefert alle Verschiebungen von
    # -(len(b)-1) bis +(len(a)-1). Wir beschränken die Suche auf
    # +/- max_lag_windows, um Rechenzeit zu sparen und Fehltreffer bei
    # sehr langen Dateien zu vermeiden.
    full_corr = np.correlate(a, b, mode="full")
    # Index in full_corr, der lag=0 entspricht:
    zero_lag_index = len(b) - 1

    lo = max(0, zero_lag_index - max_lag_windows)
    hi = min(len(full_corr), zero_lag_index + max_lag_windows + 1)

    windowed_corr = full_corr[lo:hi]
    best_index_local = int(np.argmax(windowed_corr))
    best_lag_windows = (lo + best_index_local) - zero_lag_index

    # best_lag_windows ist der Lag, um den b gegenüber a verschoben werden
    # muss (in Fenstereinheiten), damit die Korrelation maximal ist.
    # np.correlate(a, b)[k] entspricht sum_n a[n] * b[n - (k - zero_lag_index)]
    # -> positive best_lag_windows bedeutet: b muss nach RECHTS verschoben
    #    werden (später einsetzen), damit es zu a passt.
    offset_seconds = best_lag_windows * (window_ms / 1000.0)
    return offset_seconds


def determine_offset(original_path: str, original_stream_index: int,
                      extracted_aac_path: str, tmp_dir: str) -> float:
    """
    Ermittelt die zeitliche Differenz (in Sekunden) zwischen der
    Original-Audiospur und der frisch extrahierten AAC-Spur.
    """
    wav_orig = os.path.join(tmp_dir, "orig_analysis.wav")
    wav_new = os.path.join(tmp_dir, "new_analysis.wav")

    extract_audio_to_wav_mono(original_path, original_stream_index, wav_orig)
    # Extrahierte AAC-Datei hat nur einen Audiostream -> Index 0
    extract_audio_to_wav_mono(extracted_aac_path, 0, wav_new)

    samples_orig = read_wav_mono_pcm16(wav_orig)
    samples_new = read_wav_mono_pcm16(wav_new)

    env_orig = compute_rms_envelope(samples_orig, ANALYSIS_SAMPLE_RATE)
    env_new = compute_rms_envelope(samples_new, ANALYSIS_SAMPLE_RATE)

    if env_orig.size == 0 or env_new.size == 0:
        raise RuntimeError("Couldn't calculate energy profile (maybe empty audio signal?)")

    offset_seconds = compute_offset_seconds(env_orig, env_new)
    return offset_seconds


# ---------------------------------------------------------------------------
# Zusammenbau der finalen Datei
# ---------------------------------------------------------------------------

def build_filled_start_audio(original_path: str, original_audio_index: int,
                              extracted_aac_path: str, trim_duration_s: float,
                              video_duration_s: float, tmp_dir: str) -> str:
    """
    Sonderfall Linksverschiebung (> 5s):
    Die neue Spur liegt zeitlich VOR der Originalspur, d.h. ihr Anfang
    (Dauer `trim_duration_s`) gehört eigentlich vor Timestamp 0 des Videos
    und wird verworfen, damit die neue Spur synchron ab 0 beginnt. Dadurch
    wird sie um `trim_duration_s` Sekunden kürzer als die Videolänge.
    Diese fehlende Zeit entsteht am ENDE und wird mit dem entsprechenden
    Ausschnitt vom Ende der Original-Audiospur aufgefüllt (den letzten
    `trim_duration_s` Sekunden der Originalspur), sodass die resultierende
    Spur wieder genau `video_duration_s` lang ist.

    Ablauf:
        1. Anfang der neuen Spur um `trim_duration_s` kappen (-ss).
        2. Die letzten `trim_duration_s` Sekunden der Original-Audiospur
           extrahieren (Ausschnitt von video_duration_s - trim_duration_s
           bis video_duration_s).
        3. Beide Teile aneinanderhängen: [getrimmte neue Spur] + [Original-
           Ausschnitt vom Ende].

    Rückgabe: Pfad zur zusammengesetzten AAC-Datei.
    """
    # 1. Anfang der neuen Spur verwerfen, damit sie synchron bei 0 beginnt.
    trimmed_new_path = os.path.join(tmp_dir, "trimmed_new_for_fill.m4a")
    run([
        "ffmpeg", "-y",
        "-i", str(extracted_aac_path),
        "-ss", f"{trim_duration_s:.3f}",
        "-c", "copy",
        str(trimmed_new_path),
    ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    # 2. Die dadurch am Ende fehlende Zeit aus dem ENDE der Original-
    #    Audiospur holen.
    fill_start_s = max(0.0, video_duration_s - trim_duration_s)
    fill_part_path = os.path.join(tmp_dir, "fill_part_from_original_end.m4a")
    run([
        "ffmpeg", "-y",
        "-i", str(original_path),
        "-map", f"0:{original_audio_index}",
        "-ss", f"{fill_start_s:.3f}",
        "-vn", "-sn", "-dn",
        "-c:a", "aac",
        "-ar", str(AAC_SAMPLE_RATE),
        "-b:a", "192k",
        str(fill_part_path),
    ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    # 3. Zusammenfügen per concat-Filter (statt concat-Demuxer), damit
    #    unterschiedliche Encoder-Zustände/Timestamps der beiden Zuspieler
    #    sauber neu encodiert und nahtlos aneinandergereiht werden.
    combined_path = os.path.join(tmp_dir, "combined_end_filled.m4a")
    run([
        "ffmpeg", "-y",
        "-i", str(trimmed_new_path),
        "-i", str(fill_part_path),
        "-filter_complex", "[0:a][1:a]concat=n=2:v=0:a=1[aout]",
        "-map", "[aout]",
        "-c:a", "aac",
        "-ar", str(AAC_SAMPLE_RATE),
        "-b:a", "192k",
        str(combined_path),
    ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    return combined_path


def mux_synced_audio(original_path: str, original_audio_index: int,
                      extracted_aac_path: str, output_path: str,
                      offset_seconds: float, video_duration_s: float,
                      language_tag: str, tmp_dir: str) -> None:
    """
    Fügt die extrahierte (ggf. angepasste) AAC-Spur synchronisiert in eine
    Kopie von `original_path` ein und schreibt das Ergebnis nach
    `output_path`. Alle bestehenden Streams von `original_path` werden
    unverändert übernommen (Kopie), die neue Audiospur wird angehängt.

    `original_audio_index` ist der ffmpeg-Streamindex der Original-
    Audiospur in `original_path` (wird im Sonderfall Linksverschiebung
    benötigt, um den Füllteil zu extrahieren).
    `language_tag` wird als Sprach-Metadatum auf die neu eingefügte Spur
    gesetzt (das ursprünglich per CLI gesuchte Sprach-Tag).
    """
    abs_offset = abs(offset_seconds)

    # Bestimme Anzahl vorhandener Audiostreams in original_path,
    # damit wir wissen, welchen disposition-Index die neue Spur bekommt.
    orig_streams = ffprobe_streams(original_path)
    n_existing_audio = sum(1 for s in orig_streams if s.get("codec_type") == "audio")
    new_audio_index = n_existing_audio  # 0-basiert, neue Spur ist die (n+1)-te

    input_args = ["-i", str(original_path)]
    extra_input_path = extracted_aac_path
    itsoffset_value = None
    trim_start_s = 0.0

    if abs_offset <= NORMAL_THRESHOLD_S:
        # Normalfall: minimaler oder kein Versatz -> ggf. kleiner Delay,
        # kein Abschneiden nötig.
        if offset_seconds > 0:
            itsoffset_value = offset_seconds
        elif offset_seconds < 0:
            # Kleiner Linksversatz (<= 5s per Definition) -> Anfang leicht kappen
            trim_start_s = abs_offset
        log(f"Default Case: |offset|={abs_offset:.3f}s <= {NORMAL_THRESHOLD_S}s")

    elif offset_seconds > NORMAL_THRESHOLD_S:
        # Rechtsverschiebung > 5s
        if offset_seconds > RIGHT_SHIFT_HARD_LIMIT_S:
            log(
                f"Special Case right shift: offset={offset_seconds:.3f}s "
                f"> {RIGHT_SHIFT_HARD_LIMIT_S}s -> building in delay, cutting off end"
            )
        else:
            log(
                f"Right shift: offset={offset_seconds:.3f}s "
                f"(between {NORMAL_THRESHOLD_S}s and {RIGHT_SHIFT_HARD_LIMIT_S}s)"
            )
        itsoffset_value = offset_seconds

    else:
        # offset_seconds < -NORMAL_THRESHOLD_S: Linksverschiebung > 5s.
        # Anfang der neuen Spur wird verworfen (wie im Normalfall), die
        # dadurch am ENDE fehlende Zeit wird mit dem entsprechenden
        # Ausschnitt vom Ende der Original-Audiospur aufgefüllt, sodass
        # die Gesamtlänge der resultierenden Spur der Videolänge entspricht.
        log(
            f"Special Case left shift: offset={offset_seconds:.3f}s "
            f"< -{NORMAL_THRESHOLD_S}s -> Start of new stream will be cut off by "
            f"{abs_offset:.3f}s , missing segment at end will be filled with "
            f"original audio"
        )
        extra_input_path = build_filled_start_audio(
            original_path, original_audio_index, extracted_aac_path,
            abs_offset, video_duration_s, tmp_dir
        )
        # Die zusammengesetzte Spur beginnt bereits korrekt bei 0, ist
        # synchron und hat bereits die volle Videolänge -> kein
        # trim_start_s, kein itsoffset_value nötig.

    # Falls der Anfang beschnitten werden muss (nur im Normalfall bei
    # kleinem Linksversatz), erst eine getrimmte Zwischen-Datei erzeugen -
    # einfacher und robuster, als -ss direkt im finalen Mux-Kommando mit
    # -itsoffset zu kombinieren.
    if trim_start_s > 0:
        trimmed_path = os.path.join(tmp_dir, "trimmed_new_audio.m4a")
        run([
            "ffmpeg", "-y",
            "-i", str(extracted_aac_path),
            "-ss", f"{trim_start_s:.3f}",
            "-c", "copy",
            str(trimmed_path),
        ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        extra_input_path = trimmed_path

    second_input_args = []
    if itsoffset_value is not None and itsoffset_value > 0:
        second_input_args += ["-itsoffset", f"{itsoffset_value:.3f}"]
    second_input_args += ["-i", str(extra_input_path)]

    cmd = ["ffmpeg", "-y"] + input_args + second_input_args + [
        "-map", "0",
        "-map", "1:a:0",
        "-c", "copy",
        f"-disposition:a:{new_audio_index}", "0",
        f"-metadata:s:a:{new_audio_index}", f"language={language_tag}",
    ]

    # Sonderfall Rechtsverschiebung > 10s: Ende auf Videolänge kappen.
    # -t begrenzt die Gesamtlänge des Outputs auf die Videodauer, sodass
    # der Überstand am Ende der (verzögerten) neuen Audiospur abgeschnitten wird.
    if itsoffset_value is not None and itsoffset_value > RIGHT_SHIFT_HARD_LIMIT_S:
        cmd += ["-t", f"{video_duration_s:.3f}"]

    cmd += [str(output_path)]

    run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

    cmd = ["mkvpropedit", "--add-track-statistics-tags", output_path]
    run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


# ---------------------------------------------------------------------------
# Hauptprogramm
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Extrahiert eine Audiospur per Sprach-Tag aus einer zweiten "
                    "Datei, synchronisiert sie per Energieverlauf-Analyse mit der "
                    "Originalspur und fügt sie synchronisiert in das Originalvideo ein."
    )
    parser.add_argument("original", help="Pfad zur Originaldatei (Video)")
    parser.add_argument("second_file", help="Pfad zur zweiten Datei, aus der die Audiospur extrahiert wird")
    parser.add_argument("output", help="Pfad zur Ausgabedatei")
    parser.add_argument("language_tag", help="Gesuchtes Sprach-Tag, z. B. 'de', 'eng', 'ja'")
    args = parser.parse_args()

    original_path = Path(args.original)
    second_path = Path(args.second_file)
    output_path = Path(args.output)

    for p in (original_path, second_path):
        if not p.is_file():
            log(f"Error: File not found: {p}")
            sys.exit(1)

    with tempfile.TemporaryDirectory(prefix="sync_audio_") as tmp_dir:
        log(f"Searching audio stream with language tag '{args.language_tag}' in {second_path}")
        target_stream = find_audio_stream_by_language(str(second_path), args.language_tag)
        target_index = target_stream["index"]
        log(f"Gefunden: Stream-Index {target_index} "
            f"(Codec: {target_stream.get('codec_name')}, "
            f"Language: {(target_stream.get('tags', {}) or {}).get('language')})")

        extracted_aac_path = os.path.join(tmp_dir, "extracted_audio.m4a")
        log(f"Extracting audio stream to AAC ({AAC_SAMPLE_RATE} Hz) -> {extracted_aac_path}")
        extract_audio_to_aac(str(second_path), target_index, extracted_aac_path)

        log("Calculating default/first audio stream of original file for comparison")
        original_audio_stream = find_default_or_first_audio_stream(str(original_path))
        original_audio_index = original_audio_stream["index"]
        log(f"Original audio stream: Stream-Index {original_audio_index}")

        log("Calculating time shift using energy profile cross correlation")
        offset_seconds = determine_offset(
            str(original_path), original_audio_index, extracted_aac_path, tmp_dir
        )
        log(f"Calculating time shift (Offset): {offset_seconds:.3f} s "
            f"({'new stream later' if offset_seconds > 0 else 'new stream earlier' if offset_seconds < 0 else 'no shift'})")

        video_duration_s = ffprobe_format_duration(str(original_path))
        log(f"Video length (Original): {video_duration_s:.3f} s")

        log(f"Muxing synchronized audio stream in output file -> {output_path}")
        mux_synced_audio(
            str(original_path), original_audio_index, extracted_aac_path,
            str(output_path), offset_seconds, video_duration_s,
            args.language_tag, tmp_dir
        )

    log("Done muxing new file.")


if __name__ == "__main__":
    main()
PY
}
main "$@"