#!/usr/bin/env bash
# wml-assign-ports — declarative port-role assignment wizard.
#
# Industry model (OPNsense/pfSense/Fortinet): the admin EXPLICITLY assigns
# each physical NIC to a role (WAN / LAN / DMZ / optN). We bind the role to the
# NIC's PERMANENT MAC via a systemd .link file that renames it wan0/lan0/dmzN.
# After this, the role is stable across reboots and PCI renumbering — there is
# NO DHCP-probe auto-detection and therefore none of the swap/flicker bugs.
#
# Two ID methods are offered:
#   1. Link-blink ("plug a cable into the port you want as WAN now") — we watch
#      /sys/class/net/<nic>/carrier and pick the NIC whose carrier just rose.
#   2. Manual list (pick NIC by name + MAC + speed).
#
# Writes:
#   /etc/systemd/network/10-<role>.link   (Match PermanentMACAddress -> Name)
#   /etc/wml/port-roles.conf              (role=mac mapping, human-readable)
#
# Exits 0 on success. Designed to run from the installer (target chroot) AND
# from the console menu / first boot on the live system.

set -uo pipefail

# Colors (fall back to plain if not a tty)
if [[ -t 1 ]]; then
    B=$'\e[1m'; D=$'\e[2m'; G=$'\e[32m'; Y=$'\e[33m'; R=$'\e[31m'; P=$'\e[35m'; X=$'\e[0m'
else
    B=""; D=""; G=""; Y=""; R=""; P=""; X=""
fi

# Allow the installer to target a mounted root instead of /
ROOT="${WML_TARGET_ROOT:-}"
NETDIR="${ROOT}/etc/systemd/network"
ROLES_CONF="${ROOT}/etc/wml/port-roles.conf"
mkdir -p "$NETDIR" "${ROOT}/etc/wml"

