#!/usr/bin/env bash

set -euo pipefail

sink_name="${REFORGER_SURROUND_SINK_NAME:-reforger_surround_51}"
sink_description="${REFORGER_SURROUND_DESCRIPTION:-Reforger Surround 5.1}"
app_name="${REFORGER_SURROUND_APP_NAME:-Arma Reforger}"
steam_uri="${REFORGER_SURROUND_STEAM_URI:-steam://rungameid/1874880}"
wait_seconds="${REFORGER_SURROUND_WAIT_SECONDS:-90}"
state_dir="${XDG_STATE_HOME:-$HOME/.local/state}/reforger-surround"
module_id_file="${state_dir}/module-id"
previous_default_sink_file="${state_dir}/previous-default-sink"

usage() {
  cat <<'EOF'
Usage:
  reforger-surround create
  reforger-surround up
  reforger-surround launch [command...]
  reforger-surround expose
  reforger-surround pin
  reforger-surround restore
  reforger-surround down
  reforger-surround status
  reforger-surround check

Commands:
  create   Create the temporary 5.1 sink without changing the default sink.
  up       Create the 5.1 sink, save the current default, and make it default.
  launch   Run up, launch Reforger, wait for its stream, then restore.
  expose   Create the 5.1 sink and move a running Reforger stream into it.
  pin      Alias for expose.
  restore  Restore the saved previous default sink, then expose Reforger if running.
  down     Restore the saved previous default sink and unload the 5.1 sink.
  status   Show default sink, virtual sink state, and detected Reforger stream.
  check    Report whether a running Reforger stream negotiated stereo or surround.

Environment:
  REFORGER_SURROUND_SINK_NAME     Override sink name. Default: reforger_surround_51
  REFORGER_SURROUND_DESCRIPTION   Override visible sink name. Default: Reforger Surround 5.1
  REFORGER_SURROUND_APP_NAME      Override app name. Default: Arma Reforger
  REFORGER_SURROUND_STEAM_URI     Override Steam URI. Default: steam://rungameid/1874880
  REFORGER_SURROUND_WAIT_SECONDS  Override launch wait. Default: 90
EOF
}

require_command() {
  local command="$1"

  if ! command -v "${command}" >/dev/null 2>&1; then
    printf 'Missing required command: %s\n' "${command}" >&2
    exit 127
  fi
}

ensure_dependencies() {
  require_command pactl
  require_command awk
  require_command sed
}

require_pulse_connection() {
  if ! pactl info >/dev/null 2>&1; then
    printf 'Could not connect to PulseAudio/PipeWire through pactl.\n' >&2
    printf 'Check that PipeWire and pipewire-pulse are running for this user session.\n' >&2
    exit 1
  fi
}

ensure_state_dir() {
  mkdir -p "${state_dir}"
}

current_default_sink() {
  pactl get-default-sink 2>/dev/null || true
}

sink_exists() {
  pactl list sinks short | awk -v sink_name="${sink_name}" '$2 == sink_name { found = 1 } END { exit !found }'
}

find_module_id() {
  pactl list short modules | awk -v sink_name="${sink_name}" '
    $2 == "module-null-sink" && index($0, "sink_name=" sink_name) {
      print $1
      exit
    }
  '
}

stored_module_id() {
  [[ -s "${module_id_file}" ]] && sed -n '1p' "${module_id_file}"
}

create_sink() {
  ensure_state_dir

  local module_id
  module_id="$(find_module_id || true)"
  if [[ -n "${module_id}" ]]; then
    printf '%s\n' "${module_id}" > "${module_id_file}"
    printf 'Using existing 5.1 sink: %s (module %s)\n' "${sink_name}" "${module_id}"
    return 0
  fi

  if sink_exists; then
    printf '5.1 sink already exists: %s\n' "${sink_name}"
    return 0
  fi

  module_id="$(pactl load-module module-null-sink \
    "sink_name=${sink_name}" \
    "sink_properties=device.description=\"${sink_description}\"" \
    "channels=6" \
    "channel_map=front-left,front-right,front-center,lfe,rear-left,rear-right" \
    "rate=48000")"

  printf '%s\n' "${module_id}" > "${module_id_file}"
  printf 'Created 5.1 sink: %s (module %s)\n' "${sink_name}" "${module_id}"
}

save_previous_default() {
  ensure_state_dir

  local default_sink
  default_sink="$(current_default_sink)"
  if [[ -n "${default_sink}" && "${default_sink}" != "${sink_name}" ]]; then
    printf '%s\n' "${default_sink}" > "${previous_default_sink_file}"
    printf 'Saved previous default sink: %s\n' "${default_sink}"
  fi
}

set_sink_default() {
  pactl set-default-sink "${sink_name}"
  printf 'Set launch default sink: %s\n' "${sink_name}"
}

