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

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

usage() {
    cat <<'EOF'
Usage:
  ./tools/mac/device.sh detect
  ./tools/mac/device.sh backup [board]
  ./tools/mac/device.sh verify <dist/board>
  ./tools/mac/device.sh flash <dist/board> [--yes]
  ./tools/mac/device.sh restore <backup.bin> [--yes]
  ./tools/mac/device.sh monitor
  ./tools/mac/device.sh erase [--yes]
  ./tools/mac/device.sh wifi

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

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

choose_port() {
    if [[ -n "${ESPPORT:-}" ]]; then
        if [[ ! -e "$ESPPORT" && "${DEVICE_LAB_TEST_MODE:-0}" != "1" ]]; 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))]}"
}

esptool() {
    "$python_bin" -m esptool --chip esp32s3 --port "$1" "${@:2}"
}

confirm() {
    local prompt="$1"
    local assume_yes="${2:-}"
    if [[ "$assume_yes" == "--yes" ]]; then
        return
    fi
    local answer
    read -r -p "$prompt [y/N] " answer
    [[ "$answer" == "y" || "$answer" == "Y" ]] || exit 1
}

detect() {
    local port="$1"
    echo "Serial port: $port"
    esptool "$port" chip_id
    esptool "$port" flash_id
}

flash_size_bytes() {
    case "$1" in
        2MB) echo $((2 * 1024 * 1024)) ;;
        4MB) echo $((4 * 1024 * 1024)) ;;
        8MB) echo $((8 * 1024 * 1024)) ;;
        16MB) echo $((16 * 1024 * 1024)) ;;
        32MB) echo $((32 * 1024 * 1024)) ;;
        *) echo "Unsupported or undetected flash size: $1" >&2; exit 1 ;;
    esac
}

backup_flash() {
    local port="$1"
    local board="${2:-esp32-s3}"
    local info
    info="$(esptool "$port" flash_id 2>&1)"
    echo "$info"
    local size
    size="$(printf '%s\n' "$info" | sed -n 's/.*Detected flash size: \([0-9][0-9]*MB\).*/\1/p' | tail -1)"
    [[ -n "$size" ]] || { echo "Could not detect flash size" >&2; exit 1; }
    local bytes
    bytes="$(flash_size_bytes "$size")"
    local timestamp
    timestamp="$(date -u +%Y%m%d-%H%M%S)"
    local backup_dir="$repo_dir/backups"
    local image="$backup_dir/$board-$timestamp.bin"
    local metadata="$backup_dir/$board-$timestamp.json"
    mkdir -p "$backup_dir"

    echo "Reading complete $size flash to $image"
    esptool "$port" read_flash 0 "$bytes" "$image"

    local mac_info
    mac_info="$(esptool "$port" read_mac 2>&1)"
    local mac
    mac="$(printf '%s\n' "$mac_info" | sed -n 's/^MAC: //p' | tail -1)"
    "$python_bin" - "$metadata" "$board" "$size" "$mac" "$port" "$timestamp" <<'PY'
import json, sys
path, board, size, mac, port, timestamp = sys.argv[1:]
with open(path, "w", encoding="utf-8") as handle:
    json.dump({
        "board": board,
        "chip": "esp32s3",
        "mac": mac,
        "flash_size": size,
        "serial_port": port,
        "timestamp_utc": timestamp,
    }, handle, indent=2)
    handle.write("\n")
PY
    echo "Backup complete: $image"
    echo "Metadata: $metadata"
}

flash_firmware() {
    local port="$1"
    local artifact_dir="$2"
    local assume_yes="${3:-}"

    verify_artifact "$artifact_dir"
    detect "$port"
    local expected_size flash_info actual_size
    expected_size="$("$python_bin" -c 'import json,sys; print(json.load(open(sys.argv[1]))["flash_size"])' "$artifact_dir/flash.json")"
    flash_info="$(esptool "$port" flash_id 2>&1)"
    actual_size="$(printf '%s\n' "$flash_info" | sed -n 's/.*Detected flash size: \([0-9][0-9]*MB\).*/\1/p' | tail -1)"
    [[ "$actual_size" == "$expected_size" ]] || {
        echo "Flash size mismatch: artifact=$expected_size device=${actual_size:-unknown}" >&2
        exit 1
    }
    echo "Safety reminder: back up factory firmware before the first custom flash."
    confirm "Flash $artifact_dir/merged-flash.bin to $port?" "$assume_yes"
    esptool "$port" --baud 460800 write_flash 0x0 "$artifact_dir/merged-flash.bin"
    echo "Flash complete and device reset requested."
    if [[ "$assume_yes" != "--yes" ]]; then
        local answer
        read -r -p "Open serial monitor now? [Y/n] " answer
        if [[ "$answer" != "n" && "$answer" != "N" ]]; then
            monitor "$port"
        fi
    fi
}

