#!/usr/bin/env bash
# Post-flash Wi-Fi provisioning.
#
# Flashing firmware does not carry Wi-Fi credentials: the device keeps them in
# its own NVS store, and the build-time migration only fires on a genuine first
# boot (an absent store). This script re-applies credentials over the existing
# USB provisioning protocol after a flash, so a refresh does not mean re-typing
# the password by hand.
#
# Credentials are read from a file, never from argv, so the password cannot leak
# into the process list or the shell history. The default file is `.env` at the
# repo root, which is already ignored by Git:
#
#     DEVICE_WIFI_SSID=Tencent-LabWiFi
#     DEVICE_WIFI_PASSWORD=...
#
# Use `--ssid` to override just the SSID (for example a second network); the
# password still comes from the file.
#
#   ./tools/mac/provision-wifi.sh [--ssid NAME] [--credentials PATH] [--yes]
#
# Set ESPPORT=/dev/cu.usbmodem... to select a specific port non-interactively.
set -euo pipefail

tool_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
repo_dir="$(cd "$tool_dir/../.." && pwd)"
python_bin="$tool_dir/.venv/bin/python"

ssid_override=""
credentials="$repo_dir/.env"
assume_yes=""

usage() {
    cat <<'EOF'
Usage: ./tools/mac/provision-wifi.sh [options]

Re-applies Wi-Fi credentials to a flashed device over USB.

Options:
  --ssid NAME          Override the SSID (password still read from the file)
  --credentials PATH   Credentials file (default: <repo>/.env)
  --yes                Do not prompt for confirmation
  -h, --help           Show this help

The credentials file must define:
  DEVICE_WIFI_SSID=...
  DEVICE_WIFI_PASSWORD=...

Set ESPPORT=/dev/cu.usbmodem... to select a specific port non-interactively.
EOF
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --ssid)
            [[ -n "${2:-}" ]] || { echo "--ssid needs a value" >&2; exit 2; }
            ssid_override="$2"; shift 2 ;;
        --credentials)
            [[ -n "${2:-}" ]] || { echo "--credentials needs a path" >&2; exit 2; }
            credentials="$2"; shift 2 ;;
        --yes)
            assume_yes="--yes"; shift ;;
        -h|--help)
            usage; exit 0 ;;
        *)
            echo "Unknown option: $1" >&2; usage >&2; exit 2 ;;
    esac
done

if [[ ! -x "$python_bin" ]]; then
    echo "Run ./tools/mac/install.sh first." >&2
    exit 1
fi

# Emits base64 so the values survive command substitution intact: bash strips
# NUL bytes from `$(...)`, so a NUL separator would silently corrupt the split.
# `set -e` aborts the script when this helper fails, so no exit code is echoed.
load_credentials() {
    "$python_bin" - "$1" <<'PY'
import base64
import sys

path = sys.argv[1]
try:
    with open(path, "r", encoding="utf-8") as handle:
        raw = handle.read()
except OSError as exc:
    print(f"cannot read credentials file: {exc}", file=sys.stderr)
    raise SystemExit(3)

wanted = {"DEVICE_WIFI_SSID", "DEVICE_WIFI_PASSWORD"}
values = {}
for line in raw.replace("\r\n", "\n").splitlines():
    line = line.strip()
    if not line or line.startswith("#") or "=" not in line:
        continue
    key, _, value = line.partition("=")
    key = key.strip()
    if key not in wanted:
        continue
    value = value.strip()
    if len(value) >= 2 and value[0] == value[-1] and value[0] in ("'", '"'):
        value = value[1:-1]
    values[key] = value

if not values.get("DEVICE_WIFI_SSID"):
    print("credentials file has no DEVICE_WIFI_SSID", file=sys.stderr)
    raise SystemExit(3)
if not values.get("DEVICE_WIFI_PASSWORD"):
    print("credentials file has no DEVICE_WIFI_PASSWORD", file=sys.stderr)
    raise SystemExit(3)

for key in ("DEVICE_WIFI_SSID", "DEVICE_WIFI_PASSWORD"):
    print(base64.b64encode(values[key].encode("utf-8")).decode("ascii"))
PY
}