prepare_launch_default() {
  create_sink
  save_previous_default
  set_sink_default
}

up() {
  prepare_launch_default
  cat <<EOF

Launch Arma Reforger now. After the game has created its audio stream, run:
  reforger-surround restore
  reforger-surround check
EOF
}

start_game() {
  if [[ "$#" -gt 0 ]]; then
    "$@" &
    printf 'Launched command: %s\n' "$*"
    return 0
  fi

  require_command steam
  steam "${steam_uri}" >/tmp/reforger-surround-steam.log 2>&1 &
  printf 'Requested Steam launch: %s\n' "${steam_uri}"
}

wait_for_app_stream() {
  local input_id elapsed

  for ((elapsed=0; elapsed<wait_seconds; elapsed++)); do
    input_id="$(find_app_sink_input_id || true)"
    if [[ -n "${input_id}" ]]; then
      printf '%s\n' "${input_id}"
      return 0
    fi
    sleep 1
  done

  return 1
}

launch_game() {
  local input_id

  prepare_launch_default
  start_game "$@"

  if ! input_id="$(wait_for_app_stream)"; then
    printf 'Timed out waiting for application.name="%s".\n' "${app_name}" >&2
    printf 'Leaving %s as default so late stream creation can still negotiate 5.1.\n' "${sink_name}" >&2
    printf 'Run `reforger-surround restore` after the game audio stream appears.\n' >&2
    return 1
  fi

  printf 'Detected %s sink-input #%s\n' "${app_name}" "${input_id}"
  restore_with_pin
  print_check
}

expose_app() {
  local input_id sample_spec sink_after

  create_sink
  input_id="$(find_app_sink_input_id || true)"

  if [[ -z "${input_id}" ]]; then
    printf 'No running sink-input found for application.name="%s".\n' "${app_name}" >&2
    return 1
  fi

  sample_spec="$(app_stream_sample_spec "${input_id}")"
  pactl move-sink-input "${input_id}" "${sink_name}"
  sink_after="$(app_stream_sink_name "${input_id}" || true)"
  printf 'Exposed %s sink-input #%s to %s\n' "${app_name}" "${input_id}" "${sink_name}"
  printf 'Current stream sink: %s\n' "${sink_after:-unknown}"

  case "${sample_spec}" in
    *" 6ch "*)
      printf 'Stream is 6ch; Reforger should now stay on the 5.1 sink.\n'
      ;;
    *" 2ch "*)
      printf 'Warning: stream is already stereo-only. Pinning will not make Reforger renegotiate surround.\n' >&2
      ;;
    *)
      printf 'Warning: unknown stream channel count; run reforger-surround check.\n' >&2
      ;;
  esac
}

pin_app_if_running() {
  local input_id
  input_id="$(find_app_sink_input_id || true)"

  if [[ -n "${input_id}" ]]; then
    expose_app
  else
    printf 'No %s stream found to pin; restoring only the system default.\n' "${app_name}"
  fi
}

restore_default() {
  local restore_sink
  restore_sink="$(sed -n '1p' "${previous_default_sink_file}" 2>/dev/null || true)"

  if [[ -z "${restore_sink}" ]]; then
    printf 'No saved previous default sink found; nothing to restore.\n'
    return 0
  fi

  if ! pactl list sinks short | awk -v sink_name="${restore_sink}" '$2 == sink_name { found = 1 } END { exit !found }'; then
    printf 'Saved previous default sink no longer exists: %s\n' "${restore_sink}" >&2
    return 1
  fi

  pactl set-default-sink "${restore_sink}"
  printf 'Restored default sink: %s\n' "${restore_sink}"
}

restore_with_pin() {
  restore_default || true
  pin_app_if_running
}

unload_sink() {
  local module_id
  module_id="$(find_module_id || true)"
  if [[ -z "${module_id}" ]]; then
    module_id="$(stored_module_id || true)"
  fi

  if [[ -z "${module_id}" ]]; then
    printf 'No %s module found to unload.\n' "${sink_name}"
    return 0
  fi

  pactl unload-module "${module_id}" 2>/dev/null || true
  rm -f "${module_id_file}"
  printf 'Unloaded 5.1 sink module: %s\n' "${module_id}"
}

down() {
  restore_default || true
  unload_sink
}

find_app_sink_input_id() {
  pactl list sink-inputs | awk -v app_name="${app_name}" '
    /^Sink Input #/ {
      id = substr($3, 2)
    }
    index($0, "application.name = \"" app_name "\"") {
      print id
      exit
    }
  '
}

