#!/usr/bin/env bash
set -euo pipefail

BROKER1_HOST="${MESHCORE_CA_MQTT1_HOST:-mqtt1.meshcore.ca}"
BROKER2_HOST="${MESHCORE_CA_MQTT2_HOST:-mqtt.canadaverse.org}"
BROKER3_HOST="${MESHCORE_CA_MQTT3_HOST:-${MESHCORE_CA_MESHMAPPER_HOST:-mqtt.meshmapper.cc}}"
BROKER_PORT="${MESHCORE_CA_MQTT_PORT:-443}"
IATA="${MESHCORE_CA_IATA:-${CANADAVERSE_IATA:-}}"
MODE="${MESHCORE_CA_MODE:-${CANADAVERSE_CAPTURE_MODE:-auto}}"
DEVICE_TYPE="${MESHCORE_CA_DEVICE:-${CANADAVERSE_DEVICE:-}}"
RESTART_SERVICE=1
INSTALL_MCTOMQTT=0
INSTALL_PACKETCAPTURE=0

MCTOMQTT_CONFIG_DIR="${MCTOMQTT_CONFIG_DIR:-/etc/mctomqtt}"
MCTOMQTT_DROPIN="${MCTOMQTT_DROPIN:-}"
MCTOMQTT_SERVICE="${MCTOMQTT_SERVICE:-mctomqtt}"
MCTOMQTT_INSTALL_URL="${MCTOMQTT_INSTALL_URL:-https://raw.githubusercontent.com/Cisien/meshcoretomqtt/main/install.sh}"

PACKETCAPTURE_DIR="${PACKETCAPTURE_DIR:-$HOME/.meshcore-packet-capture}"
PACKETCAPTURE_ENV_FILE="${PACKETCAPTURE_ENV_FILE:-}"
PACKETCAPTURE_SERVICE="${PACKETCAPTURE_SERVICE:-meshcore-capture}"
PACKETCAPTURE_INSTALL_URL="${PACKETCAPTURE_INSTALL_URL:-https://raw.githubusercontent.com/agessaman/meshcore-packet-capture/main/install.sh}"

usage() {
  cat <<EOF
Add the current MeshCore MQTT broker set to a host-side MeshCore observer install.

Supported installs:
  - USB serial MeshCore node via Cisien/meshcoretomqtt
  - Companion radio via meshcore-packet-capture (.env.local)

Usage:
  MESHCORE_CA_IATA=YOW bash <(curl -fsSL https://canadaverse.org/scripts/add-canadaverse-meshcore-broker.sh)
  bash <(curl -fsSL https://canadaverse.org/scripts/add-canadaverse-meshcore-broker.sh) --iata YOW --device serial-host

Options:
  --iata CODE             Region/IATA code.
  --device TYPE           serial-host | companion
  --mode auto|mctomqtt|env
  --install-mctomqtt      Install Cisien/meshcoretomqtt first if /etc/mctomqtt is missing.
  --install-packetcapture Install meshcore-packet-capture first if the companion install is missing.
  --no-restart            Patch config only.
  --config-dir PATH       Default: /etc/mctomqtt
  --dropin PATH           Default: config.d/20-meshcore-ca.toml, or existing Canadaverse drop-in.
  --service NAME          Default: mctomqtt
  --dir PATH              Default: \$HOME/.meshcore-packet-capture
  --env-file PATH         Explicit .env.local path.
  --packet-service NAME   Default: meshcore-capture
EOF
}

say() {
  printf '[Canadaverse] %s\n' "$*"
}

have_cmd() {
  command -v "$1" >/dev/null 2>&1
}

prompt_yes_no() {
  local prompt="$1"
  local default="${2:-y}"
  local answer=""
  [ -t 0 ] || return 1
  if [ "$default" = "y" ]; then
    printf '%s [Y/n] ' "$prompt"
  else
    printf '%s [y/N] ' "$prompt"
  fi
  read -r answer
  answer="$(printf '%s' "${answer:-$default}" | tr '[:upper:]' '[:lower:]')"
  case "$answer" in
    y|yes) return 0 ;;
    *) return 1 ;;
  esac
}

backup_path() {
  local original="$1"
  local stamp candidate n
  stamp="$(date +%Y%m%d%H%M%S)"
  candidate="${original}.bak.${stamp}"
  n=1
  while [ -e "$candidate" ]; do
    candidate="${original}.bak.${stamp}.${n}"
    n=$((n + 1))
  done
  printf '%s\n' "$candidate"
}

