#!/usr/bin/env bash
# ============================================================
#  kahf-install.sh — Install Kahf Family Safety on OpenWrt
# ============================================================
#
#  What it does:
#    1. Connects to your OpenWrt router over SSH
#    2. Verifies it really is OpenWrt
#    3. Downloads the Kahf auto-update agent from s3.kahf.co
#    4. Runs it once to install OAF + family-safe packages
#    5. Sets up a nightly cron job for automatic updates
#
#  Requirements on the machine running this script:
#    - bash, ssh
#    - sshpass   (for password authentication)
#      macOS:  brew install sshpass
#      Ubuntu: sudo apt install sshpass
#      Fedora: sudo dnf install sshpass
#    - SSH key auth is used automatically if sshpass is absent
#
#  Usage:
#    Interactive:      ./kahf-install.sh
#    With arguments:   ./kahf-install.sh --ip 192.168.1.1 --user root --password mypass
#    Skip password:    ./kahf-install.sh --ip 192.168.1.1   (key auth or prompted)
#    Dev channel:      ./kahf-install.sh --channel dev

set -euo pipefail

# ── Colours ──────────────────────────────────────────────────────────────────
if [ -t 1 ] && tput colors &>/dev/null && [ "$(tput colors)" -ge 8 ]; then
    C_RESET='\033[0m'
    C_BOLD='\033[1m'
    C_BLUE='\033[1;34m'
    C_GREEN='\033[1;32m'
    C_YELLOW='\033[1;33m'
    C_RED='\033[1;31m'
    C_CYAN='\033[1;36m'
else
    C_RESET=''; C_BOLD=''; C_BLUE=''; C_GREEN=''; C_YELLOW=''; C_RED=''; C_CYAN=''
fi

info()    { printf "${C_BLUE}[*]${C_RESET} %s\n" "$*"; }
ok()      { printf "${C_GREEN}[✓]${C_RESET} %s\n" "$*"; }
warn()    { printf "${C_YELLOW}[!]${C_RESET} %s\n" "$*"; }
err()     { printf "${C_RED}[✗]${C_RESET} %s\n" "$*" >&2; }
die()     { err "$*"; echo "" >&2; exit 1; }
step()    { printf "\n${C_BOLD}${C_CYAN}──── %s ────${C_RESET}\n" "$*"; }
banner()  { printf "${C_BOLD}%s${C_RESET}\n" "$*"; }

# ── Defaults ─────────────────────────────────────────────────────────────────
ROUTER_IP=""
ROUTER_USER="root"
ROUTER_PASS=""
CHANNEL="prod"
SSH_PORT=22
SKIP_CRON=0
DRY_RUN=0

S3_BASE_PROD="https://s3.kahf.co/router-prod"
S3_BASE_DEV="https://s3.kahf.co/router-dev"
UPDATE_SCRIPT_NAME="oaf-update.sh"
INSTALL_PATH="/usr/sbin/oaf-update.sh"
CRON_FILE="/etc/crontabs/root"
CRON_ENTRY="0 3 * * * /usr/sbin/oaf-update.sh >> /var/log/oaf-update.log 2>&1"

# ── Argument parsing ─────────────────────────────────────────────────────────
while [ $# -gt 0 ]; do
    case "$1" in
        --ip|-i)          ROUTER_IP="$2";   shift 2 ;;
        --user|-u)        ROUTER_USER="$2"; shift 2 ;;
        --password|-p)    ROUTER_PASS="$2"; shift 2 ;;
        --port)           SSH_PORT="$2";    shift 2 ;;
        --channel|-c)     CHANNEL="$2";     shift 2 ;;
        --skip-cron)      SKIP_CRON=1;      shift   ;;
        --dry-run)        DRY_RUN=1;        shift   ;;
        --help|-h)
            sed -n '/^#/!q;s/^# \{0,1\}//p' "$0"
            exit 0
            ;;
        *) die "Unknown option: $1  (run with --help for usage)" ;;
    esac
done

# Resolve S3 base URL from channel
case "${CHANNEL}" in
    dev)  S3_BASE="${S3_BASE_DEV}" ;;
    prod) S3_BASE="${S3_BASE_PROD}" ;;
    *)    die "Unknown channel '${CHANNEL}'. Use 'prod' or 'dev'." ;;