if [[ ! -f "$credentials" ]]; then
    echo "Credentials file not found: $credentials" >&2
    echo "Create it with DEVICE_WIFI_SSID and DEVICE_WIFI_PASSWORD (it is git-ignored)." >&2
    exit 1
fi

creds="$(load_credentials "$credentials")"
ssid="$(printf '%s\n' "$creds" | sed -n '1p' | "$python_bin" -c 'import base64,sys; sys.stdout.write(base64.b64decode(sys.stdin.read().strip()).decode("utf-8"))')"
password="$(printf '%s\n' "$creds" | sed -n '2p' | "$python_bin" -c 'import base64,sys; sys.stdout.write(base64.b64decode(sys.stdin.read().strip()).decode("utf-8"))')"
if [[ -n "$ssid_override" ]]; then
    ssid="$ssid_override"
fi
if [[ -z "$ssid" || -z "$password" ]]; then
    echo "SSID and password must both be non-empty." >&2
    exit 1
fi

# Port selection mirrors tools/mac/device.sh so both accept ESPPORT identically.
choose_port() {
    if [[ -n "${ESPPORT:-}" ]]; then
        if [[ ! -e "$ESPPORT" ]]; then
            echo "ESPPORT does not exist: $ESPPORT" >&2
            exit 1
        fi
        printf '%s\n' "$ESPPORT"
        return
    fi

    shopt -s nullglob
    local ports=(/dev/cu.usbmodem* /dev/cu.usbserial* /dev/cu.SLAB_USBtoUART* /dev/cu.wchusbserial*)
    shopt -u nullglob
    if [[ ${#ports[@]} -eq 0 ]]; then
        echo "No ESP32 serial port found. Check the USB data cable and download mode." >&2
        exit 1
    fi
    if [[ ${#ports[@]} -eq 1 ]]; then
        printf '%s\n' "${ports[0]}"
        return
    fi

    echo "Multiple serial ports found:" >&2
    local index=1
    for port in "${ports[@]}"; do
        echo "  $index) $port" >&2
        index=$((index + 1))
    done
    local selection
    read -r -p "Select port [1-${#ports[@]}]: " selection
    if [[ ! "$selection" =~ ^[0-9]+$ ]] || (( selection < 1 || selection > ${#ports[@]} )); then
        echo "Invalid selection" >&2
        exit 1
    fi
    printf '%s\n' "${ports[$((selection - 1))]}"
}

port="$(choose_port)"
echo "Device: M5Stack StopWatch"
echo "Port:   $port"
echo "SSID:   $ssid"

if [[ "$assume_yes" != "--yes" ]]; then
    local_answer=""
    read -r -p "Apply this network to the device? [y/N] " local_answer
    [[ "$local_answer" == "y" || "$local_answer" == "Y" ]] || exit 1
fi

# The password is fed on stdin, never as an argument or an environment value:
# `ps` cannot see it, `/proc/<pid>/environ` cannot see it, and the device
# protocol never echoes it back. Port, SSID and a base64 password travel as
# three lines, so an embedded newline in the password cannot break the framing.
# The helper is a real file (not a heredoc) so stdin stays free for the payload.
helper="$(mktemp "${TMPDIR:-/tmp}/provision-wifi.XXXXXX.py")"
trap 'rm -f "$helper"' EXIT
cat > "$helper" <<'PY'
import base64
import sys

# argv[1] is the tools directory (not a file path) so `import wifi_ctl`
# resolves regardless of the caller's working directory.
sys.path.insert(0, sys.argv[1])
import wifi_ctl

raw = sys.stdin.buffer.read().decode("utf-8")
lines = raw.split("\n")
port = lines[0]
ssid = lines[1]
password = base64.b64decode(lines[2].strip()).decode("utf-8")

result = wifi_ctl.provision(port, ssid, password)
verdict = result.get("result", "")
sys.stderr.write(wifi_ctl.redact(wifi_ctl.format_result(result)) + "\n")
raise SystemExit(0 if verdict == "VERIFIED" else 1)
PY

{
    printf '%s\n' "$port"
    printf '%s\n' "$ssid"
    printf '%s\n' "$(printf '%s' "$password" | "$python_bin" -c 'import base64,sys; sys.stdout.write(base64.b64encode(sys.stdin.buffer.read()).decode("ascii"))')"
} | "$python_bin" "$helper" "$tool_dir"