#!/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"

#################################
##### 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; }

nm_ensure_client() {
    # DHCP client (for ICS); IPv6 disabled; short timeout; prefer routes via DHCP
    if nmcli -t -f NAME connection show | grep -F -q -e "^$CLIENT_NAME$"; then
        nmcli connection modify "$CLIENT_NAME" \
            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 || true
    else
        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 || true
    fi
}

nm_ensure_shared() {
    # Pi serves DHCP+NAT on 10.12.194.0/28; IPv6 disabled; do autoconnect
    if nmcli -t -f NAME connection show | grep -F -q -e "^$SHARED_NAME$"; then
        nmcli connection modify "$SHARED_NAME" \
            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 || true
    else
        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 || true
    fi
}

nm_remove_profiles() {
    nmcli -t -f NAME connection show | grep -Fxq "$CLIENT_NAME" && \
        nmcli connection delete "$CLIENT_NAME" >/dev/null 2>&1 || true
    nmcli -t -f NAME connection show | grep -Fxq "$SHARED_NAME" && \
        nmcli connection delete "$SHARED_NAME" >/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
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 [ "$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
        nmcli connection reload >/dev/null 2>&1 || true
        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
    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

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