app_stream_sink_name() {
  local input_id="$1"

  pactl list sink-inputs | awk -v target="Sink Input #${input_id}" '
    $0 == target {
      in_input = 1
      next
    }
    in_input && $1 == "Sink:" {
      print $2
      exit
    }
    in_input && /^Sink Input #/ {
      exit
    }
  ' | while read -r sink_id; do
    pactl list sinks short | awk -v sink_id="${sink_id}" '$1 == sink_id { print $2; exit }'
  done
}

app_stream_lines() {
  local input_id="$1"

  pactl list sink-inputs | awk -v target="Sink Input #${input_id}" '
    $0 == target {
      in_input = 1
      next
    }
    in_input && ($1 == "Sink:" || ($1 == "Sample" && $2 == "Specification:") || ($1 == "Channel" && $2 == "Map:")) {
      print
    }
    in_input && /^Sink Input #/ {
      exit
    }
  '
}

app_stream_sample_spec() {
  local input_id="$1"

  pactl list sink-inputs | awk -v target="Sink Input #${input_id}" '
    $0 == target {
      in_input = 1
      next
    }
    in_input && $1 == "Sample" && $2 == "Specification:" {
      print $0
      exit
    }
    in_input && /^Sink Input #/ {
      exit
    }
  '
}

app_stream_channel_map() {
  local input_id="$1"

  pactl list sink-inputs | awk -v target="Sink Input #${input_id}" '
    $0 == target {
      in_input = 1
      next
    }
    in_input && $1 == "Channel" && $2 == "Map:" {
      print $0
      exit
    }
    in_input && /^Sink Input #/ {
      exit
    }
  '
}

print_check() {
  local input_id sample_spec channel_map sink
  input_id="$(find_app_sink_input_id || true)"

  if [[ -z "${input_id}" ]]; then
    printf 'No running sink-input found for application.name="%s".\n' "${app_name}" >&2
    return 1
  fi

  sample_spec="$(app_stream_sample_spec "${input_id}")"
  channel_map="$(app_stream_channel_map "${input_id}")"
  sink="$(app_stream_sink_name "${input_id}" || true)"

  printf 'Detected %s sink-input #%s\n' "${app_name}" "${input_id}"
  printf 'Sink: %s\n' "${sink:-unknown}"
  printf '%s\n' "${sample_spec}"
  printf '%s\n' "${channel_map}"

  case "${sample_spec}" in
    *" 6ch "*)
      printf 'Result: success, Reforger negotiated surround output.\n'
      printf 'Expected useful channels: FL, FR, FC, RL, RR. LFE may be silent.\n'
      ;;
    *" 2ch "*)
      printf 'Result: stereo-only. Restart Reforger with the 5.1 sink selected/default before launch.\n'
      return 2
      ;;
    *)
      printf 'Result: unknown channel count; inspect the sample specification and channel map above.\n'
      return 3
      ;;
  esac
}

print_status() {
  local module_id input_id
  module_id="$(find_module_id || true)"
  input_id="$(find_app_sink_input_id || true)"

  printf 'default sink: %s\n' "$(current_default_sink)"
  printf '5.1 sink: %s\n' "${sink_name}"
  printf '5.1 module: %s\n' "${module_id:-not loaded}"
  printf 'saved previous default: %s\n' "$(sed -n '1p' "${previous_default_sink_file}" 2>/dev/null || printf 'none')"

  if sink_exists; then
    pactl list sinks short | awk -v sink_name="${sink_name}" '$2 == sink_name { print }'
  fi

  if [[ -z "${input_id}" ]]; then
    printf '%s stream: not found\n' "${app_name}"
  else
    printf '%s stream: sink-input #%s\n' "${app_name}" "${input_id}"
    app_stream_lines "${input_id}"
  fi
}

main() {
  local command="${1:-}"
  if [[ -z "${command}" ]]; then
    usage >&2
    exit 1
  fi
  shift || true

  case "${command}" in
    -h|--help|help)
      usage
      ;;
    create)
      ensure_dependencies
      require_pulse_connection
      create_sink
      ;;
    up)
      ensure_dependencies
      require_pulse_connection
      up
      ;;
    launch)
      ensure_dependencies
      require_pulse_connection
      launch_game "$@"
      ;;
    expose)
      ensure_dependencies
      require_pulse_connection
      expose_app
      ;;
    pin)
      ensure_dependencies
      require_pulse_connection
      expose_app
      ;;
    restore)
      ensure_dependencies
      require_pulse_connection
      restore_with_pin
      ;;
    down)
      ensure_dependencies
      require_pulse_connection
      down
      ;;
    status)
      ensure_dependencies
      require_pulse_connection
      print_status
      ;;
    check)
      ensure_dependencies
      require_pulse_connection
      print_check
      ;;
    *)
      usage >&2
      exit 1
      ;;
  esac
}

main "$@"