esac
UPDATE_SCRIPT_URL="${S3_BASE}/${UPDATE_SCRIPT_NAME}"

# ── Banner ───────────────────────────────────────────────────────────────────
echo ""
banner "╔══════════════════════════════════════════════════════════════╗"
banner "║        Kahf Family Safety — OpenWrt Installer               ║"
banner "║   Installs OAF + family-safe packages via auto-update agent ║"
banner "╚══════════════════════════════════════════════════════════════╝"
echo ""

# ── Step 0: Interactive prompts ───────────────────────────────────────────────
step "Configuration"

if [ -z "${ROUTER_IP}" ]; then
    printf "  Router IP address [192.168.1.1]: "
    read -r ROUTER_IP
    ROUTER_IP="${ROUTER_IP:-192.168.1.1}"
fi

if [ -z "${ROUTER_PASS}" ]; then
    if command -v sshpass &>/dev/null; then
        printf "  Root password (leave blank to use SSH key): "
        read -rs ROUTER_PASS
        echo ""
    else
        info "sshpass not found — will attempt SSH key authentication."
        info "Install sshpass if you want password login:"
        info "  macOS:  brew install sshpass"
        info "  Ubuntu: sudo apt install sshpass"
    fi
fi

echo ""
info "Router:  ${ROUTER_USER}@${ROUTER_IP}:${SSH_PORT}"
info "Channel: ${CHANNEL}  (${S3_BASE})"
[ "${DRY_RUN}" = "1" ] && warn "DRY RUN — no changes will be made to the router"

# ── Helpers: SSH wrapper ──────────────────────────────────────────────────────
SSH_BASE_OPTS="-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o ConnectTimeout=10 -o LogLevel=ERROR"

_ssh_cmd() {
    if command -v sshpass &>/dev/null && [ -n "${ROUTER_PASS}" ]; then
        SSHPASS="${ROUTER_PASS}" sshpass -e \
            ssh ${SSH_BASE_OPTS} -p "${SSH_PORT}" "${ROUTER_USER}@${ROUTER_IP}" "$@"
    else
        ssh ${SSH_BASE_OPTS} -o BatchMode=yes \
            -p "${SSH_PORT}" "${ROUTER_USER}@${ROUTER_IP}" "$@"
    fi
}

# Run a command on the router; return its exit code without aborting the script
ssh_run() {
    _ssh_cmd "$@" 2>&1
}

# Run a command on the router; die on failure
ssh_run_strict() {
    local label="${1}"; shift
    if ! _ssh_cmd "$@" 2>&1; then
        die "Remote command failed: ${label}"
    fi
}

# ── Step 1: Dependency check ─────────────────────────────────────────────────
step "Dependency check"

if ! command -v ssh &>/dev/null; then
    die "ssh not found. Please install OpenSSH."
fi
ok "ssh found"

if command -v sshpass &>/dev/null; then
    ok "sshpass found"
elif [ -n "${ROUTER_PASS}" ]; then
    die "Password supplied but sshpass is not installed.\n  macOS:  brew install sshpass\n  Ubuntu: sudo apt install sshpass"
fi

# ── Step 2: Router reachability ───────────────────────────────────────────────
step "Router reachability"

info "Checking ${ROUTER_IP}:${SSH_PORT}..."
if command -v nc &>/dev/null; then
    if ! nc -z -w 5 "${ROUTER_IP}" "${SSH_PORT}" 2>/dev/null; then
        die "Cannot reach ${ROUTER_IP}:${SSH_PORT}.\n  • Is the router powered on?\n  • Is your machine on the same network as the router?\n  • Correct IP? (router admin page usually shows it)"
    fi
elif command -v ping &>/dev/null; then
    if ! ping -c 1 -W 3 "${ROUTER_IP}" &>/dev/null; then
        warn "${ROUTER_IP} did not respond to ping (may be blocked). Attempting SSH anyway..."
    fi
fi
ok "${ROUTER_IP}:${SSH_PORT} is reachable"

# ── Step 3: SSH authentication ────────────────────────────────────────────────
step "SSH authentication"