normalize_device_type() {
  case "$(printf '%s' "$1" | tr '[:upper:]' '[:lower:]')" in
    1|serial|serial-host|serial_host|usb|usb-host|server|repeater|roomserver|room-server|mctomqtt)
      printf 'serial-host'
      ;;
    2|companion|packetcapture|packet-capture|meshcore-packet-capture|env)
      printf 'companion'
      ;;
    *)
      return 1
      ;;
  esac
}

upper_iata() {
  printf '%s' "$1" | tr '[:lower:]' '[:upper:]'
}

require_iata() {
  if [ -z "$IATA" ] && [ -t 0 ]; then
    printf 'Enter 3-character region/IATA code: '
    read -r IATA
  fi
  if [ -z "$IATA" ]; then
    cat >&2 <<EOF
Missing region/IATA code.

Run with:
  MESHCORE_CA_IATA=YOW bash <(curl -fsSL https://canadaverse.org/scripts/add-canadaverse-meshcore-broker.sh)
EOF
    exit 1
  fi
  IATA="$(upper_iata "$IATA")"
  if ! printf '%s' "$IATA" | grep -Eq '^[A-Z0-9]{3}$'; then
    echo "IATA/region code must be exactly 3 letters or numbers, got: $IATA" >&2
    exit 1
  fi
  say "Region selected: $IATA"
}

as_root() {
  if [ "$(id -u)" -eq 0 ]; then
    "$@"
  elif have_cmd sudo; then
    sudo "$@"
  else
    echo "This action needs root permissions and sudo is not available: $*" >&2
    exit 1
  fi
}

restart_systemd_if_present() {
  local service="$1"
  if [ "$RESTART_SERVICE" -ne 1 ] || ! have_cmd systemctl; then
    return 0
  fi
  if systemctl list-unit-files --type=service 2>/dev/null | awk '{print $1}' | grep -qx "$service.service"; then
    say "Restarting systemd service: $service"
    as_root systemctl restart "$service"
  else
    say "No $service systemd service found; restart the process manually."
  fi
}

packetcapture_env_path() {
  if [ -n "$PACKETCAPTURE_ENV_FILE" ]; then
    printf '%s\n' "$PACKETCAPTURE_ENV_FILE"
  else
    printf '%s/.env.local\n' "$PACKETCAPTURE_DIR"
  fi
}

detect_mctomqtt_install() {
  [ -d "$MCTOMQTT_CONFIG_DIR" ] && return 0
  have_cmd systemctl && systemctl list-unit-files --type=service 2>/dev/null | awk '{print $1}' | grep -qx "$MCTOMQTT_SERVICE.service"
}

detect_packetcapture_install() {
  local env_file base_dir
  env_file="$(packetcapture_env_path)"
  base_dir="$(dirname "$env_file")"
  [ -f "$env_file" ] || [ -f "$base_dir/packet_capture.py" ]
}

install_mctomqtt() {
  detect_mctomqtt_install && return 0
  if [ "$INSTALL_MCTOMQTT" -ne 1 ]; then
    if prompt_yes_no "meshcoretomqtt is not installed. Install it now?" "y"; then
      INSTALL_MCTOMQTT=1
    else
      return 0
    fi
  fi
  [ "$INSTALL_MCTOMQTT" -eq 1 ] || return 0
  have_cmd curl || { echo "curl is required for --install-mctomqtt." >&2; exit 1; }
  local tmp
  tmp="$(mktemp)"
  say "Installing Cisien/meshcoretomqtt from: $MCTOMQTT_INSTALL_URL"
  curl -fsSL "$MCTOMQTT_INSTALL_URL" -o "$tmp"
  as_root bash "$tmp"
  rm -f "$tmp"
}

install_packetcapture() {
  detect_packetcapture_install && return 0
  if [ "$INSTALL_PACKETCAPTURE" -ne 1 ]; then
    if prompt_yes_no "meshcore-packet-capture is not installed. Install it now?" "y"; then
      INSTALL_PACKETCAPTURE=1
    else
      return 0
    fi
  fi
  [ "$INSTALL_PACKETCAPTURE" -eq 1 ] || return 0
  have_cmd curl || { echo "curl is required for --install-packetcapture." >&2; exit 1; }
  local tmp
  tmp="$(mktemp)"
  say "Installing meshcore-packet-capture from: $PACKETCAPTURE_INSTALL_URL"
  curl -fsSL "$PACKETCAPTURE_INSTALL_URL" -o "$tmp"
  bash "$tmp"
  rm -f "$tmp"
}

select_device_type() {
  if [ -n "$DEVICE_TYPE" ]; then
    DEVICE_TYPE="$(normalize_device_type "$DEVICE_TYPE" || true)"
    [ -n "$DEVICE_TYPE" ] || { echo "--device must be serial-host or companion" >&2; exit 2; }
    return 0
  fi
  case "$MODE" in
    mctomqtt) DEVICE_TYPE="serial-host"; return 0 ;;
    env) DEVICE_TYPE="companion"; return 0 ;;
  esac
  if detect_mctomqtt_install; then
    DEVICE_TYPE="serial-host"
  elif detect_packetcapture_install; then
    DEVICE_TYPE="companion"
  else
    DEVICE_TYPE="serial-host"
  fi
}

