#!/bin/sh

# ----- root + platform checks -----
if [ "$(id -u)" -ne 0 ]; then
    printf '\033[31mPlease run this script as root\033[0m\n' >&2
    exit 1
fi

if ! [ -r /proc/device-tree/model ] || ! grep -q 'Raspberry Pi' /proc/device-tree/model; then
    printf '\033[31mThis script is for Raspberry Pi devices only\033[0m\n' >&2
    exit 1
fi

# ----- Settings -----
NM_IFACE="${USB_GADGET_IFACE:-usb0}"

CLIENT_NAME='USB Gadget (client)'
SHARED_NAME='USB Gadget (shared)'
SERVICE="rpi-usb-gadget-ics.service"

SHARED_ADDR="10.12.194.1/28"
CLIENT_DHCP_TIMEOUT=6
CLIENT_ROUTE_METRIC=100
NM_WAIT_DOWN=5
NM_WAIT_UP=10

# modules-load + overlay
MODS_CONF="/etc/modules-load.d/usb-gadget.conf"
CFG_FW=/boot/firmware/config.txt
CFG_LEGACY=/boot/config.txt
OVERLAY_LINE='dtoverlay=dwc2,dr_mode=peripheral'
DNSMASQ_DIR=/etc/NetworkManager/dnsmasq-shared.d
DNSMASQ_SNIPPET="$DNSMASQ_DIR/90-rpi-usb-gadget-lease.conf"

# A list of supported devices (for informational purposes)
SUPPORTED_DEVICES="$(printf '%s\n' 'Raspberry Pi Zero' 'Raspberry Pi Zero W' 'Raspberry Pi Zero 2 W' 'Raspberry Pi 3 Model A+' 'Raspberry Pi 4 Model B' 'Raspberry Pi 5 Model B' 'Raspberry Pi 500' 'Raspberry Pi Compute Module 0' 'Raspberry Pi Compute Module 5')"

#################################
##### Systemd service setup #####
#################################

have_systemctl() { command -v systemctl >/dev/null 2>&1; }

svc_enable_now() {
    have_systemctl || return 0
    systemctl daemon-reload || true
    systemctl enable --now "$SERVICE" 2>/dev/null || true
}

svc_disable_now() {
    have_systemctl || return 0
    systemctl disable --now "$SERVICE" 2>/dev/null || true
}

svc_status() {
    have_systemctl || { printf '(no systemd)\n'; return; }
    if systemctl --quiet is-enabled "$SERVICE"; then en=yes; else en=no; fi
    if systemctl --quiet is-active  "$SERVICE"; then ac=yes; else ac=no; fi
    printf "ICS watcher: enabled=$en active=$ac\n"
}

################################
# NetworkManager (nmcli) setup #
################################

have_nmcli() { command -v nmcli >/dev/null 2>&1; }

nm_down() { nmcli -w "$NM_WAIT_DOWN" con down "$1" >/dev/null 2>&1 || true; }
nm_up()   { nmcli -w "$NM_WAIT_UP"   con up   "$1" >/dev/null 2>&1 || true; }

# Return the first UUID whose connection.id == $1
nm_first_uuid_by_name() {
    nmcli -t -f NAME,UUID connection show 2>/dev/null \
      | awk -F: -v n="$1" '$1==n { print $2; exit }'
}

# Delete ALL connections with connection.id == $1
nm_delete_all_by_name() {
    nmcli -t -f NAME,UUID connection show 2>/dev/null \
      | awk -F: -v n="$1" '$1==n { print $2 }' \
      | while IFS= read -r u; do
            [ -n "$u" ] && nmcli connection delete uuid "$u" >/dev/null 2>&1 || true
        done
}

nm_ensure_client() {
    # Ensure exactly one "client" profile with desired settings
    uuid=$(nm_first_uuid_by_name "$CLIENT_NAME")

    # DHCP client (for ICS); IPv6 disabled; short timeout; prefer routes via DHCP
    if [ -z "$uuid" ]; then
        nm_delete_all_by_name "$CLIENT_NAME"  # paranoia cleanup
        nmcli connection add type ethernet ifname "$NM_IFACE" con-name "$CLIENT_NAME" \
            connection.autoconnect no \
            connection.autoconnect-priority 100 \
            ipv4.method auto \
            ipv4.may-fail yes \
            ipv4.route-metric "$CLIENT_ROUTE_METRIC" \
            ipv4.dhcp-timeout "$CLIENT_DHCP_TIMEOUT" \
            ipv6.method disabled >/dev/null 2>&1 || true
    else
        # modify by UUID so we don't hit the "multiple connections with same name" warning
        nmcli connection modify uuid "$uuid" \
            connection.type ethernet \
            connection.interface-name "$NM_IFACE" \
            connection.autoconnect no \
            connection.autoconnect-priority 100 \
            ipv4.method auto \
            ipv4.may-fail yes \
            ipv4.route-metric "$CLIENT_ROUTE_METRIC" \
            ipv4.dhcp-timeout "$CLIENT_DHCP_TIMEOUT" \
            ipv6.method disabled >/dev/null 2>&1 || true

        # remove any other duplicates that might exist
        nmcli -t -f NAME,UUID connection show 2>/dev/null \
          | awk -F: -v n="$CLIENT_NAME" '$1==n { print $2 }' \
          | tail -n +2 | while IFS= read -r du; do
                [ -n "$du" ] && nmcli connection delete uuid "$du" >/dev/null 2>&1 || true
            done
    fi
}