AUTH_OK=0
if command -v sshpass &>/dev/null && [ -n "${ROUTER_PASS}" ]; then
    info "Testing password authentication..."
    if SSHPASS="${ROUTER_PASS}" sshpass -e \
        ssh ${SSH_BASE_OPTS} -p "${SSH_PORT}" "${ROUTER_USER}@${ROUTER_IP}" \
        "echo __auth_ok__" 2>&1 | grep -q "__auth_ok__"; then
        AUTH_OK=1
        ok "Password authentication successful"
    else
        die "Authentication failed.\n  • Wrong password?\n  • Is root login enabled on this router?\n  • Check: System → Administration → SSH Access in LuCI"
    fi
else
    info "Testing SSH key authentication..."
    if ssh ${SSH_BASE_OPTS} -o BatchMode=yes \
        -p "${SSH_PORT}" "${ROUTER_USER}@${ROUTER_IP}" \
        "echo __auth_ok__" 2>&1 | grep -q "__auth_ok__"; then
        AUTH_OK=1
        ok "SSH key authentication successful"
    else
        err "SSH key authentication failed."
        if ! command -v sshpass &>/dev/null; then
            die "No SSH key set up and sshpass is not installed.\n  Install sshpass to use password auth:\n    macOS:  brew install sshpass\n    Ubuntu: sudo apt install sshpass\n  Then re-run with --password"
        else
            die "SSH key authentication failed. Re-run with --password <yourpassword>"
        fi
    fi
fi

[ "${AUTH_OK}" = "1" ] || die "Could not authenticate — bailing out."

# ── Step 4: OpenWrt detection ────────────────────────────────────────────────
step "OpenWrt detection"

info "Checking for OpenWrt..."
OWRT_RELEASE="$(ssh_run "cat /etc/openwrt_release 2>/dev/null || echo NOT_OPENWRT")"

if echo "${OWRT_RELEASE}" | grep -q "NOT_OPENWRT"; then
    die "This does not appear to be an OpenWrt router.\n  /etc/openwrt_release not found.\n  This installer only supports OpenWrt routers."
fi

if ! echo "${OWRT_RELEASE}" | grep -qi "openwrt"; then
    die "Unexpected /etc/openwrt_release content — cannot confirm this is OpenWrt."
fi

DISTRIB_RELEASE="$(echo "${OWRT_RELEASE}" | grep "^DISTRIB_RELEASE=" | cut -d'=' -f2 | tr -d '"')"
DISTRIB_TARGET="$(echo "${OWRT_RELEASE}" | grep "^DISTRIB_TARGET=" | cut -d'=' -f2 | tr -d '"')"
DISTRIB_ARCH="$(echo "${OWRT_RELEASE}" | grep "^DISTRIB_ARCH=" | cut -d'=' -f2 | tr -d '"')"

ok "Confirmed OpenWrt router"
info "  Version : ${DISTRIB_RELEASE:-unknown}"
info "  Target  : ${DISTRIB_TARGET:-unknown}"
info "  Arch    : ${DISTRIB_ARCH:-unknown}"

# ── Step 5: Readiness checks ─────────────────────────────────────────────────
step "Router readiness"

# Package manager
PKG_MGR="$(ssh_run "command -v apk >/dev/null 2>&1 && echo apk || command -v opkg >/dev/null 2>&1 && echo opkg || echo none")"
PKG_MGR="$(echo "${PKG_MGR}" | tr -d '[:space:]')"
case "${PKG_MGR}" in
    apk)  ok "Package manager: apk (OpenWrt 25.x+)" ;;
    opkg) ok "Package manager: opkg (OpenWrt 23.x/24.x)" ;;
    *)    die "No supported package manager found (neither apk nor opkg).\n  This installer requires OpenWrt 23.05 or newer." ;;
esac

# Disk space: need at least 4 MB free in /tmp
TMP_FREE_KB="$(ssh_run "df /tmp 2>/dev/null | awk 'NR==2{print \$4}'" | tr -d '[:space:]')"
if [ -n "${TMP_FREE_KB}" ] && [ "${TMP_FREE_KB}" -lt 4096 ] 2>/dev/null; then
    die "Less than 4 MB free in /tmp (${TMP_FREE_KB} KB). Free up space on the router."
