#!/usr/bin/env bash
set -u

# -------------------------------------------------
# CONFIGURATION
# -------------------------------------------------
WIFI_5G="cintra-5g"
WIFI_24G="cintra-5g"

declare -A BT_DEVICES
BT_DEVICES[TAH8506]="00:1E:7C:CA:77:00"
BT_DEVICES[SHB7150]="00:1E:7C:30:F3:69"
BT_DEVICES[HD450BT]="00:16:94:33:91:51"
BT_DEVICES[XTREME]="41:42:5D:FA:26:0A" # BT Speaker
BT_DEVICES[EM03]="F1:5A:D5:16:57:9D"   # BT Speaker/powerbank
BT_DEVICES[X5]="47:0D:FF:C9:62:E1"     # Bone Headphones thinkplus-X5

BT_AUTO_DEVICES=("TAH8506" "SHB7150")

BT_SERVICE="bluetooth"
BT_RETRIES=6
BT_DELAY=3

# Polkit rule that lets the current user start/stop bluetooth.service
# without a password, so this script runs sudo-free. Installed on first
# run via ensure_polkit_rule().
BT_POLKIT_RULE_FILE="/etc/polkit-1/rules.d/49-btswitch.rules"
# Sentinel file to track that the rule was installed, since the user
# may not have read/execute access to the polkit rules directory.
BT_POLKIT_SENTINEL="$HOME/.config/btswitch/.polkit_rule_installed"

# -------------------------------------------------
# HELPERS
# -------------------------------------------------

wifi_switch_needed() {
  [[ -n "$1" && -n "$2" && "$1" != "$2" ]]
}

ensure_polkit_rule() {
  if [[ -f "$BT_POLKIT_SENTINEL" ]]; then
    return 0
  fi

  cat >&2 << EOF

NOTE: This script needs to start/stop the bluetooth service, which normally
requires root. A one-time polkit rule can be installed so "$USER" can do this
without sudo. You will be prompted for your password once to write:

  $BT_POLKIT_RULE_FILE

EOF
  read -r -p "Install the polkit rule now? [y/N] " answer
  case "$answer" in
    y | Y | yes | YES) ;;
    *)
      echo "Skipping. The script cannot control bluetooth.service without it." >&2
      return 1
      ;;
  esac

  local tmpfile
  tmpfile=$(mktemp) || return 1

  cat > "$tmpfile" << POLKIT
// Allow user "$USER" to start/stop bluetooth.service without a password.
// Generated by btswitch on $(date -Is).
polkit.addRule(function(action, subject) {
    if (action.id === "org.freedesktop.systemd1.manage-units" &&
        subject.user === "$USER") {
        try {
            var unit = action.lookup("unit");
            if (unit === "$BT_SERVICE.service") {
                return polkit.Result.YES;
            }
        } catch (e) {}
    }
});
POLKIT

  sudo cp "$tmpfile" "$BT_POLKIT_RULE_FILE" || {
    echo "ERROR: Failed to write $BT_POLKIT_RULE_FILE." >&2
    rm -f "$tmpfile"
    return 1
  }
  rm -f "$tmpfile"

  sudo chown root:polkitd "$BT_POLKIT_RULE_FILE" 2>/dev/null \
    || sudo chown root:root "$BT_POLKIT_RULE_FILE"
  sudo chmod 644 "$BT_POLKIT_RULE_FILE"

  mkdir -p "$(dirname "$BT_POLKIT_SENTINEL")"
  touch "$BT_POLKIT_SENTINEL"

  echo "✔ Installed polkit rule: $BT_POLKIT_RULE_FILE" >&2
  echo "  (polkitd picks it up automatically; no restart needed.)" >&2
  return 0
}

systemctl_bt() {
  ensure_polkit_rule || {
    echo "ERROR: Cannot control bluetooth.service without the polkit rule." >&2
    echo "Run 'sudo systemctl $* $BT_SERVICE' manually, or re-run this script and install the rule." >&2
    exit 1
  }
  if ! systemctl "$@" "$BT_SERVICE"; then
    echo "ERROR: systemctl $* $BT_SERVICE failed." >&2
    echo "If polkit prompted for a password, the rule file may be missing." >&2
    echo "Remove $BT_POLKIT_SENTINEL and re-run to reinstall it." >&2
    exit 1
  fi
}

bt_info() {
  bluetoothctl info "$1" 2> /dev/null
}

bt_is_paired() {
  bt_info "$1" | grep -q "Paired: yes"
}