choose_mctomqtt_dropin() {
  if [ -n "$MCTOMQTT_DROPIN" ]; then
    return 0
  fi
  if [ -f "$MCTOMQTT_CONFIG_DIR/config.d/20-canadaverse.toml" ]; then
    MCTOMQTT_DROPIN="$MCTOMQTT_CONFIG_DIR/config.d/20-canadaverse.toml"
  elif [ -f "$MCTOMQTT_CONFIG_DIR/config.d/zz-canadaverse.toml" ]; then
    MCTOMQTT_DROPIN="$MCTOMQTT_CONFIG_DIR/config.d/zz-canadaverse.toml"
  else
    MCTOMQTT_DROPIN="$MCTOMQTT_CONFIG_DIR/config.d/20-meshcore-ca.toml"
  fi
}

patch_mctomqtt() {
  require_iata
  install_mctomqtt
  choose_mctomqtt_dropin
  if [ ! -d "$MCTOMQTT_CONFIG_DIR" ]; then
    echo "meshcoretomqtt config directory was not found: $MCTOMQTT_CONFIG_DIR" >&2
    exit 1
  fi

  local dropin_dir tmp backup
  dropin_dir="$(dirname "$MCTOMQTT_DROPIN")"
  tmp="$(mktemp)"
  cat > "$tmp" <<EOF
# MeshCore broker drop-in for Cisien/meshcoretomqtt.
# Generated by add-canadaverse-meshcore-broker.sh.

[general]
iata = "$IATA"

[[broker]]
name = "meshcore-ca-1"
enabled = true
server = "$BROKER1_HOST"
port = $BROKER_PORT
transport = "websockets"
websocket_path = "/mqtt"
keepalive = 60
qos = 0
retain = true

[broker.tls]
enabled = true
verify = true

[broker.auth]
method = "token"
audience = "$BROKER1_HOST"

[[broker]]
name = "canadaverse"
enabled = true
server = "$BROKER2_HOST"
port = $BROKER_PORT
transport = "websockets"
websocket_path = "/mqtt"
keepalive = 60
qos = 0
retain = true

[broker.tls]
enabled = true
verify = true

[broker.auth]
method = "token"
audience = "$BROKER2_HOST"

[[broker]]
name = "meshmapper"
enabled = true
server = "$BROKER3_HOST"
port = $BROKER_PORT
transport = "websockets"
websocket_path = "/mqtt"
keepalive = 60
qos = 0
retain = true

[broker.tls]
enabled = true
verify = true

[broker.auth]
method = "token"
audience = "$BROKER3_HOST"
EOF

  as_root mkdir -p "$dropin_dir"
  if [ -f "$MCTOMQTT_DROPIN" ]; then
    backup="$(backup_path "$MCTOMQTT_DROPIN")"
    as_root cp -p "$MCTOMQTT_DROPIN" "$backup"
    say "Backup written: $backup"
  fi
  as_root install -m 0644 "$tmp" "$MCTOMQTT_DROPIN"
  rm -f "$tmp"
  say "Patched: $MCTOMQTT_DROPIN"
  say "Brokers: $BROKER1_HOST:$BROKER_PORT, $BROKER2_HOST:$BROKER_PORT, $BROKER3_HOST:$BROKER_PORT"
  restart_systemd_if_present "$MCTOMQTT_SERVICE"
}

env_value() {
  local file="$1"
  local key="$2"
  awk -F= -v key="$key" '$1 == key { v=$0; sub("^[^=]*=", "", v); print v }' "$file" | tail -n 1
}

upsert_env() {
  local file="$1"
  local key="$2"
  local value="$3"
  local tmp
  tmp="$(mktemp)"
  awk -v key="$key" -v value="$value" '
    BEGIN { done=0 }
    $0 ~ "^" key "=" {
      if (!done) {
        print key "=" value
        done=1
      }
      next
    }
    { print }
    END {
      if (!done) print key "=" value
    }
  ' "$file" > "$tmp"
  cat "$tmp" > "$file"
  rm -f "$tmp"
}