nm_ensure_shared() {
    uuid=$(nm_first_uuid_by_name "$SHARED_NAME")

    # Pi serves DHCP+NAT on 10.12.194.0/28; IPv6 disabled; do autoconnect
    if [ -z "$uuid" ]; then
        nm_delete_all_by_name "$SHARED_NAME"
        nmcli connection add type ethernet ifname "$NM_IFACE" con-name "$SHARED_NAME" \
            connection.autoconnect yes \
            connection.autoconnect-priority 10 \
            ipv4.method shared \
            ipv4.addresses "$SHARED_ADDR" \
            ipv6.method disabled >/dev/null 2>&1 || true
    else
        nmcli connection modify uuid "$uuid" \
            connection.type ethernet \
            connection.interface-name "$NM_IFACE" \
            connection.autoconnect yes \
            connection.autoconnect-priority 10 \
            ipv4.method shared \
            ipv4.addresses "$SHARED_ADDR" \
            ipv6.method disabled >/dev/null 2>&1 || true

        nmcli -t -f NAME,UUID connection show 2>/dev/null \
          | awk -F: -v n="$SHARED_NAME" '$1==n { print $2 }' \
          | tail -n +2 | while IFS= read -r du; do
                [ -n "$du" ] && nmcli connection delete uuid "$du" >/dev/null 2>&1 || true
            done
    fi
}

nm_remove_profiles() {
    nm_delete_all_by_name "$CLIENT_NAME"
    nm_delete_all_by_name "$SHARED_NAME"
    nmcli connection reload >/dev/null 2>&1 || true
}

nm_dnsmasq_setup() {
    mkdir -p "$DNSMASQ_DIR" && chmod 0755 "$DNSMASQ_DIR"
    : "${DNSMASQ_SNIPPET:=$DNSMASQ_DIR/90-rpi-usb-gadget-lease.conf}"

    cat >"$DNSMASQ_SNIPPET" <<'EOF'
# Autogenerated by rpi-usb-gadget – safe to delete

# controls server behaviour (send NAKs immediately, faster failover)
dhcp-authoritative
# 2 minute leases (clients should renew every 1 minute)
dhcp-option=51,120
# suppresses dnsmasq’s DHCP chatter in logs
quiet-dhcp
EOF

    # If shared is currently active, bounce it so the new dnsmasq picks this up
    if nmcli -t -f NAME con show --active | grep -F -q -e "^$SHARED_NAME$"; then
        nm_down "$SHARED_NAME"
        nm_up   "$SHARED_NAME"
    fi
}

nm_dnsmasq_remove() {
    rm -f "$DNSMASQ_SNIPPET" || true
    
    # Bounce only if the shared profile is active
    if nmcli -t -f NAME con show --active | grep -F -q -e "^$SHARED_NAME$"; then
        nm_down "$SHARED_NAME"
        nm_up   "$SHARED_NAME"
    fi
}

nm_status() {
  printf ':: NetworkManager:\n'
  if ! have_nmcli; then
    printf '  (nmcli not found)\n'; return
  fi

  # Pull details for the gadget iface
  st=""
  fullst=$(nmcli -g GENERAL.STATE device show "$NM_IFACE" 2>/dev/null)
  st=$(printf '%s' "$fullst" | sed -n 's/.*(\(.*\)).*/\1/p')   # extract text in ()
  [ -n "$st" ] || st="$fullst"
  active=$(nmcli -g GENERAL.CONNECTION device show "$NM_IFACE" 2>/dev/null)
  ip4=$(nmcli -g IP4.ADDRESS device show "$NM_IFACE" 2>/dev/null | head -n1)
  gw=$(nmcli -g IP4.GATEWAY device show "$NM_IFACE" 2>/dev/null)

  printf "  iface:        %s\n" "$NM_IFACE"
  printf "  link:         %s\n" "${st:--}"
  printf "  active prof:  %s\n" "${active:--}"
  printf "  IPv4:         %s\n" "${ip4:--}"
  printf "  gateway:      %s\n" "${gw:--}"

  # List our two profiles with autoconnect + priority
  LC_ALL=C nmcli -t -f NAME,AUTOCONNECT,AUTOCONNECT-PRIORITY connection show \
    | awk -F: -v c1="$CLIENT_NAME" -v c2="$SHARED_NAME" '
    BEGIN { IGNORECASE=1 }
    $1==c1 || $1==c2 {
        v = tolower($2)
        prio = ($3 == "" ? 0 : $3)
        mode = (v=="yes" || v=="on" || v=="true" || v=="1") ? "auto" : "manual"
        printf "  profile:      %s [%s, prio %s]\n", $1, mode, prio
    }'
}