verify_artifact() {
    local artifact_dir="$1"
    [[ -f "$artifact_dir/flash.json" ]] || { echo "Missing $artifact_dir/flash.json" >&2; exit 1; }
    [[ -f "$artifact_dir/merged-flash.bin" ]] || { echo "Missing $artifact_dir/merged-flash.bin" >&2; exit 1; }

    local expected_chip
    expected_chip="$("$python_bin" -c 'import json,sys; print(json.load(open(sys.argv[1]))["chip"])' "$artifact_dir/flash.json")"
    [[ "$expected_chip" == "esp32s3" ]] || { echo "Artifact is for $expected_chip, not esp32s3" >&2; exit 1; }
    local expected_sha actual_sha
    expected_sha="$("$python_bin" -c 'import json,sys; print(json.load(open(sys.argv[1])).get("merged_image_sha256", ""))' "$artifact_dir/flash.json")"
    [[ -n "$expected_sha" ]] || { echo "Artifact metadata has no merged image SHA-256" >&2; exit 1; }
    actual_sha="$("$python_bin" -c 'import hashlib,sys; print(hashlib.sha256(open(sys.argv[1], "rb").read()).hexdigest())' "$artifact_dir/merged-flash.bin")"
    [[ "$actual_sha" == "$expected_sha" ]] || { echo "Artifact SHA-256 mismatch; download it again" >&2; exit 1; }
    echo "Artifact verified: $artifact_dir/merged-flash.bin"
}

restore_flash() {
    local port="$1"
    local image="$2"
    local assume_yes="${3:-}"
    [[ -f "$image" ]] || { echo "Backup not found: $image" >&2; exit 1; }
    detect "$port"
    local flash_info size bytes image_bytes
    flash_info="$(esptool "$port" flash_id 2>&1)"
    size="$(printf '%s\n' "$flash_info" | sed -n 's/.*Detected flash size: \([0-9][0-9]*MB\).*/\1/p' | tail -1)"
    [[ -n "$size" ]] || { echo "Could not detect flash size" >&2; exit 1; }
    bytes="$(flash_size_bytes "$size")"
    image_bytes="$(wc -c < "$image" | tr -d ' ')"
    [[ "$image_bytes" == "$bytes" ]] || {
        echo "Backup size mismatch: file=$image_bytes bytes device=$bytes bytes" >&2
        exit 1
    }
    confirm "Restore the full image $image to $port?" "$assume_yes"
    esptool "$port" --baud 460800 write_flash 0x0 "$image"
    echo "Restore complete."
}

monitor() {
    local port="$1"
    echo "Monitoring $port at 115200 baud. Exit with Ctrl-]."
    "$python_bin" -m serial.tools.miniterm "$port" 115200 --raw
}

# ---- Wi-Fi provisioning over the existing USB Serial protocol -------------
# The device speaks a versioned framed protocol:
#   WFP/1 <verb> [key=value ...]   ->   WFP/1 OK|ERROR ... / DATA ... / END
# The Python helper owns the read loop so timeouts and framing stay consistent
# across every subcommand. A password is typed with hidden input and is never
# echoed by the device or printed by this script.
wifi_ctl() {
    local port="$1"; shift
    "$python_bin" "$tool_dir/wifi_ctl.py" "$port" "$@"
}

wifi_interactive() {
    local port="$1"
    echo "Device: M5Stack StopWatch"
    echo "Port:   $port"
    echo
    while true; do
        echo "Known networks:"
        wifi_ctl "$port" list || true
        echo
        cat <<'EOF'
Actions:
1. Add network
2. Forget network
3. Forget all
4. Show status
5. Exit
EOF
        local choice
        read -r -p "Select [1-5]: " choice
        case "$choice" in
            1)
                local ssid
                read -r -p "SSID: " ssid
                [[ -n "$ssid" ]] || { echo "SSID is required" >&2; continue; }
                # The password is prompted for by wifi_ctl.py through getpass, so
                # it never appears in this shell's argv or environment.
                wifi_ctl "$port" add --ssid "$ssid"
                ;;
            2)
                local ssid
                read -r -p "SSID to forget: " ssid
                [[ -n "$ssid" ]] || { echo "SSID is required" >&2; continue; }
                wifi_ctl "$port" remove --ssid "$ssid"
                ;;
            3)
                local answer
                read -r -p "Forget ALL known networks? [y/N] " answer
                [[ "$answer" == "y" || "$answer" == "Y" ]] || continue
                wifi_ctl "$port" clear
                ;;
            4) wifi_ctl "$port" status ;;
            5) return 0 ;;
            *) echo "Invalid selection" >&2 ;;
        esac
        echo
    done
}

erase_flash() {
    local port="$1"
    local assume_yes="${2:-}"
    echo "WARNING: this erases the entire flash, including factory firmware and NVS."
    confirm "Erase all flash on $port?" "$assume_yes"
    esptool "$port" erase_flash
}

require_tools
command="${1:-}"
case "$command" in
    detect)
        port="$(choose_port)"; detect "$port" ;;
    backup)
        port="$(choose_port)"; backup_flash "$port" "${2:-esp32-s3}" ;;
    verify)
        [[ -n "${2:-}" ]] || { usage; exit 2; }
        verify_artifact "$(cd "$2" && pwd)" ;;
    flash)
        [[ -n "${2:-}" ]] || { usage; exit 2; }
        port="$(choose_port)"; flash_firmware "$port" "$(cd "$2" && pwd)" "${3:-}" ;;
    restore)
        [[ -n "${2:-}" ]] || { usage; exit 2; }
        port="$(choose_port)"; restore_flash "$port" "$(cd "$(dirname "$2")" && pwd)/$(basename "$2")" "${3:-}" ;;
    monitor)
        port="$(choose_port)"; monitor "$port" ;;
    wifi)
        port="$(choose_port)"; wifi_interactive "$port" ;;
    erase)
        port="$(choose_port)"; erase_flash "$port" "${2:-}" ;;
    *)
        usage; exit 2 ;;
esac