#!/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}"
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 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.
  restore  Restore the saved previous default sink, leaving the 5.1 sink alive.
  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
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}"
}

up() {
  create_sink
  save_previous_default
  set_sink_default
  cat <<EOF

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

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}"
}

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

  case "${command}" in
    -h|--help|help)
      usage
      ;;
    create)
      ensure_dependencies
      require_pulse_connection
      create_sink
      ;;
    up)
      ensure_dependencies
      require_pulse_connection
      up
      ;;
    restore)
      ensure_dependencies
      require_pulse_connection
      restore_default
      ;;
    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 "$@"