set_packetcapture_slot() {
  local file="$1"
  local slot="$2"
  local host="$3"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_ENABLED" "true"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_SERVER" "$host"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_PORT" "$BROKER_PORT"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_TRANSPORT" "websockets"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_WEBSOCKET_PATH" "/mqtt"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_USE_TLS" "true"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_TLS_VERIFY" "true"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_USE_AUTH_TOKEN" "true"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_TOKEN_AUDIENCE" "$host"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_KEEPALIVE" "120"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_QOS" "0"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_RETAIN" "true"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_TOPIC_STATUS" 'meshcore/{IATA}/{PUBLIC_KEY}/status'
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_TOPIC_PACKETS" 'meshcore/{IATA}/{PUBLIC_KEY}/packets'
}

disable_packetcapture_slot() {
  local file="$1"
  local slot="$2"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_ENABLED" "false"
  upsert_env "$file" "PACKETCAPTURE_MQTT${slot}_SERVER" ""
}

patch_env_capture() {
  local env_file current_iata backup n
  env_file="$(packetcapture_env_path)"
  install_packetcapture
  mkdir -p "$(dirname "$env_file")"
  [ -f "$env_file" ] || install -m 0600 /dev/null "$env_file"
  current_iata="$(env_value "$env_file" "PACKETCAPTURE_IATA" | tr -d '\r' || true)"
  if [ -z "$IATA" ]; then
    IATA="$current_iata"
  fi
  require_iata

  backup="$(backup_path "$env_file")"
  cp -p "$env_file" "$backup"
  chmod 0600 "$backup"

  upsert_env "$env_file" "PACKETCAPTURE_IATA" "$IATA"
  set_packetcapture_slot "$env_file" 1 "$BROKER1_HOST"
  set_packetcapture_slot "$env_file" 2 "$BROKER2_HOST"
  set_packetcapture_slot "$env_file" 3 "$BROKER3_HOST"
  for n in 4 5 6; do
    disable_packetcapture_slot "$env_file" "$n"
  done
  chmod 0600 "$env_file"
  say "Patched: $env_file"
  say "Backup written: $backup"
  say "Brokers: $BROKER1_HOST:$BROKER_PORT, $BROKER2_HOST:$BROKER_PORT, $BROKER3_HOST:$BROKER_PORT"
  restart_systemd_if_present "$PACKETCAPTURE_SERVICE"
}

while [ "$#" -gt 0 ]; do
  case "$1" in
    --iata) IATA="${2:?Missing value for --iata}"; shift 2 ;;
    --device) DEVICE_TYPE="${2:?Missing value for --device}"; shift 2 ;;
    --mode) MODE="${2:?Missing value for --mode}"; shift 2 ;;
    --install-mctomqtt) INSTALL_MCTOMQTT=1; shift ;;
    --install-packetcapture) INSTALL_PACKETCAPTURE=1; shift ;;
    --no-restart) RESTART_SERVICE=0; shift ;;
    --restart) RESTART_SERVICE=1; shift ;;
    --config-dir) MCTOMQTT_CONFIG_DIR="${2:?Missing value for --config-dir}"; shift 2 ;;
    --dropin) MCTOMQTT_DROPIN="${2:?Missing value for --dropin}"; shift 2 ;;
    --service) MCTOMQTT_SERVICE="${2:?Missing value for --service}"; shift 2 ;;
    --dir) PACKETCAPTURE_DIR="${2:?Missing value for --dir}"; shift 2 ;;
    --env-file) PACKETCAPTURE_ENV_FILE="${2:?Missing value for --env-file}"; shift 2 ;;
    --packet-service) PACKETCAPTURE_SERVICE="${2:?Missing value for --packet-service}"; shift 2 ;;
    --slot) shift 2 ;; # Kept for compatibility; slots 1, 2, and 3 are now managed together.
    --help|-h) usage; exit 0 ;;
    *) echo "Unknown option: $1" >&2; usage >&2; exit 2 ;;
  esac
done

case "$MODE" in
  auto|mctomqtt|env) ;;
  *) echo "--mode must be auto, mctomqtt, or env" >&2; exit 2 ;;
esac

select_device_type
say "Device path: $DEVICE_TYPE"
if [ "$DEVICE_TYPE" = "serial-host" ]; then
  patch_mctomqtt
else
  patch_env_capture
fi
say "Done."