#!/usr/bin/env bash
# Flash stock OpenWrt on a GL.iNet GL-MT3000 (Beryl AX) via SSH from archy
#
# Usage: ./flash-mt3000-openwrt.sh <router-password> [router-ip]
#
# Defaults to 192.168.8.1. Runs entirely from archy (archipelago).
# Requires: curl, sshpass, ssh, sha256sum, scp
#
# What it does:
#   1. Verifies tools and connectivity
#   2. Downloads the stock OpenWrt sysupgrade image
#   3. Verifies SHA256 checksum
#   4. Uploads image to the router via SCP
#   5. Flashes via sysupgrade (no settings preserved)
#   6. Waits for reboot and verifies new firmware
#
# After flash, stock OpenWrt is at 192.168.1.1 (root, no password).

set -euo pipefail

PASSWORD="${1:?Usage: $0 <router-password> [router-ip]}"
ROUTER="${2:-192.168.8.1}"
SSH_USER="root"
OPENWRT_VERSION="24.10.2"
IMAGE_NAME="openwrt-24.10.2-mediatek-filogic-glinet_gl-mt3000-squashfs-sysupgrade.bin"
IMAGE_URL="https://downloads.openwrt.org/releases/${OPENWRT_VERSION}/targets/mediatek/filogic/${IMAGE_NAME}"
DOWNLOAD_DIR="/tmp/openwrt-flash"
LOCAL_IMAGE="${DOWNLOAD_DIR}/${IMAGE_NAME}"
NEW_IP="192.168.1.1"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

log()  { echo -e "${GREEN}[+]${NC} $*"; }
warn() { echo -e "${YELLOW}[!]${NC} $*"; }
err()  { echo -e "${RED}[-]${NC} $*" >&2; }

SSH_CMD="sshpass -p ${PASSWORD} ssh -o StrictHostKeyChecking=no -o ConnectTimeout=10 ${SSH_USER}@${ROUTER}"
SCP_CMD="sshpass -p ${PASSWORD} scp -o StrictHostKeyChecking=no -o ConnectTimeout=10"

# --- Step 1: Check tools ---
log "Checking required tools..."
MISSING=()
for cmd in curl sshpass ssh sha256sum scp; do
    command -v "$cmd" >/dev/null 2>&1 || MISSING+=("$cmd")