bt_is_trusted() {
  bt_info "$1" | grep -q "Trusted: yes"
}

bt_is_connected() {
  bt_info "$1" | grep -q "Connected: yes"
}

bt_is_mac_address() {
  [[ "$1" =~ ^([0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}$ ]]
}

bt_resolve_device() {
  local device="$1"
  if bt_is_mac_address "$device"; then
    echo "$device"
  else
    echo "${BT_DEVICES[$device]:-}"
  fi
}

bt_get_all_macs() {
  local devices=("${!1}")
  for dev in "${devices[@]}"; do
    local mac
    mac=$(bt_resolve_device "$dev")
    [[ -n "$mac" ]] && echo "$mac"
  done
}

# -------------------------------------------------
# BLUETOOTH CONTROL
# -------------------------------------------------

start_bluetooth() {
  echo "Starting Bluetooth service..."
  systemctl_bt start

  echo "Waiting for bluetoothd to become ready..."
  for i in {1..10}; do
    if bluetoothctl show | grep -q "Powered: yes"; then
      bluetoothctl pairable on > /dev/null 2>&1
      return 0
    fi
    bluetoothctl power on > /dev/null 2>&1
    sleep 1
  done

  echo "ERROR: Bluetooth adapter not found or could not be powered on."
  return 1
}

bt_pair_device() {
  local mac="$1"

  if bt_is_paired "$mac"; then
    return 0
  fi

  echo "✖ Device $mac is not paired. Automatic pairing is not supported."
  echo "  Please pair it manually, then re-run this command:"
  echo ""
  echo "      bluetoothctl"
  echo "      [bluetoothctl]# pair $mac"
  echo "      [bluetoothctl]# trust $mac"
  echo "      [bluetoothctl]# connect $mac"
  echo "      [bluetoothctl]# quit"
  echo ""
  echo "  (Put the device in pairing mode first, if needed.)"
  return 1
}

bt_ensure_connected() {
  local mac="$1"
  echo "----------------------------------------"
  echo "Checking device: $mac"

  if ! bt_is_trusted "$mac"; then
    echo "→ Trusting $mac"
    bluetoothctl trust "$mac" > /dev/null
  fi

  for ((i = 1; i <= BT_RETRIES; i++)); do
    if bt_is_connected "$mac"; then
      echo "✔ Connected: $mac"
      return 0
    fi

    echo "→ Attempt $i/$BT_RETRIES: Connecting to $mac..."

    if ! bt_is_paired "$mac"; then
      bt_pair_device "$mac" || return 1
    fi

    bluetoothctl connect "$mac" > /dev/null 2>&1
    sleep "$BT_DELAY"

    if bt_is_connected "$mac"; then
      echo "✔ Success!"
      return 0
    fi
  done

  echo "✖ Failed to connect $mac after $BT_RETRIES attempts"
  return 1
}

bt_safe_disconnect() {
  local mac="$1"
  if bt_is_connected "$mac"; then
    echo "Disconnecting $mac..."
    bluetoothctl disconnect "$mac" > /dev/null 2>&1
  else
    echo "$mac is already offline."
  fi
}

bt_stop_all() {
  local devices=("${!1}")
  for dev in "${devices[@]}"; do
    local mac
    mac=$(bt_resolve_device "$dev")
    [[ -n "$mac" ]] && bt_safe_disconnect "$mac"
  done
}

bt_start_all() {
  local devices=("${!1}")
  for dev in "${devices[@]}"; do
    local mac
    mac=$(bt_resolve_device "$dev")
    [[ -n "$mac" ]] && bt_ensure_connected "$mac"
  done
}

bt_status_all() {
  local devices=("${!1}")
  for dev in "${devices[@]}"; do
    local mac
    mac=$(bt_resolve_device "$dev")
    if [[ -n "$mac" ]]; then
      echo "--- $dev ($mac) ---"
      if bt_is_connected "$mac"; then
        echo "Status: CONNECTED"
      elif bt_is_paired "$mac"; then
        echo "Status: Paired (Offline)"
      else
        echo "Status: Not Paired / Unknown"
      fi
    fi
  done
}

bt_scan() {
  local scan_duration=30
  declare -A seen_devices
  local device_count=0

  echo "Scanning for Bluetooth devices for $scan_duration seconds..."
  echo "Press Ctrl+C to stop early"
  echo "----------------------------------------"

  trap 'echo ""; echo "Scan stopped"; exit 0' INT TERM

  local start_time=$(date +%s)
  local end_time=$((start_time + scan_duration))

  while true; do
    local current_time=$(date +%s)
    if ((current_time >= end_time)); then
      break
    fi

    while IFS= read -r line; do
      if [[ $line =~ ^[[:space:]]*([0-9A-F:]{17})[[:space:]]+(.+)$ ]]; then
        local mac="${BASH_REMATCH[1]}"
        local name="${BASH_REMATCH[2]}"

        if [[ -z "${seen_devices[$mac]:-}" ]]; then
          seen_devices[$mac]="$name"
          ((device_count++))
          echo "[$device_count] $mac - $name"
        fi
      fi
    done < <(hcitool scan 2> /dev/null | tail -n +2)

    sleep 1
  done

  echo "----------------------------------------"
  echo "Scan complete. Found $device_count unique device(s):"
  for mac in "${!seen_devices[@]}"; do
    echo "  $mac - ${seen_devices[$mac]}"
  done

  trap - INT TERM
}

bt_list() {
  echo "=== BT_DEVICES ==="
  for name in "${!BT_DEVICES[@]}"; do
    echo "  $name: ${BT_DEVICES[$name]}"
  done
  echo ""
  echo "=== BT_AUTO_DEVICES ==="
  for name in "${BT_AUTO_DEVICES[@]}"; do
    echo "  $name: ${BT_DEVICES[$name]}"
  done
}

# -------------------------------------------------
# MAIN
# -------------------------------------------------

DEVICES=()
TIMED_MINUTES=0

parse_args() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      -d | --device)
        DEVICES+=("$2")
        shift 2
        ;;
      -t | --timed)
        TIMED_MINUTES="$2"
        shift 2
        ;;
      start | stop | status | scan | list)
        COMMAND="$1"
        shift
        ;;
      *)
        echo "Unknown option: $1"
        echo "Usage: $0 [-d|--device DEVICE] [-t|--timed MINUTES] {start|stop|status|scan|list}"
        exit 1
        ;;
    esac
  done
}

