#!/usr/bin/env bash
# Flash stock OpenWrt on a GL.iNet GL-MT3000 (Beryl AX)
#
# Usage:
#   ./flash-mt3000-openwrt.sh <router-password> [router-ip] [--image /path/to/image.bin]
#
# Defaults to 192.168.8.1.
# Requires: curl, sshpass, ssh, sha256sum
# Optional: scp (falls back to pipe if unavailable on router)
#
# What it does:
#   1. Verifies tools and connectivity
#   2. Downloads the stock OpenWrt sysupgrade image (or uses a pre-downloaded one)
#   3. Verifies SHA256 checksum
#   4. Uploads image to the router (SCP or pipe)
#   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

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

# --- Parse args ---
PASSWORD=""
ROUTER="192.168.8.1"
LOCAL_IMAGE=""

while [[ $# -gt 0 ]]; do
    case "$1" in
        --image)
            LOCAL_IMAGE="$2"
            shift 2
            ;;
        --router)
            ROUTER="$2"
            shift 2
            ;;
        *)
            if [ -z "$PASSWORD" ]; then
                PASSWORD="$1"
            else
                ROUTER="$1"
            fi
            shift
            ;;
    esac
done

if [ -z "$PASSWORD" ]; then
    echo "Usage: $0 <router-password> [router-ip] [--image /path/to/sysupgrade.bin]"
    echo ""
    echo "Options:"
    echo "  --image PATH   Use a pre-downloaded sysupgrade image instead of downloading"
    echo "  --router IP    Router IP (default: 192.168.8.1)"
    exit 1
fi

SSH_USER="root"
OPENWRT_VERSION="24.10.2"
IMAGE_NAME="openwrt-${OPENWRT_VERSION}-mediatek-filogic-glinet_gl-mt3000-squashfs-sysupgrade.bin"
IMAGE_URL="https://downloads.openwrt.org/releases/${OPENWRT_VERSION}/targets/mediatek/filogic/${IMAGE_NAME}"
CHECKSUM_URL="https://downloads.openwrt.org/releases/${OPENWRT_VERSION}/targets/mediatek/filogic/sha256sums"
DOWNLOAD_DIR="/tmp/openwrt-flash"
DEFAULT_IMAGE="${DOWNLOAD_DIR}/${IMAGE_NAME}"
NEW_IP="192.168.1.1"

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

# --- Step 1: Check tools ---
log "Checking required tools..."
MISSING=()
for cmd in curl sshpass ssh sha256sum; 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: Get image ---
if [ -n "$LOCAL_IMAGE" ]; then
    # User provided a local image
    if [ ! -f "$LOCAL_IMAGE" ]; then
        err "Image file not found: ${LOCAL_IMAGE}"
        exit 1
    fi
    log "Using provided image: ${LOCAL_IMAGE}"
else
    # Download the image
    log "Downloading OpenWrt ${OPENWRT_VERSION} sysupgrade image..."
    mkdir -p "$DOWNLOAD_DIR"

    if [ -f "$DEFAULT_IMAGE" ]; then
        warn "Image already exists locally, skipping download."
    else
        # Test if HTTPS is being intercepted (common with GL.iNet routers)
        CERT_ISSUE=false
        CERT_INFO=$(curl -v --connect-timeout 5 "https://downloads.openwrt.org" 2>&1 || true)
        if echo "$CERT_INFO" | grep -qi "GLiNet\|gl-inet\|self-signed"; then
            CERT_ISSUE=true
            warn "HTTPS interception detected — the router is MITM-ing TLS connections."
            warn "This is the same issue that breaks Tailscale."
        fi

        if [ "$CERT_ISSUE" = true ]; then
            warn "Attempting download via HTTP..."
            HTTP_URL="${IMAGE_URL/https:/http:}"
            if ! curl -fSL --progress-bar --connect-timeout 10 -o "$DEFAULT_IMAGE" "$HTTP_URL" 2>/dev/null; then
                err ""
                err "Download failed. The router is intercepting HTTPS and HTTP is also blocked."
                err ""
                err "Workaround: download the image on a machine NOT behind this router,"
                err "then run this script with --image /path/to/${IMAGE_NAME}"
                err ""
                err "Direct download URL:"
                err "  ${IMAGE_URL}"
                err ""
                exit 1
            fi
        else
            if ! curl -fSL --progress-bar -o "$DEFAULT_IMAGE" "$IMAGE_URL"; then
                err "Download failed from ${IMAGE_URL}"
                exit 1
            fi
        fi
    fi
    LOCAL_IMAGE="$DEFAULT_IMAGE"
fi

log "Image ready: $(ls -lh "$LOCAL_IMAGE" | awk '{print $5}')"

# --- Step 4: Verify checksum ---
log "Verifying SHA256 checksum..."
EXPECTED_HASH=""
CHECKSUM_ATTEMPTS=(
    "https://downloads.openwrt.org/releases/${OPENWRT_VERSION}/targets/mediatek/filogic/sha256sums"
    "http://downloads.openwrt.org/releases/${OPENWRT_VERSION}/targets/mediatek/filogic/sha256sums"
)

for url in "${CHECKSUM_ATTEMPTS[@]}"; do
    EXPECTED_HASH=$(curl -fsSL --connect-timeout 5 "$url" 2>/dev/null | grep "${IMAGE_NAME}" | awk '{print $1}' || true)
    if [ -n "$EXPECTED_HASH" ]; then
        break
    fi
done

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 (network issue?), 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 2>/dev/null | grep DISTRIB_DESCRIPTION" || true)
log "Current firmware: ${CURRENT_FW}"

# --- Step 6: Upload image ---
log "Uploading image to router..."

# Try SCP first, fall back to pipe (GL.iNet firmware lacks sftp-server)
if command -v scp >/dev/null 2>&1 && \
   sshpass -p "$PASSWORD" scp -o StrictHostKeyChecking=no -o ConnectTimeout=5 -o BatchMode=yes "$LOCAL_IMAGE" "${SSH_USER}@${ROUTER}:/tmp/openwrt-sysupgrade.bin" 2>/dev/null; then
    log "Upload complete (via SCP)."
else
    log "SCP unavailable on router, using pipe transfer..."
    cat "$LOCAL_IMAGE" | $SSH_CMD "cat > /tmp/openwrt-sysupgrade.bin"
    log "Upload complete (via pipe)."
fi

# 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 ---
echo ""
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 90); do
    if ping -c1 -W2 "$NEW_IP" >/dev/null 2>&1; then
        log "Router is up at ${NEW_IP}!"
        break
    fi
    if [ $i -eq 90 ]; then
        warn "Router not responding at ${NEW_IP} after 3 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/90)" $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=accept-new 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}"