#################################
# Main script logic starts here #
#################################

TURN_ON=true
FORCE=false
case "$1" in
    on)      TURN_ON=true ;;
    off)     TURN_ON=false ;;
    toggle)
        if [ -f "$MODS_CONF" ]; then
            TURN_ON=false
        fi
        ;;
    status)
        if [ -f "$MODS_CONF" ]; then
            printf '\033[33mUSB Gadget mode is on\033[0m\n'
        else
            printf '\033[33mUSB Gadget mode is off\033[0m\n'
        fi
        nm_status
        svc_status
        exit 0
        ;;
    help|"")
        printf "Usage: rpi-usb-gadget [on|off|toggle|status|help]\n"
        exit 1
        ;;
    *)
        printf "Usage: rpi-usb-gadget [on|off|toggle|status|help]\n"
        exit 1
        ;;
esac

if [ "$2" = "-f" ] || [ "$2" = "--force" ]; then
    FORCE=true
fi

check_supported() {
    # Read and sanitize the model string
    model=$(tr -d '\0' < /proc/device-tree/model | tr -d '\r')
    
    # Strip "Rev ..." and trailing spaces
    base_model=$(printf '%s' "$model" \
        | sed -E 's/[[:space:]]+Rev[[:space:]].*$//' \
        | sed -E 's/^[[:space:]]+|[[:space:]]+$//g')

    # Check if the cleaned base model is exactly in the list or matches a prefix
    if printf '%s\n' "$SUPPORTED_DEVICES" | grep -Fxq "$base_model"; then
        return 0
    fi

    printf '\033[33mWarning: Your device model is not in the list of supported devices:\033[0m\n'
    printf '%s\n' "$SUPPORTED_DEVICES" | sed 's/^/  - /'
    printf '\033[33mYour device model (raw): %s\033[0m\n' "$model"
    printf '\033[33mDetected model (cleaned): %s\033[0m\n' "$base_model"

    if [ "$FORCE" = false ]; then
        while true; do
            printf 'Do you want to continue? (y/n): '
            read -r ans
            case "$ans" in
                y|Y) return 0 ;;
                n|N) return 1 ;;
                *)   printf 'Please answer y or n.\n' ;;
            esac
        done
    else
        printf '\033[33mForce flag detected; proceeding despite unsupported device.\033[0m\n'
        return 0
    fi
}

if [ "$TURN_ON" = false ]; then
    printf 'Turning \033[31moff\033[0m USB Gadget mode\n'

    # stop ICS watcher first
    svc_disable_now

    # stop NM bits first
    if have_nmcli; then
        nm_down "$SHARED_NAME"
        nm_down "$CLIENT_NAME"
        nm_remove_profiles
        nm_dnsmasq_remove
    fi

    # kernel gadget off
    rm -f "$MODS_CONF"
    sed -i "\|^$OVERLAY_LINE\$|d" "$CFG_FW" 2>/dev/null || true
    sed -i "\|^$OVERLAY_LINE\$|d" "$CFG_LEGACY" 2>/dev/null || true

else
    check_supported || exit 1

    printf 'Turning \033[32mon\033[0m USB Gadget mode\n'
    printf "g_ether\n" > "$MODS_CONF"
    # ensure overlay present once
    sed -i "\|^$OVERLAY_LINE\$|d" "$CFG_FW" 2>/dev/null || true
    sed -i "\|^$OVERLAY_LINE\$|d" "$CFG_LEGACY" 2>/dev/null || true
    printf '%s\n' "$OVERLAY_LINE" >> "$CFG_FW" 2>/dev/null || printf '%s\n' "$OVERLAY_LINE" >> "$CFG_LEGACY"

    # NetworkManager: dnsmasq + client profile + shared profile
    if have_nmcli; then
        nm_dnsmasq_setup
        nm_ensure_client
        nm_ensure_shared
        nmcli connection reload >/dev/null 2>&1 || true

        # Make sure NetworkManager actually manages the gadget iface
        nmcli device set "$NM_IFACE" managed yes >/dev/null 2>&1 || true

        # Try client first; ICS-watch will flip to shared if no DHCP/gateway
        nm_down "$CLIENT_NAME"
        nm_up   "$SHARED_NAME"

        # start ICS watcher so it can auto-switch client/shared
        svc_enable_now
    else
        printf '\033[33mNetworkManager (nmcli) not found; skipping NM auto-config.\033[0m\n'
    fi
fi

printf "Reboot to apply changes\n"