parse_args "$@"

if [[ -z "${COMMAND:-}" ]]; then
  echo "Usage: $0 [-d|--device DEVICE] [-t|--timed MINUTES] {start|stop|status|scan|list}"
  exit 1
fi

if [[ ${#DEVICES[@]} -gt 0 ]]; then
  for dev in "${DEVICES[@]}"; do
    mac=$(bt_resolve_device "$dev")
    if [[ -z "$mac" ]]; then
      echo "ERROR: Unknown device '$dev'. Valid aliases: ${!BT_DEVICES[@]}"
      exit 1
    fi
  done
fi

case "$COMMAND" in
  start)
    if wifi_switch_needed "$WIFI_5G" "$WIFI_24G"; then
      nmcli connection up "$WIFI_5G"
    fi

    start_bluetooth || exit 1

    if [[ ${#DEVICES[@]} -gt 0 ]]; then
      bt_start_all DEVICES[@]
    else
      bt_start_all BT_AUTO_DEVICES[@]
    fi

    if [[ "$TIMED_MINUTES" -gt 0 ]]; then
      echo "Will disconnect after $TIMED_MINUTES minute(s)..."
      sleep $((TIMED_MINUTES * 60))
      echo "Time elapsed. Disconnecting..."

      if [[ ${#DEVICES[@]} -gt 0 ]]; then
        bt_stop_all DEVICES[@]
      else
        bt_stop_all BT_AUTO_DEVICES[@]
      fi

      echo "Stopping Bluetooth service..."
      systemctl_bt stop

      if wifi_switch_needed "$WIFI_5G" "$WIFI_24G"; then
        nmcli connection up "$WIFI_24G"
      fi
    fi
    ;;

  stop)
    if [[ ${#DEVICES[@]} -gt 0 ]]; then
      bt_stop_all DEVICES[@]
    else
      bt_stop_all BT_AUTO_DEVICES[@]
    fi

    echo "Stopping Bluetooth service..."
    systemctl_bt stop

    if wifi_switch_needed "$WIFI_5G" "$WIFI_24G"; then
      nmcli connection up "$WIFI_24G"
    fi
    ;;

  status)
    bluetoothctl show | grep -E "Name|Powered|Discoverable"

    if [[ ${#DEVICES[@]} -gt 0 ]]; then
      bt_status_all DEVICES[@]
    else
      bt_status_all BT_AUTO_DEVICES[@]
    fi
    ;;

  scan)
    bt_scan
    ;;

  list)
    bt_list
    ;;
esac