done
if [ ${#MISSING[@]} -gt 0 ]; then
    err "Missing tools: ${MISSING[*]}"
    exit 1
fi
log "All tools found."

# --- Step 2: Check router connectivity ---
log "Checking router connectivity at ${ROUTER}..."
if ! $SSH_CMD "echo ok" >/dev/null 2>&1; then
    err "Cannot SSH into ${ROUTER}. Is the router reachable and is the password correct?"
    exit 1
fi
log "Router reachable."

# Verify it's a GL-MT3000
BOARD=$($SSH_CMD "cat /tmp/sysinfo/board_name 2>/dev/null" || true)
if [[ "$BOARD" != *"mt3000"* ]]; then
    err "Router board is '${BOARD}', expected GL-MT3000. Aborting."
    exit 1
fi
log "Confirmed GL-MT3000 (board: ${BOARD})."

# --- Step 3: Download image ---
log "Downloading OpenWrt ${OPENWRT_VERSION} sysupgrade image..."
mkdir -p "$DOWNLOAD_DIR"
if [ -f "$LOCAL_IMAGE" ]; then
    warn "Image already exists locally, skipping download."
else
    if ! curl -fSL --progress-bar -o "$LOCAL_IMAGE" "$IMAGE_URL"; then
        err "Download failed from ${IMAGE_URL}"
        exit 1
    fi
fi
log "Downloaded: $(ls -lh "$LOCAL_IMAGE" | awk '{print $5}')"

# --- Step 4: Verify checksum ---
log "Verifying SHA256 checksum..."
    EXPECTED_HASH=$(curl -fsSL "https://downloads.openwrt.org/releases/${OPENWRT_VERSION}/targets/mediatek/filogic/sha256sums" 2>/dev/null | grep "${IMAGE_NAME}" | awk '{print $1}' || true)
if [ -n "$EXPECTED_HASH" ]; then
    ACTUAL_HASH=$(sha256sum "$LOCAL_IMAGE" | awk '{print $1}')
    if [ "$EXPECTED_HASH" != "$ACTUAL_HASH" ]; then
        err "Checksum mismatch!"
        err "  Expected: ${EXPECTED_HASH}"
        err "  Actual:   ${ACTUAL_HASH}"
        exit 1
    fi
    log "Checksum OK: ${ACTUAL_HASH:0:16}..."
else
    warn "Could not fetch expected checksum, skipping verification."
fi

# --- Step 5: Pre-flash checks ---
log "Running pre-flash checks on router..."

# Check free space on /tmp
FREE_KB=$($SSH_CMD "df /tmp | tail -1 | awk '{print \$4}'")
IMAGE_SIZE_KB=$(( $(stat -c%s "$LOCAL_IMAGE") / 1024 ))
if [ "$FREE_KB" -lt "$((IMAGE_SIZE_KB + 10240))" ]; then
    err "Not enough space on /tmp. Need ~$((IMAGE_SIZE_KB/1024))MB, have $((FREE_KB/1024))MB."
    exit 1
fi
log "Free space on /tmp: $((FREE_KB/1024))MB, image: $((IMAGE_SIZE_KB/1024))MB — OK."

# Verify current firmware
CURRENT_FW=$($SSH_CMD "cat /etc/openwrt_release | grep DISTRIB_DESCRIPTION" || true)
log "Current firmware: ${CURRENT_FW}"

# --- Step 6: Upload image ---
log "Uploading image to router..."
$SCP_CMD "$LOCAL_IMAGE" ${SSH_USER}@${ROUTER}:/tmp/openwrt-sysupgrade.bin
log "Upload complete."

# Verify uploaded file
REMOTE_HASH=$($SSH_CMD "sha256sum /tmp/openwrt-sysupgrade.bin | awk '{print \$1}'")
LOCAL_HASH=$(sha256sum "$LOCAL_IMAGE" | awk '{print $1}')
if [ "$REMOTE_HASH" != "$LOCAL_HASH" ]; then
    err "Remote file hash mismatch after upload!"
    $SSH_CMD "rm -f /tmp/openwrt-sysupgrade.bin"
    exit 1
fi
log "Remote file verified."

# --- Step 7: Flash ---
log "============================================"
log "FLASHING STOCK OPENWRT ${OPENWRT_VERSION}"
log "The router will reboot. This takes ~3-5 min."
log "After flash, OpenWrt will be at ${NEW_IP}"
log "============================================"
echo ""

$SSH_CMD "sysupgrade -n /tmp/openwrt-sysupgrade.bin" &
SCP_PID=$!

# sysupgrade disconnects SSH, wait for it
wait $SCP_PID 2>/dev/null || true

log "Flash initiated. Waiting for router to reboot..."
sleep 15

# --- Step 8: Wait for new router ---
log "Waiting for stock OpenWrt at ${NEW_IP}..."
for i in $(seq 1 60); do
    if ping -c1 -W2 "$NEW_IP" >/dev/null 2>&1; then
        log "Router is up at ${NEW_IP}!"
        break
    fi
    if [ $i -eq 60 ]; then
        warn "Router not responding at ${NEW_IP} after 2 minutes."
        warn "It may still be booting. Try: ssh root@${NEW_IP}"
        warn "If unreachable, hold reset for U-Boot recovery at 192.168.1.1"
        exit 1
    fi
    printf "\r  Waiting... (%d/60)" $i
    sleep 2
done
echo ""

# --- Step 9: Verify new firmware ---
log "Verifying new firmware..."
for attempt in 1 2 3; do
    if ssh -o ConnectTimeout=10 -o StrictHostKeyChecking=no root@${NEW_IP} "cat /etc/openwrt_release" 2>/dev/null; then
        echo ""
        log "============================================"
        log "FLASH COMPLETE"
        log "Stock OpenWrt is running at ${NEW_IP}"
        log "Login: root (no password)"
        log "Run 'passwd' to set a root password."
        log "============================================"
        exit 0
    fi
    sleep 5
done

warn "Could not verify firmware. Router is reachable but SSH may still be starting."
warn "Try: ssh root@${NEW_IP}"