#!/bin/bash

# omarchy:summary=Restart audio services and recover stuck USB audio devices.
# omarchy:examples=omarchy restart audio | omarchy-restart-audio

services=(wireplumber.service pipewire.service pipewire-pulse.service)

restart_audio_services() {
  echo -e "Restarting audio services...\n"

  if timeout 25s systemctl --user restart "${services[@]}"; then
    return 0
  fi

  echo -e "\nAudio services did not restart cleanly. Forcing stuck services down...\n"
  systemctl --user cancel >/dev/null 2>&1 || true
  systemctl --user kill --kill-whom=all --signal=KILL "${services[@]}" >/dev/null 2>&1 || true
  systemctl --user reset-failed "${services[@]}" >/dev/null 2>&1 || true

  timeout 25s systemctl --user start pipewire.service pipewire-pulse.service wireplumber.service
}

wpctl_healthy() {
  timeout 5s wpctl status >/dev/null 2>&1
}

usb_device_for_card() {
  local card="$1"
  local path

  path=$(readlink -f "/sys/class/sound/card$card" 2>/dev/null) || return 1

  while [[ $path != "/" ]]; do
    if [[ -f $path/idVendor && -f $path/idProduct && -f $path/busnum && -f $path/devnum ]]; then
      basename "$path"
      return 0
    fi

    path=$(dirname "$path")
  done

  return 1
}

append_unique() {
  local value="$1"
  shift
  local existing

  [[ -z $value ]] && return

  for existing in "$@"; do
    [[ $existing == "$value" ]] && return
  done

  printf '%s\n' "$value"
}

audio_process_blocked_on_usb() {
  local pid
  local name
  local wchan

  for pid in /proc/[0-9]*; do
    [[ -r $pid/comm && -r $pid/wchan ]] || continue

    name=$(cat "$pid/comm" 2>/dev/null || true)
    [[ $name == "wireplumber" || $name == "pipewire" ]] || continue

    wchan=$(cat "$pid/wchan" 2>/dev/null || true)
    [[ $wchan == usb_* || $wchan == *usb* ]] && return 0
  done

  return 1
}

usb_audio_devices() {
  local card
  local device
  local devices=()
  local unique

  for card in /sys/class/sound/card*; do
    [[ -e $card && $card =~ /card([0-9]+)$ ]] || continue

    device=$(usb_device_for_card "${BASH_REMATCH[1]}" || true)
    unique=$(append_unique "$device" "${devices[@]}")
    [[ -n $unique ]] && devices+=("$unique")
  done

  printf '%s\n' "${devices[@]}"
}

clear_usb_audio_defaults() {
  local state_dir="${XDG_STATE_HOME:-$HOME/.local/state}/wireplumber"

  [[ -d $state_dir ]] || return 0

  if [[ -f $state_dir/default-nodes ]]; then
    cp "$state_dir/default-nodes" "$state_dir/default-nodes.bak.$(date +%s)" 2>/dev/null || true
    sed -i \
      -e '/^default\.configured\.audio\.sink=alsa_output\.usb-/d' \
      -e '/^default\.configured\.audio\.sink\.[0-9]\+=alsa_output\.usb-/d' \
      "$state_dir/default-nodes" 2>/dev/null || true
  fi

  if [[ -f $state_dir/default-routes ]]; then
    cp "$state_dir/default-routes" "$state_dir/default-routes.bak.$(date +%s)" 2>/dev/null || true
    sed -i '/^alsa_card\.usb-/d' "$state_dir/default-routes" 2>/dev/null || true
  fi
}

stuck_usb_audio_devices() {
  local status
  local state
  local owner_pid
  local owner_name
  local card
  local device
  local devices=()
  local unique

  for status in /proc/asound/card*/pcm*p/sub*/status; do
    [[ -e $status ]] || continue

    state=$(awk '/^state:/ { print $2; exit }' "$status")
    [[ $state == "SETUP" ]] || continue

    owner_pid=$(awk '/^owner_pid[[:space:]]*:/ { print $3; exit }' "$status")
    [[ -n $owner_pid && -r /proc/$owner_pid/comm ]] || continue

    owner_name=$(cat "/proc/$owner_pid/comm" 2>/dev/null || true)
    [[ $owner_name == "wireplumber" || $owner_name == "pipewire" ]] || continue
    [[ $status =~ /card([0-9]+)/ ]] || continue

    card="${BASH_REMATCH[1]}"
    device=$(usb_device_for_card "$card" || true)
    unique=$(append_unique "$device" "${devices[@]}")
    [[ -n $unique ]] && devices+=("$unique")
  done

  printf '%s\n' "${devices[@]}"
}

reset_usb_audio_device() {
  local device="$1"
  local sysfs="/sys/bus/usb/devices/$device"
  local busnum
  local devnum
  local bus_device
  local product

  [[ -d $sysfs ]] || return 1

  if ! omarchy-cmd-present usbreset; then
    echo "usbreset is not installed; replug or power-cycle $device to recover it."
    return 1
  fi

  busnum=$(cat "$sysfs/busnum" 2>/dev/null) || return 1
  devnum=$(cat "$sysfs/devnum" 2>/dev/null) || return 1
  product=$(cat "$sysfs/product" 2>/dev/null || echo "$device")
  bus_device=$(printf '%03d/%03d' "$busnum" "$devnum")

  echo -e "\nResetting stuck USB audio device: $product ($bus_device)..."

  if timeout 25s sudo usbreset "$bus_device"; then
    return 0
  fi

  echo "Could not reset $product automatically. Replug or power-cycle it, then run omarchy-restart-audio again."
  return 1
}

recover_stuck_usb_audio() {
  local devices=()
  local device
  local detected=()
  local reset_count=0

  while IFS= read -r device; do
    [[ -n $device ]] && detected+=("$device")
  done < <(stuck_usb_audio_devices)

  if (( ${#detected[@]} == 0 )) && audio_process_blocked_on_usb; then
    echo "Audio service is blocked in USB I/O; resetting USB audio devices..."
    clear_usb_audio_defaults
    while IFS= read -r device; do
      [[ -n $device ]] && detected+=("$device")
    done < <(usb_audio_devices)
  fi

  for device in "${detected[@]}"; do
    [[ -n $(append_unique "$device" "${devices[@]}") ]] && devices+=("$device")
  done

  (( ${#devices[@]} > 0 )) || return 1

  for device in "${devices[@]}"; do
    reset_usb_audio_device "$device" && reset_count=$((reset_count + 1))
  done

  (( reset_count > 0 ))
}

restart_audio_services || true

if ! wpctl_healthy; then
  echo -e "\nPipeWire is still not responding. Checking for stuck USB audio devices..."

  if recover_stuck_usb_audio; then
    echo -e "\nRestarting audio services after USB audio reset...\n"
    restart_audio_services || true
  else
    echo "No stuck USB audio device was detected automatically."
  fi
fi

sleep 2

echo -e "\nAudio status:\n"
if ! timeout 5s wpctl status; then
  echo "Audio services are still not responding. Try replugging or power-cycling the selected USB audio device."
  exit 1
fi