fi
ok "Disk space: /tmp has ${TMP_FREE_KB:-?} KB free"

# Internet connectivity from the router (needed to reach S3)
info "Checking router internet access (s3.kahf.co)..."
INET_CHECK="$(ssh_run "wget -q --timeout=10 -O /dev/null https://s3.kahf.co 2>&1; echo exit:\$?")"
if echo "${INET_CHECK}" | grep -q "exit:0"; then
    ok "Router can reach s3.kahf.co"
else
    warn "Router may not have internet access. The installer will continue but"
    warn "package downloads may fail if s3.kahf.co is unreachable from the router."
fi

# ── Step 6: Install the auto-update agent ────────────────────────────────────
step "Installing auto-update agent"

info "Downloading ${UPDATE_SCRIPT_NAME} from ${S3_BASE}..."

if [ "${DRY_RUN}" = "1" ]; then
    ok "[DRY RUN] Would download ${UPDATE_SCRIPT_URL} → ${INSTALL_PATH} and run it"
else
    INSTALL_CMD="$(cat <<EOF
set -e
wget -q --timeout=30 -O '${INSTALL_PATH}' '${UPDATE_SCRIPT_URL}' || {
    echo "ERROR: Failed to download ${UPDATE_SCRIPT_URL}"
    echo "       Check that the router has internet access and that s3.kahf.co is reachable."
    exit 1
}
chmod +x '${INSTALL_PATH}'
echo "AUTO_UPDATE_AGENT_READY"
EOF
)"
    DOWNLOAD_OUT="$(ssh_run "sh" <<< "${INSTALL_CMD}")"

    if echo "${DOWNLOAD_OUT}" | grep -q "AUTO_UPDATE_AGENT_READY"; then
        ok "Auto-update agent installed at ${INSTALL_PATH}"
    else
        err "Download output:"
        echo "${DOWNLOAD_OUT}" | sed 's/^/    /'
        die "Failed to install auto-update agent. See output above."
    fi

    # Write the channel config so the update script knows where to pull from
    if [ "${CHANNEL}" != "prod" ]; then
        info "Writing channel config (${CHANNEL})..."
        ssh_run "sh -c 'echo UPDATE_CHANNEL=${CHANNEL} > /etc/oaf-update.conf'"
        ok "Channel set to: ${CHANNEL}"
    else
        # Remove any dev channel config so it defaults to prod
        ssh_run "rm -f /etc/oaf-update.conf" || true
        ok "Channel: prod (default)"
    fi
fi

# ── Step 7: Run the update agent ─────────────────────────────────────────────
step "Running initial package install"

info "This will download and install:"
info "  • kmod-oaf     — kernel traffic inspection module"
info "  • appfilter    — application-level traffic filter daemon"
info "  • luci-app-oaf — web UI for traffic filtering"
info "  • family-safe  — DNS-based family safety (content filtering)"
echo ""
info "Download + install may take 2–5 minutes depending on connection speed..."
echo ""

if [ "${DRY_RUN}" = "1" ]; then
    ok "[DRY RUN] Would run: sh ${INSTALL_PATH} --force"
else
    UPDATE_EXIT=0
    ssh_run "sh '${INSTALL_PATH}' --force" || UPDATE_EXIT=$?

    if [ "${UPDATE_EXIT}" -ne 0 ]; then
        warn "Update agent exited with code ${UPDATE_EXIT}."
        warn "Check the output above — a kmod kernel mismatch is expected on first run"
        warn "and does not prevent the other packages from installing."
    fi
fi

# ── Step 8: Verify installation ───────────────────────────────────────────────
step "Verifying installation"

if [ "${DRY_RUN}" = "1" ]; then
    ok "[DRY RUN] Would verify appfilter + family-safe are installed"
else
    VERIFY_SCRIPT="$(cat <<'EOF'
errors=0
if [ "$(command -v apk)" ]; then
    PKG_LIST="$(apk info 2>/dev/null)"
    CHECK_CMD="echo \"${PKG_LIST}\" | grep -q"
else
    PKG_LIST="$(opkg list-installed 2>/dev/null)"
    CHECK_CMD="echo \"${PKG_LIST}\" | grep -q"
fi

