#!/bin/bash
set -Eeuo pipefail

TMP_DIR="/tmp/chr-installer.$$"
MNT_DIR="$TMP_DIR/mnt"
IMG_ZIP="$TMP_DIR/chr.img.zip"
IMG_FILE="$TMP_DIR/chr.img"
LOOP_DEV=""

cleanup() {
    set +e
    sync
    if mountpoint -q "$MNT_DIR" 2>/dev/null; then
        umount "$MNT_DIR"
    fi
    if [ -n "${LOOP_DEV:-}" ]; then
        losetup -d "$LOOP_DEV" 2>/dev/null
    fi
    rm -rf "$TMP_DIR"
}
trap cleanup EXIT

err() {
    echo "ERROR: $*" >&2
    exit 1
}

info() {
    echo "[INFO] $*"
}

warn() {
    echo "[WARN] $*" >&2
}

require_root() {
    [ "$(id -u)" -eq 0 ] || err "This script must be run as root."
}

require_cmd() {
    command -v "$1" >/dev/null 2>&1 || err "Required command not found: $1"
}

download_file() {
    local url="$1"
    local out="$2"

    if command -v curl >/dev/null 2>&1; then
        curl -fL --retry 3 --connect-timeout 15 -o "$out" "$url" || err "Failed to download: $url"
    elif command -v wget >/dev/null 2>&1; then
        wget -O "$out" "$url" || err "Failed to download: $url"
    else
        err "Neither curl nor wget is installed."
    fi
}

extract_zip() {
    local zip_file="$1"
    local out_file="$2"

    if command -v unzip >/dev/null 2>&1; then
        unzip -p "$zip_file" > "$out_file" || err "Failed to extract zip file."
    else
        err "unzip is required but not installed."
    fi
}

yesno() {
    local prompt="$1"
    local answer
    while true; do
        read -r -p "$prompt [y/N]: " answer
        case "$answer" in
            y|Y) return 0 ;;
            n|N|"") return 1 ;;
            *) echo "Please answer y or n." ;;
        esac
    done
}

choose_from_menu() {
    local title="$1"
    shift
    local options=("$@")
    local i choice

    echo
    echo "$title"
    for i in "${!options[@]}"; do
        printf "  %d) %s\n" "$((i+1))" "${options[$i]}"
    done

    while true; do
        read -r -p "Enter choice number: " choice
        if [[ "$choice" =~ ^[0-9]+$ ]] && [ "$choice" -ge 1 ] && [ "$choice" -le "${#options[@]}" ]; then
            CHOSEN_INDEX=$((choice-1))
            return 0
        fi
        echo "Invalid choice."
    done
}

detect_arch() {
    ARCH="$(uname -m)"
    case "$ARCH" in
        x86_64|i386|i486|i586|i686)
            ARCH_FAMILY="x86"
            ;;
        aarch64)
            ARCH_FAMILY="arm64"
            ;;
        *)
            err "Unsupported architecture: $ARCH"
            ;;
    esac
}

detect_boot_mode() {
    if [ -d /sys/firmware/efi ]; then
        BOOT_MODE="UEFI"
    else
        BOOT_MODE="BIOS"
    fi
}

detect_default_net() {
    ETH="$(ip route show default 2>/dev/null | awk '/^default/ {print $5; exit}')"
    GATEWAY="$(ip route show default 2>/dev/null | awk '/^default/ {print $3; exit}')"

    if [ -n "${ETH:-}" ]; then
        ADDRESS="$(ip -o -4 addr show dev "$ETH" 2>/dev/null | awk '{print $4; exit}')"
    else
        ADDRESS=""
    fi

    DNS="$(awk '/^nameserver/ {print $2; exit}' /etc/resolv.conf 2>/dev/null || true)"
    [ -n "${DNS:-}" ] || DNS="8.8.8.8"
}

ask_network() {
    detect_default_net

    echo
    echo "Network configuration:"
    read -r -p "IP address/CIDR [${ADDRESS:-none}]: " INPUT
    ADDRESS="${INPUT:-$ADDRESS}"
    [ -n "${ADDRESS:-}" ] || err "IP address cannot be empty."

    read -r -p "Gateway [${GATEWAY:-none}]: " INPUT
    GATEWAY="${INPUT:-$GATEWAY}"
    [ -n "${GATEWAY:-}" ] || err "Gateway cannot be empty."

    read -r -p "DNS [${DNS:-8.8.8.8}]: " INPUT
    DNS="${INPUT:-$DNS}"
    [ -n "${DNS:-}" ] || err "DNS cannot be empty."
}

select_storage() {
    local disks=()
    local labels=()
    local line name size model label root_src root_disk

    root_src="$(findmnt -n -o SOURCE / || true)"
    root_disk="$(lsblk -no PKNAME "$root_src" 2>/dev/null | head -n1 || true)"

    while read -r name size model; do
        [ -n "$name" ] || continue
        label="/dev/$name | size=$size"
        [ -n "${model:-}" ] && label="$label | model=$model"
        if [ -n "${root_disk:-}" ] && [ "$name" = "$root_disk" ]; then
            label="$label | CURRENT_SYSTEM_DISK"
        fi
        disks+=("$name")
        labels+=("$label")
    done < <(lsblk -d -n -o NAME,SIZE,MODEL | awk '$1 !~ /^loop/')

    [ "${#disks[@]}" -gt 0 ] || err "No disk found."

    choose_from_menu "Available storage devices:" "${labels[@]}"
    STORAGE="${disks[$CHOSEN_INDEX]}"
}