# --- enumerate physical NICs (skip virtual) ---
# Includes already-renamed role NICs (wan0/lan0/dmzN) so re-assignment works.
phys_nics() {
    for d in /sys/class/net/*; do
        n=$(basename "$d")
        case "$n" in
            lo|docker*|virbr*|veth*|tap*|wg*|br-*|tun*|tailscale*) continue ;;
        esac
        [[ -e "$d/device" ]] || continue
        echo "$n"
    done | sort
}

# permanent (burned-in) MAC — falls back to current MAC for virtual NICs
perm_mac() {
    local nic="$1" mac
    mac=$(ethtool -P "$nic" 2>/dev/null | awk '{print $NF}')
    if [[ -z "$mac" || "$mac" == "00:00:00:00:00:00" ]]; then
        mac=$(cat "/sys/class/net/$nic/address" 2>/dev/null)
    fi
    echo "$mac"
}

nic_speed() {
    local s
    s=$(cat "/sys/class/net/$1/speed" 2>/dev/null)
    [[ -n "$s" && "$s" -gt 0 ]] 2>/dev/null && echo "${s} Mbps" || echo "link down"
}

nic_carrier() { cat "/sys/class/net/$1/carrier" 2>/dev/null || echo 0; }

# --- link-blink detection: return the NIC whose carrier rises ---
# Snapshot carriers, ask user to plug in, poll until one transitions 0->1.
detect_by_blink() {
    local timeout="${1:-60}"
    declare -A before
    local nics
    mapfile -t nics < <(phys_nics)
    for n in "${nics[@]}"; do
        ip link set "$n" up 2>/dev/null
        before[$n]=$(nic_carrier "$n")
    done
    local deadline=$(( $(date +%s) + timeout ))
    while [[ $(date +%s) -lt $deadline ]]; do
        for n in "${nics[@]}"; do
            local now
            now=$(nic_carrier "$n")
            if [[ "${before[$n]}" == "0" && "$now" == "1" ]]; then
                # settle to avoid bounce
                sleep 1
                if [[ "$(nic_carrier "$n")" == "1" ]]; then
                    echo "$n"
                    return 0
                fi
            fi
            before[$n]=$now
        done
        sleep 1
    done
    return 1
}

# --- write the MAC-bound .link file + roles.conf entry ---
assign_role() {
    local role="$1" nic="$2" mac
    mac=$(perm_mac "$nic")
    if [[ -z "$mac" ]]; then
        echo "${R}Could not read MAC for $nic — cannot bind role.${X}"
        return 1
    fi
    # Priority prefix: wan0=10, lan0=20, dmz0=30, optN=40+
    local prio
    case "$role" in
        wan0) prio=10 ;;
        lan0) prio=20 ;;
        dmz0) prio=30 ;;
        *)    prio=40 ;;
    esac
    cat > "${NETDIR}/${prio}-${role}.link" <<EOF
# WML OS port role: ${role} -> ${nic} (MAC ${mac})
[Match]
PermanentMACAddress=${mac}

[Link]
Name=${role}
EOF
    chmod 644 "${NETDIR}/${prio}-${role}.link"
    # Append to human-readable roles.conf (idempotent: drop old line for role)
    if [[ -f "$ROLES_CONF" ]]; then
        grep -v "^${role}=" "$ROLES_CONF" > "${ROLES_CONF}.tmp" 2>/dev/null || true
        mv "${ROLES_CONF}.tmp" "$ROLES_CONF"
    fi
    echo "${role}=${mac}" >> "$ROLES_CONF"
    echo "${G}  ✓ ${role} bound to ${nic} (${mac})${X}"
}

prompt() { local p="$1" d="${2:-}"; local a; read -rp "$p" a; echo "${a:-$d}"; }

# --- main wizard ---
main() {
    clear 2>/dev/null || true
    echo
    echo "  ${B}${P}WML OS — Interface Assignment${X}"
    echo "  ${D}Bind each physical port to a role. Roles are stable across reboots.${X}"
    echo

    local nics
    mapfile -t nics < <(phys_nics)
    if [[ ${#nics[@]} -eq 0 ]]; then
        echo "  ${R}No physical network interfaces detected.${X}"
        return 1
    fi

    echo "  ${B}Detected interfaces:${X}"
    for n in "${nics[@]}"; do
        printf "    %-10s MAC %s   %s   carrier:%s\n" \
            "$n" "$(perm_mac "$n")" "$(nic_speed "$n")" \
            "$([[ "$(nic_carrier "$n")" == "1" ]] && echo "${G}up${X}" || echo "${D}down${X}")"
    done
    echo

    # How many roles? WAN is required; LAN required; extra LAN/DMZ optional.
    local method
    echo "  ${B}How do you want to identify ports?${X}"
    echo "    ${B}1)${X} Plug-and-detect (recommended) — plug a cable into the port for each role"
    echo "    ${B}2)${X} Manual — pick from the list by name/MAC"
    method=$(prompt "  Selection [1]: " 1)

    # Fresh roles.conf + clear old role .link files
    : > "$ROLES_CONF"
    rm -f "${NETDIR}"/10-wan0.link "${NETDIR}"/20-lan0.link "${NETDIR}"/3*-dmz*.link "${NETDIR}"/4*-opt*.link 2>/dev/null

    local -a ROLE_LIST=(wan0 lan0)
    # Ask if they want a DMZ / extra LANs
    local more
    more=$(prompt "  Add a DMZ port? [y/N]: " n)
    [[ "$more" =~ ^[Yy] ]] && ROLE_LIST+=(dmz0)

    declare -A used_nic
    for role in "${ROLE_LIST[@]}"; do
        echo
        local label="${role^^}"
        echo "  ${B}${P}Assign ${label}${X}"
        local chosen=""
        if [[ "$method" == "1" ]]; then
            echo "  ${Y}Unplug the ${label} cable, then plug it into the desired port now.${X}"
            echo "  ${D}Waiting up to 60s for a link to appear...${X}"
            chosen=$(detect_by_blink 60) || {
                echo "  ${R}No link detected in 60s. Falling back to manual.${X}"
                method_fallback=1
            }
            if [[ -n "$chosen" ]]; then
                if [[ -n "${used_nic[$chosen]:-}" ]]; then
                    echo "  ${R}$chosen is already assigned to ${used_nic[$chosen]}. Pick a different port.${X}"
                    chosen=""
                fi
            fi
        fi
        # Manual fallback / method 2
        while [[ -z "$chosen" ]]; do
            echo "  Available ports:"
            local i=1; local -a avail=()
            for n in "${nics[@]}"; do
                [[ -n "${used_nic[$n]:-}" ]] && continue
                printf "    ${B}%d)${X} %-10s %s  %s\n" "$i" "$n" "$(perm_mac "$n")" "$(nic_speed "$n")"
                avail+=("$n"); i=$((i+1))
            done
            local pick
            pick=$(prompt "  Select ${label} port [1]: " 1)
            if [[ "$pick" =~ ^[0-9]+$ ]] && (( pick >= 1 && pick <= ${#avail[@]} )); then
                chosen="${avail[$((pick-1))]}"
            else
                echo "  ${R}Invalid selection.${X}"
            fi
        done
        assign_role "$role" "$chosen" && used_nic[$chosen]="$role"
    done

    echo
    echo "  ${B}Assignment summary:${X}"
    cat "$ROLES_CONF" | sed 's/^/    /'
    echo
    echo "  ${G}Roles written. They take effect after networkd reloads / reboot.${X}"
    echo "  ${D}Stored in $ROLES_CONF and ${NETDIR}/*.link${X}"
    return 0
}

main "$@"