echo "${PKG_LIST}" | grep -q "appfilter"  && echo "INSTALLED:appfilter"  || { echo "MISSING:appfilter"; errors=$((errors+1)); }
echo "${PKG_LIST}" | grep -q "family-safe" && echo "INSTALLED:family-safe" || { echo "MISSING:family-safe"; errors=$((errors+1)); }
echo "${PKG_LIST}" | grep -q "luci-app-oaf" && echo "INSTALLED:luci-app-oaf" || echo "MISSING:luci-app-oaf (optional)"

pgrep oafd >/dev/null 2>&1 && echo "RUNNING:oafd" || echo "NOT_RUNNING:oafd"

exit ${errors}
EOF
)"
    VERIFY_OUT="$(ssh_run "sh" <<< "${VERIFY_SCRIPT}")"
    VERIFY_EXIT=$?

    while IFS= read -r line; do
        case "${line}" in
            INSTALLED:*)   ok "Installed: ${line#INSTALLED:}" ;;
            MISSING:*)     warn "Missing:   ${line#MISSING:}" ;;
            RUNNING:*)     ok "Running:   ${line#RUNNING:}" ;;
            NOT_RUNNING:*) warn "Not running: ${line#NOT_RUNNING:}" ;;
        esac
    done <<< "${VERIFY_OUT}"

    if [ "${VERIFY_EXIT}" -ne 0 ]; then
        warn "Some packages may not have installed. This can happen if:"
        warn "  • The kernel module (kmod-oaf) version doesn't match the router firmware"
        warn "  • The router doesn't have internet access"
        warn "  • The router uses a kernel version we haven't built for yet"
        warn ""
        warn "Family-safe DNS filtering may still work even without kmod-oaf."
        warn "Try re-running: ssh root@${ROUTER_IP} sh ${INSTALL_PATH} --force"
    fi
fi

# ── Step 9: Set up nightly cron ───────────────────────────────────────────────
if [ "${SKIP_CRON}" = "0" ]; then
    step "Nightly auto-update cron"

    if [ "${DRY_RUN}" = "1" ]; then
        ok "[DRY RUN] Would add cron: ${CRON_ENTRY}"
    else
        CRON_SCRIPT="$(cat <<EOF
mkdir -p "\$(dirname '${CRON_FILE}')"
touch '${CRON_FILE}'
if grep -q "oaf-update" '${CRON_FILE}' 2>/dev/null; then
    echo "CRON_ALREADY_SET"
else
    echo '${CRON_ENTRY}' >> '${CRON_FILE}'
    echo "CRON_ADDED"
fi
/etc/init.d/cron reload 2>/dev/null || /etc/init.d/cron restart 2>/dev/null || true
EOF
)"
        CRON_OUT="$(ssh_run "sh" <<< "${CRON_SCRIPT}")"
        if echo "${CRON_OUT}" | grep -q "CRON_ADDED"; then
            ok "Cron job added — packages will auto-update nightly at 03:00 UTC"
        elif echo "${CRON_OUT}" | grep -q "CRON_ALREADY_SET"; then
            ok "Cron job already configured"
        fi
    fi
fi

# ── Done ─────────────────────────────────────────────────────────────────────
echo ""
banner "╔══════════════════════════════════════════════════════════════╗"
if [ "${DRY_RUN}" = "1" ]; then
    banner "║   Dry run complete — no changes were made                   ║"
else
    banner "║   Kahf Family Safety installation complete!                 ║"
fi
banner "╚══════════════════════════════════════════════════════════════╝"
echo ""
ok "Router IP     : ${ROUTER_IP}"
ok "Update agent  : ${INSTALL_PATH}"
ok "S3 channel    : ${CHANNEL}  (${S3_BASE})"
[ "${SKIP_CRON}" = "0" ] && ok "Auto-update   : nightly at 03:00 UTC"
echo ""
info "To manually trigger an update:"
info "  ssh ${ROUTER_USER}@${ROUTER_IP} sh ${INSTALL_PATH} --force"
echo ""
info "To view update logs on the router:"
info "  ssh ${ROUTER_USER}@${ROUTER_IP} cat /var/log/oaf-update.log"
echo ""
info "To access the router web interface (LuCI):"
info "  http://${ROUTER_IP}"
echo ""