validate_platform_vs_version() {
    local selected_arch="$1"

    if [ "$ARCH_FAMILY" != "$selected_arch" ]; then
        err "Selected image architecture ($selected_arch) does not match system architecture ($ARCH_FAMILY)."
    fi
}

resolve_image_url() {
    local normal_url="$1"
    local bios_url="$2"
    local major="$3"

    if [ "$BOOT_MODE" = "BIOS" ] && [ "$major" = "7" ]; then
        [ -n "$bios_url" ] || err "No BIOS legacy image URL defined for this v7 version."
        SELECTED_URL="$bios_url"
    else
        [ -n "$normal_url" ] || err "No normal image URL defined."
        SELECTED_URL="$normal_url"
    fi
}

download_and_prepare_image() {
    mkdir -p "$TMP_DIR" "$MNT_DIR"

    info "Downloading image..."
    download_file "$SELECTED_URL" "$IMG_ZIP"

    info "Extracting image..."
    extract_zip "$IMG_ZIP" "$IMG_FILE"

    [ -s "$IMG_FILE" ] || err "Extracted image is empty."
}

mount_image_partition() {
    local partition_suffix

    LOOP_DEV="$(losetup -Pf --show "$IMG_FILE")" || err "Failed to create loop device."
    sleep 1

    if [ "$SELECTED_MAJOR" -ge 7 ]; then
        partition_suffix="p2"
    else
        partition_suffix="p1"
    fi

    mount "${LOOP_DEV}${partition_suffix}" "$MNT_DIR" || err "Failed to mount image partition ${LOOP_DEV}${partition_suffix}."
}

write_autorun() {
    local autorun_path="$MNT_DIR/rw/autorun.scr"

    [ -d "$MNT_DIR/rw" ] || err "rw directory not found inside image."

    cat > "$autorun_path" <<EOF
/ip dns set servers=$DNS
/ip address add address=$ADDRESS interface=ether1
/ip route add gateway=$GATEWAY
EOF

    sync
    info "autorun.scr created successfully."
}

write_image_to_disk() {
    echo
    warn "ALL DATA ON /dev/$STORAGE WILL BE LOST."
    yesno "Do you want to continue?" || err "Operation aborted by user."

    info "Writing image to /dev/$STORAGE ..."
    dd if="$IMG_FILE" of="/dev/$STORAGE" bs=4M conv=fsync status=progress || err "dd failed."

    sync
    info "Image written successfully."
}

reboot_system() {
    echo
    info "Rebooting system..."
    echo 1 > /proc/sys/kernel/sysrq 2>/dev/null || true
    echo b > /proc/sysrq-trigger 2>/dev/null || reboot -f
}

VERSIONS=(
    "RouterOS CHR 7.18.2 Stable|7|x86|http://iso.viphost.cloud/ISO/chr/chr-7.21.3.img.zip|http://iso.viphost.cloud/ISO/chr/chr-7.21.3-legacy-bios.img.zip"
    "RouterOS CHR 7.22.3 Stable|7|x86|http://iso.viphost.cloud/ISO/chr/chr-7.22.3.img.zip|http://iso.viphost.cloud/ISO/chr/chr-7.22.3-legacy-bios.img.zip"
)

main() {
    require_root
    require_cmd ip
    require_cmd lsblk
    require_cmd dd
    require_cmd losetup
    require_cmd mount
    require_cmd umount
    require_cmd awk
    require_cmd findmnt
    require_cmd unzip

    [ "${#VERSIONS[@]}" -gt 0 ] || err "No versions are defined in the script."

    detect_arch
    detect_boot_mode

    echo "System architecture: $ARCH"
    echo "Boot mode: $BOOT_MODE"

    local labels=()
    local entry

    for entry in "${VERSIONS[@]}"; do
        labels+=("$(echo "$entry" | cut -d'|' -f1)")
    done

    choose_from_menu "Available MikroTik versions:" "${labels[@]}"
    SELECTED_ENTRY="${VERSIONS[$CHOSEN_INDEX]}"

    SELECTED_LABEL="$(echo "$SELECTED_ENTRY" | cut -d'|' -f1)"
    SELECTED_MAJOR="$(echo "$SELECTED_ENTRY" | cut -d'|' -f2)"
    SELECTED_ARCH="$(echo "$SELECTED_ENTRY" | cut -d'|' -f3)"
    SELECTED_NORMAL_URL="$(echo "$SELECTED_ENTRY" | cut -d'|' -f4)"
    SELECTED_BIOS_URL="$(echo "$SELECTED_ENTRY" | cut -d'|' -f5)"

    validate_platform_vs_version "$SELECTED_ARCH"
    resolve_image_url "$SELECTED_NORMAL_URL" "$SELECTED_BIOS_URL" "$SELECTED_MAJOR"

    echo
    echo "Selected version: $SELECTED_LABEL"
    echo "Selected image URL: $SELECTED_URL"

    select_storage
    ask_network
    download_and_prepare_image
    mount_image_partition
    write_autorun

    if mountpoint -q "$MNT_DIR"; then
        umount "$MNT_DIR"
    fi
    if [ -n "${LOOP_DEV:-}" ]; then
        losetup -d "$LOOP_DEV"
        LOOP_DEV=""
    fi

    write_image_to_disk
    reboot_system
}

main "$@"