#!/usr/bin/env bash
#
# AIS-catcher Installation Script
#
# Description:
#   Automated installation script for AIS-catcher on Debian and Ubuntu systems.
#   Supports both package-based and source-based installations.
#
# Usage:
#   sudo ./aiscatcher-install [OPTIONS]
#
# Options:
#   -p, --package    Install from pre-built package (faster, recommended)
#   --no-systemd     Skip systemd service installation (useful for containers)
#   -h, --help       Display this help message
#
# Requirements:
#   - Must be run as root (use sudo)
#   - Debian-based system (Debian/Ubuntu/Raspbian)
#   - Internet connection for downloading packages/source
#
# Author: AIS-catcher Project
# License: See LICENSE file
#

set -euo pipefail

# Display usage information
usage() {
    cat << EOF
AIS-catcher Installation Script

Usage: sudo $0 [OPTIONS]

Options:
  -p, --package              Install from pre-built package (faster, recommended)
  -b, --branch               Specify git branch to clone (default: main)
  --no-systemd               Skip systemd service installation (useful for containers)
  --no-user                  Skip user creation and run as root (useful for containers)
  --set-reboot-on-failure [BURST [SEC]]  Enable reboot watchdog (default: 3 restarts in 1800s)
  --unset-reboot-on-failure              Disable automatic reboot on repeated failures
  --set-auto-restart                     Enable automatic service restart on crash (Restart=always)
  --unset-auto-restart                   Disable automatic service restart on crash (Restart=no)
  -h, --help                 Display this help message

Examples:
  sudo $0 --package                    # Install from pre-built package
  sudo $0                              # Build and install from source (main branch)
  sudo $0 -b develop                   # Build from develop branch
  sudo $0 --no-systemd --no-user       # Docker/container installation
  sudo $SCRIPT_NAME --set-reboot-on-failure      # Enable reboot watchdog (no reinstall needed)
  sudo $SCRIPT_NAME --unset-reboot-on-failure    # Disable reboot watchdog (no reinstall needed)
  sudo $SCRIPT_NAME --set-auto-restart           # Enable auto-restart on crash (no reinstall needed)
  sudo $SCRIPT_NAME --unset-auto-restart         # Disable auto-restart on crash (no reinstall needed)

For more information, visit: https://docs.aiscatcher.org/
EOF
    exit 0
}

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    _hint_args=$(printf '%q ' "$@")
    if [[ -t 1 ]]; then
        echo -e "\033[1;31mError: This script must be run as root\033[0m" >&2
        echo -e "\033[1;33mPlease run: sudo $0 ${_hint_args% }\033[0m" >&2
    else
        echo "Error: This script must be run as root" >&2
        echo "Please run: sudo $0 ${_hint_args% }" >&2
    fi
    exit 1
fi

# Constants
SCRIPT_NAME=$(basename "$0")
readonly SCRIPT_NAME
LOG_FILE="/var/log/${SCRIPT_NAME%.*}.log"
readonly LOG_FILE
TMP_DIR=$(mktemp -d -t AIS-catcher-XXXXXX)
readonly TMP_DIR
readonly CONFIG_DIR="/etc/AIS-catcher"
readonly PLUGIN_DIR="${CONFIG_DIR}/plugins"
readonly DBMS_DIR="${CONFIG_DIR}/DBMS"
readonly README_DIR="${CONFIG_DIR}/README"
readonly LIB_DIR="/usr/lib/ais-catcher"
readonly CONFIG_FILE="${CONFIG_DIR}/config.json"
readonly CMD_FILE="${CONFIG_DIR}/config.cmd"
readonly SERVICE_NAME="ais-catcher.service"
readonly SERVICE_FILE="/etc/systemd/system/${SERVICE_NAME}"
readonly REBOOT_SERVICE_NAME="ais-catcher-reboot.service"
readonly REBOOT_SERVICE_FILE="/etc/systemd/system/${REBOOT_SERVICE_NAME}"
readonly REBOOT_BURST_DEFAULT=3
readonly REBOOT_INTERVAL_DEFAULT=1800

# Global variables
INSTALL_PACKAGE=false
SKIP_SYSTEMD=false
SKIP_USER=false
BRANCH="main"
SERVICE_STATE=""
SET_REBOOT=""
REBOOT_BURST="${REBOOT_BURST_DEFAULT}"
REBOOT_INTERVAL="${REBOOT_INTERVAL_DEFAULT}"
SET_AUTO_RESTART=""
INSTALL_REQUESTED=false

# Color constants (using ANSI escape codes)
if [[ -t 1 ]]; then
    readonly COLOR_RESET="\033[0m"
    readonly COLOR_BOLD="\033[1m"
    readonly COLOR_RED="\033[0;31m"
    readonly COLOR_GREEN="\033[0;32m"
    readonly COLOR_YELLOW="\033[0;33m"
    readonly COLOR_BLUE="\033[0;34m"
    readonly COLOR_CYAN="\033[0;36m"
    readonly COLOR_BOLD_RED="\033[1;31m"
    readonly COLOR_BOLD_GREEN="\033[1;32m"
    readonly COLOR_BOLD_YELLOW="\033[1;33m"
    readonly COLOR_BOLD_BLUE="\033[1;34m"
else
    readonly COLOR_RESET=""
    readonly COLOR_BOLD=""
    readonly COLOR_RED=""
    readonly COLOR_GREEN=""
    readonly COLOR_YELLOW=""
    readonly COLOR_BLUE=""
    readonly COLOR_CYAN=""
    readonly COLOR_BOLD_RED=""
    readonly COLOR_BOLD_GREEN=""
    readonly COLOR_BOLD_YELLOW=""
    readonly COLOR_BOLD_BLUE=""
fi

# Logger function with color support
log() {
    local level="$1"
    shift
    local color=""
    local timestamp="$(date +'%Y-%m-%d %H:%M:%S')"
    
    case "$level" in
        ERROR)
            color="${COLOR_BOLD_RED}"
            ;;
        WARN)
            color="${COLOR_BOLD_YELLOW}"
            ;;
        INFO)
            color="${COLOR_BOLD_GREEN}"
            ;;
        *)
            color="${COLOR_RESET}"
            ;;
    esac
    
    # Print to terminal with color; %s ensures backslashes in message text are not interpreted
    printf "${COLOR_CYAN}[${timestamp}]${COLOR_RESET} ${color}[${level}]${COLOR_RESET} %s\n" "$*"
    # Log to file without color codes
    printf '[%s] [%s] %s\n' "${timestamp}" "${level}" "$*" >> "$LOG_FILE"
}

# Error handler
error_exit() {
    log "ERROR" "$1"
    # Try to restore service state before exiting
    if [[ -n "$SERVICE_STATE" ]]; then
        log "INFO" "Attempting to restore service state before exit"
        restore_service_state || log "WARN" "Failed to restore service state during error exit"
    fi
    exit 1
}

# Cleanup function
cleanup() {
    log "INFO" "Cleaning up temporary directory ${TMP_DIR}"
    rm -rf "$TMP_DIR"
}

# Cancel any pending reboot and stop the watchdog service.
# Safe to call unconditionally — harmless if nothing is pending or running.
cancel_reboot_watchdog() {
    shutdown -c 2>/dev/null || true
    systemctl stop "$REBOOT_SERVICE_NAME" 2>/dev/null || true
    log "INFO" "Reboot watchdog cancelled (pending shutdown cleared, service stopped)"
}

# Replace a directive if present in FILE, otherwise insert it under its section header.
sed_replace_or_insert() {
    local key="$1" val="$2" section="$3" file="$4"
    if grep -q "^${key}=" "$file"; then
        sed -i "s/^${key}=.*/${key}=${val}/" "$file"
    else
        sed -i "/^\[${section}\]/a ${key}=${val}" "$file"
    fi
}

# Reload systemd and reset the service failure counter.
reload_and_reset_service() {
    systemctl daemon-reload || error_exit "Failed to reload systemd"
    log "INFO" "Resetting failure counter for ${SERVICE_NAME}"
    systemctl reset-failed "$SERVICE_NAME" 2>/dev/null || log "WARN" "Failed to reset failure counter"
}

# Log whether new policy settings will take effect now or at next restart.
# Neither toggle handler restarts the service, so we just check current state for messaging.
log_service_policy_notice() {
    local policy_desc="$1"
    if systemctl is-active --quiet "$SERVICE_NAME" 2>/dev/null; then
        log "INFO" "${SERVICE_NAME} is running — ${policy_desc} will apply from the next restart"
    else
        log "INFO" "${SERVICE_NAME} was not running — configuration updated"
    fi
}

################################################################################
# REBOOT-ON-FAILURE TOGGLE (standalone, no reinstall needed)
################################################################################

# Enable or disable the reboot watchdog by updating StartLimitBurst/Interval in the service file
handle_reboot_setting() {
    if [[ ! -f "$SERVICE_FILE" ]]; then
        error_exit "Service file not found at ${SERVICE_FILE} — is AIS-catcher installed?"
    fi

    if [[ "$SET_REBOOT" == "true" ]]; then
        sed_replace_or_insert StartLimitBurst    "${REBOOT_BURST}"    Unit "$SERVICE_FILE"
        sed_replace_or_insert StartLimitIntervalSec "${REBOOT_INTERVAL}" Unit "$SERVICE_FILE"
        # Add OnFailure= if not already present
        if ! grep -q "^OnFailure=" "$SERVICE_FILE"; then
            sed -i "/^After=network.target/a OnFailure=${REBOOT_SERVICE_NAME}" "$SERVICE_FILE"
        fi
        log "INFO" "Reboot on failure enabled (burst: ${REBOOT_BURST}, interval: ${REBOOT_INTERVAL}s)"
    else
        sed_replace_or_insert StartLimitBurst    0 Unit "$SERVICE_FILE"
        sed_replace_or_insert StartLimitIntervalSec 0 Unit "$SERVICE_FILE"
        # Remove OnFailure= line — with burst=0 it fires on every crash, not just limit-hit
        sed -i "/^OnFailure=/d" "$SERVICE_FILE"
        log "INFO" "Reboot on failure disabled"
    fi

    reload_and_reset_service

    # No restart needed: Restart=, StartLimitBurst=, OnFailure= are systemd policy directives
    # that take effect on the next crash/start — daemon-reload is sufficient.
    log_service_policy_notice "new watchdog settings"
}

################################################################################
# AUTO-RESTART TOGGLE (standalone, no reinstall needed)
################################################################################

# Enable or disable automatic service restart by updating Restart= in the service file
handle_auto_restart_setting() {
    if [[ ! -f "$SERVICE_FILE" ]]; then
        error_exit "Service file not found at ${SERVICE_FILE} — is AIS-catcher installed?"
    fi

    if [[ "$SET_AUTO_RESTART" == "true" ]]; then
        sed_replace_or_insert Restart always Service "$SERVICE_FILE"
        log "INFO" "Auto-restart enabled (Restart=always)"
    else
        sed_replace_or_insert Restart no Service "$SERVICE_FILE"
        log "INFO" "Auto-restart disabled (Restart=no)"
    fi

    reload_and_reset_service

    # No restart needed: Restart= is a systemd policy directive that takes effect on the
    # next crash/start — daemon-reload is sufficient and avoids unnecessary interruption.
    log_service_policy_notice "new restart policy"
}

################################################################################
# UTILITY FUNCTIONS (Used by both package and source installations)
################################################################################

# Check for conflicting existing service
check_existing_service() {
    if systemctl is-enabled aiscatcher.service &>/dev/null; then
        log "ERROR" "Detected existing aiscatcher.service. Please disable it first with:"
        log "ERROR" "systemctl disable aiscatcher.service"
        log "ERROR" "Then run this script again."
        exit 1
    fi
}

# Detect system information (OS, version, architecture)
get_system_info() {
    local os version arch
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        if [ "${ID:-}" = "raspbian" ]; then
            os="debian"
        else
            os="${ID:-}"
        fi
        version="${VERSION_CODENAME:-}"
        [[ -z "$os" ]]      && error_exit "Could not determine OS ID from /etc/os-release"
        [[ -z "$version" ]] && error_exit "Could not determine VERSION_CODENAME from /etc/os-release"
    else
        os=$(uname -s)  || error_exit "Failed to determine OS name (uname -s)"
        version=$(uname -r) || error_exit "Failed to determine OS version (uname -r)"
    fi
    arch=$(dpkg --print-architecture) || error_exit "Failed to determine system architecture (dpkg --print-architecture)"

    echo "${os}_${version}_${arch}"
}

download_webassets_common() {
    local webassets_url="https://github.com/jvde-github/webassets/archive/refs/heads/main.zip"
    local temp_zip
    temp_zip=$(mktemp --suffix=.zip)

    log "INFO" "Downloading webassets"
    if curl -fsL "$webassets_url" -o "$temp_zip"; then
        if [[ -d "$CONFIG_DIR/webassets" ]]; then
            log "INFO" "Removing old webassets"
            rm -rf "$CONFIG_DIR/webassets"
        fi
        
        log "INFO" "Extracting webassets"
        unzip -q "$temp_zip" -d "$CONFIG_DIR"
        mv "$CONFIG_DIR/webassets-main" "$CONFIG_DIR/webassets"
        rm "$temp_zip"
        
        log "INFO" "Webassets updated successfully"
    else
        log "ERROR" "Failed to download webassets"
        return 1
    fi
}

################################################################################
# SOURCE INSTALLATION FUNCTIONS
################################################################################

# Install build dependencies for source compilation
install_source_build_dependencies() {
    local required_deps="git cmake build-essential pkg-config libssl-dev zlib1g-dev libusb-1.0-0-dev curl unzip"
    local optional_deps="libzmq3-dev libpq-dev libsqlite3-dev libairspy-dev libairspyhf-dev libhackrf-dev"
    
    log "INFO" "Installing required build dependencies: $required_deps"
    apt-get install -y $required_deps || error_exit "Failed to install required dependencies"
    
    log "INFO" "Installing optional dependencies (failures allowed): $optional_deps"
    for pkg in $optional_deps; do
        if apt-get install -y "$pkg"; then
            log "INFO" "Installed $pkg"
        else
            log "WARN" "Could not install $pkg - corresponding feature may be unavailable"
        fi
    done
}

# Install librtlsdr from source
install_librtlsdr_from_source() {
    log "INFO" "Installing librtlsdr from source"
    git clone https://gitea.osmocom.org/sdr/rtl-sdr --depth 1 || error_exit "Failed to clone rtl-sdr repository"
    cd rtl-sdr || error_exit "Failed to change directory to rtl-sdr"
    mkdir build && cd build || error_exit "Failed to create and enter build directory"
    cmake ../ -DCMAKE_INSTALL_PREFIX="${LIB_DIR}" -DDETACH_KERNEL_DRIVER=ON -DINSTALL_UDEV_RULES=ON || error_exit "Failed to run cmake"
    make || error_exit "Failed to build rtl-sdr"
    make install || error_exit "Failed to install rtl-sdr"
    # ldconfig not needed: libraries are in a custom prefix and the RPATH in the AIS-catcher
    # binary points directly to ${LIB_DIR}/lib, so the dynamic linker finds them without it.
    cd ../.. || error_exit "Failed to return to parent directory"
}

# Install libhydrasdr from source
install_libhydrasdr_from_source() {
    log "INFO" "Installing libhydrasdr from source"
    git clone https://github.com/hydrasdr/rfone_host.git --depth 1 || error_exit "Failed to clone rfone_host repository"
    cd rfone_host/libhydrasdr || error_exit "Failed to change directory to rfone_host/libhydrasdr"
    mkdir build && cd build || error_exit "Failed to create and enter build directory"
    cmake .. -DCMAKE_INSTALL_PREFIX="${LIB_DIR}" -DINSTALL_UDEV_RULES=ON || error_exit "Failed to run cmake"
    make || error_exit "Failed to build libhydrasdr"
    make install || error_exit "Failed to install libhydrasdr"
    cd ../../.. || error_exit "Failed to return to parent directory"
}

# Install NMEA2000 library from source
install_nmea2000_from_source() {
    log "INFO" "Building NMEA2000 library from source"
    git clone https://github.com/jvde-github/NMEA2000.git --depth 1 || error_exit "Failed to clone NMEA2000 repository"
    cd NMEA2000/src || error_exit "Failed to change directory to NMEA2000/src"
    
    # Build as static library
    log "INFO" "Building NMEA2000 as static library"
    g++ -O3 -c N2kMsg.cpp N2kStream.cpp N2kMessages.cpp N2kTimer.cpp NMEA2000.cpp N2kGroupFunctionDefaultHandlers.cpp N2kGroupFunction.cpp -I. || error_exit "Failed to compile NMEA2000"
    ar rcs libnmea2000.a *.o || error_exit "Failed to create NMEA2000 static library"
    
    cd ../.. || error_exit "Failed to return to parent directory"
}

# Install SDR libraries from source (only for source build)
install_sdr_libraries_from_source() {
    # Create library directory
    mkdir -p "${LIB_DIR}" || error_exit "Failed to create library directory"
    
    # Install rtlsdr and hydrasdr from source (need custom cmake flags for udev rules)
    install_librtlsdr_from_source
    install_libhydrasdr_from_source
    
    # Install NMEA2000 library
    install_nmea2000_from_source
}

# Build AIS-catcher from source
build_ais_catcher_from_source() {
    log "INFO" "Building AIS-catcher from branch: ${BRANCH}"

    git clone https://github.com/jvde-github/AIS-catcher.git --depth 1 --branch "${BRANCH}" || error_exit "Failed to download AIS-catcher (branch: ${BRANCH})"
    cd AIS-catcher

    # Set environment variables to find custom libraries
    export PKG_CONFIG_PATH="${LIB_DIR}/lib/pkgconfig:${PKG_CONFIG_PATH:-}"
    export CMAKE_PREFIX_PATH="${LIB_DIR}:${CMAKE_PREFIX_PATH:-}"
    
    # Build with rpath to custom library directory
    mkdir -p build && cd build
    cmake .. -DCMAKE_BUILD_TYPE=Release \
             -DCMAKE_INSTALL_RPATH="${LIB_DIR}/lib;${LIB_DIR}" \
             -DCMAKE_BUILD_RPATH="${LIB_DIR}/lib;${LIB_DIR}" \
             -DNMEA2000_INCLUDE="${TMP_DIR}/NMEA2000/src" \
             -DNMEA2000_LIB="${TMP_DIR}/NMEA2000/src/libnmea2000.a" || error_exit "Failed to run cmake"
    make || error_exit "Failed to build AIS-catcher"
    
    cd ../..
    install -m 755 AIS-catcher/build/AIS-catcher /usr/bin/AIS-catcher || error_exit "Failed to install AIS-catcher executable"
}

# Install additional files (plugins, DBMS, README, LICENSE) from source
install_source_additional_files() {
    # Ensure the config directory exists before copying files into it.
    # Do not rely on mkdir -p "$PLUGIN_DIR" creating the parent as a side-effect.
    mkdir -p "$CONFIG_DIR" || error_exit "Failed to create config directory"

    log "INFO" "Installing plugins"
    mkdir -p "$PLUGIN_DIR" || error_exit "Failed to create plugin directory"
    cp "${TMP_DIR}/AIS-catcher/plugins/"* "$PLUGIN_DIR" || error_exit "Failed to copy plugins"

    log "INFO" "Installing DBMS"
    mkdir -p "$DBMS_DIR" || error_exit "Failed to create DBMS directory"
    cp "${TMP_DIR}/AIS-catcher/Source/DBMS/create.sql" "$DBMS_DIR" || error_exit "Failed to copy DBMS"

    log "INFO" "Installing README"
    mkdir -p "$README_DIR" || error_exit "Failed to create README directory"
    cp "${TMP_DIR}/AIS-catcher/README.md" "$README_DIR" || error_exit "Failed to copy README"

    log "INFO" "Copying LICENSE"
    cp "${TMP_DIR}/AIS-catcher/LICENSE" /etc/AIS-catcher/LICENSE || error_exit "Failed to copy LICENSE"
}

################################################################################
# PACKAGE INSTALLATION FUNCTIONS
################################################################################

# Check if a pre-built package is available for this system
# Args:
#   $1: system_info string (format: os_version_arch)
# Returns:
#   0 if package found (prints package name), 1 otherwise
check_package_availability() {
    local system_info=$1
    local packages=(
        "ais-catcher_debian_bookworm_amd64.deb"
        "ais-catcher_debian_bookworm_arm64.deb"
        "ais-catcher_debian_bookworm_armhf.deb"
        "ais-catcher_debian_bullseye_amd64.deb"
        "ais-catcher_debian_bullseye_arm64.deb"
        "ais-catcher_debian_bullseye_armhf.deb"
        "ais-catcher_debian_trixie_amd64.deb"
        "ais-catcher_debian_trixie_arm64.deb"
        "ais-catcher_debian_trixie_armhf.deb"
        "ais-catcher_ubuntu_focal_amd64.deb"
        "ais-catcher_ubuntu_focal_arm64.deb"
        "ais-catcher_ubuntu_jammy_amd64.deb"
        "ais-catcher_ubuntu_jammy_arm64.deb"
        "ais-catcher_ubuntu_jammy_armhf.deb"
        "ais-catcher_ubuntu_noble_amd64.deb"
        "ais-catcher_ubuntu_noble_arm64.deb"
        "ais-catcher_ubuntu_noble_armhf.deb"
        "ais-catcher_ubuntu_plucky_amd64.deb"
        "ais-catcher_ubuntu_plucky_arm64.deb"
        "ais-catcher_ubuntu_plucky_armhf.deb"
        "ais-catcher_ubuntu_questing_amd64.deb"
        "ais-catcher_ubuntu_questing_arm64.deb"
        "ais-catcher_ubuntu_questing_armhf.deb"
        "ais-catcher_ubuntu_resolute_amd64.deb"
        "ais-catcher_ubuntu_resolute_arm64.deb"
        "ais-catcher_ubuntu_resolute_armhf.deb"
    )

    for package in "${packages[@]}"; do
        if [[ $package == *"$system_info"* ]]; then
            echo "$package"
            return 0
        fi
    done

    return 1
}

# Download and install a pre-built Debian package by filename.
# Separated from install_package_method so the download/install logic can be
# tested or reused independently of the method-selection layer.
download_and_install_deb() {
    local package=$1
    local download_url="https://github.com/jvde-github/AIS-catcher/releases/download/Edge/$package"
   
    log "INFO" "Installing AIS-catcher from package: $package"

    # Create a temporary directory with appropriate permissions
    local temp_dir=$(mktemp -d)
    chmod 755 "$temp_dir"

    # Download the package
    log "INFO" "Downloading package..."
    if ! curl -fsL -o "$temp_dir/$(basename "$download_url")" "$download_url"; then
        rm -rf "$temp_dir"
        log "WARN" "Failed to download package"
        return 1
    fi

    # Ensure the downloaded file has correct permissions
    chmod 644 "$temp_dir/$package"

    # Use --allow-downgrades to handle version conflicts and better error handling
    log "INFO" "Installing package and dependencies..."
    local install_success=false
    if apt-get install -y --allow-downgrades "$temp_dir/$package"; then
        install_success=true
        log "INFO" "Package installed successfully"
    else
        log "WARN" "Standard package installation failed, trying with additional flags..."
        if apt-get install -y --allow-downgrades --allow-change-held-packages "$temp_dir/$package"; then
            install_success=true
            log "INFO" "Package installed successfully with additional flags"
        fi
    fi

    # Clean up temporary directory
    rm -rf "$temp_dir" || log "WARN" "Failed to remove temporary directory: $temp_dir"

    if [[ "$install_success" != true ]]; then
        log "WARN" "Failed to install package after all attempts"
        return 1
    fi

    log "INFO" "AIS-catcher package installation completed"
}

# Resolve which .deb package to use for this system and drive download_and_install_deb.
install_package_method() {
    local system_info
    system_info=$(get_system_info)
    log "INFO" "Detected system: $system_info"

    # Declare separately so the assignment exit-code is not masked by 'local'.
    local package
    package=$(check_package_availability "$system_info") || true

    if [ -z "$package" ]; then
        log "WARN" "No pre-built package available for $system_info"
        return 1
    fi

    log "INFO" "Package available: $package"

    # Install minimal dependencies for package installation and webassets download
    log "INFO" "Installing minimal dependencies for package installation"
    apt-get install -y curl unzip || error_exit "Failed to install minimal dependencies"

    download_and_install_deb "$package" || return 1
    log "INFO" "AIS-catcher installation from package completed successfully"
}

################################################################################
# CONFIGURATION FUNCTIONS (Common to both installation methods)
################################################################################

# Create and configure aiscatcher system user
setup_aiscatcher_user() {
    if [ "$SKIP_USER" = true ]; then
        log "INFO" "Skipping user creation (--no-user flag specified), will run as root"
        return 0
    fi
    
    log "INFO" "Setting up aiscatcher system user"
    
    # Create system user if it doesn't exist
    if ! id -u aiscatcher &>/dev/null; then
        log "INFO" "Creating aiscatcher system user"
        useradd --system --no-create-home --shell /usr/sbin/nologin aiscatcher || error_exit "Failed to create aiscatcher user"
    else
        log "INFO" "User aiscatcher already exists"
    fi
    
    # Add user to necessary groups for SDR device access
    for group in plugdev dialout; do
        if getent group "$group" &>/dev/null; then
            if ! id -nG aiscatcher | grep -qw "$group"; then
                log "INFO" "Adding aiscatcher to $group group"
                usermod -a -G "$group" aiscatcher || log "WARN" "Failed to add aiscatcher to $group group"
            fi
        else
            log "WARN" "Group $group does not exist, skipping"
        fi
    done
}

# Setup configuration
setup_configuration() {
    log "INFO" "Setting up configuration files"
    mkdir -p "$CONFIG_DIR" || error_exit "Failed to create config directory"

    download_webassets_common || error_exit "Failed to download webassets"

    if [[ ! -f "$CONFIG_FILE" ]]; then
        echo '{"config":"aiscatcher","version":1,"screen":0,"verbose":true,"verbose_time":60,"sharing":true,"udp":[],"tcp":[],"server":{"active":false,"plugin_dir":"/etc/AIS-catcher/plugins","cdn":"/etc/AIS-catcher/webassets","file":"/etc/AIS-catcher/stat.bin","realtime":false,"geojson":false,"prome":false,"port":"8100","backup":"10","context":"settings"},"rtlsdr":{"tuner":"auto","bandwidth":"192K","sample_rate":"1536K","biastee":false,"rtlagc":true},"airspyhf":{"threshold":"low","preamp":false},"rtltcp":{"protocol":"none","sample_rate":"288K","bandwidth":"0","rtlagc":false}}' | tee "$CONFIG_FILE" > /dev/null || error_exit "Failed to create config.json"
    fi
    
    if [[ ! -f "$CMD_FILE" ]]; then
        cat << EOF | tee "$CMD_FILE" > /dev/null || error_exit "Failed to create config.cmd"
# AIS-catcher configuration

EOF
    fi
}

################################################################################
# SYSTEMD SERVICE FUNCTIONS (Common to both installation methods)
################################################################################

# Install the reboot watchdog poll script
install_reboot_script() {
    local script_path="${LIB_DIR}/wait-and-reboot.sh"

    mkdir -p "$LIB_DIR" || error_exit "Failed to create lib directory"

    cat > "$script_path" << 'SCRIPT'
#!/bin/bash
# Schedule a reboot in 5 minutes; cancel it if AIS-catcher recovers in time.
# Checks every 10 seconds; cancels the pending shutdown if the service comes
# back active (e.g. after manual "systemctl reset-failed + start" intervention).
# To cancel a pending reboot manually: shutdown -c
REBOOT_MINUTES=5
shutdown -r "+${REBOOT_MINUTES}" "AIS-catcher failed repeatedly - rebooting in ${REBOOT_MINUTES} minutes. To cancel: shutdown -c"
for i in $(seq 30); do
    sleep 10
    if systemctl is-active --quiet ais-catcher.service; then
        shutdown -c
        exit 0
    fi
done
SCRIPT

    chmod 755 "$script_path" || error_exit "Failed to set permissions on reboot script"
    log "INFO" "Reboot watchdog script installed to ${script_path}"
}

# Setup systemd service
setup_systemd_service() {
    log "INFO" "Setting up systemd service"

    install_reboot_script

    # Write the reboot watchdog unit — fires only when OnFailure is triggered (StartLimitBurst > 0)
    # Type=simple: systemd spawns the script and does not wait for it to exit, which avoids
    # hitting TimeoutStartSec (90s default) during the 5-minute cancellation window.
    cat > "$REBOOT_SERVICE_FILE" << EOF
[Unit]
Description=Reboot system due to repeated AIS-catcher failures

[Service]
Type=simple
ExecStart=${LIB_DIR}/wait-and-reboot.sh
EOF
    log "INFO" "Reboot watchdog service written to ${REBOOT_SERVICE_FILE}"

    # Preserve existing burst/interval/restart from current service file if present
    local burst=0
    local interval=0
    local restart_policy="always"
    if [[ -f "$SERVICE_FILE" ]]; then
        burst=$(grep "^StartLimitBurst=" "$SERVICE_FILE" | cut -d= -f2 || echo 0)
        interval=$(grep "^StartLimitIntervalSec=" "$SERVICE_FILE" | cut -d= -f2 || echo 0)
        local existing_restart
        existing_restart=$(grep "^Restart=" "$SERVICE_FILE" | cut -d= -f2 || true)
        burst=${burst:-0}
        interval=${interval:-0}
        [[ "$burst" =~ ^[0-9]+$ ]] || burst=0
        [[ "$interval" =~ ^[0-9]+$ ]] || interval=0
        restart_policy=${existing_restart:-always}
        log "INFO" "Preserving existing service limits: burst=${burst} interval=${interval}s restart=${restart_policy}"
    fi

    # Apply explicitly requested toggle overrides when combined with an installation
    if [[ "$SET_REBOOT" == "true" ]]; then
        burst="$REBOOT_BURST"
        interval="$REBOOT_INTERVAL"
        log "INFO" "Overriding watchdog settings: burst=${burst} interval=${interval}s"
    elif [[ "$SET_REBOOT" == "false" ]]; then
        burst=0
        interval=0
        log "INFO" "Disabling watchdog as requested"
    fi
    if [[ "$SET_AUTO_RESTART" == "true" ]]; then
        restart_policy="always"
        log "INFO" "Overriding auto-restart: always"
    elif [[ "$SET_AUTO_RESTART" == "false" ]]; then
        restart_policy="no"
        log "INFO" "Disabling auto-restart as requested"
    fi

    # Build main service file content
    # OnFailure= is only meaningful when StartLimitBurst > 0; with burst=0 the service
    # restarts indefinitely via Restart=always and OnFailure would fire on every crash.
    local on_failure_line=""
    if [[ "${burst}" -gt 0 ]]; then
        on_failure_line="OnFailure=${REBOOT_SERVICE_NAME}"
    fi
    local service_content="[Unit]
Description=AIS-catcher Service
After=network.target
${on_failure_line:+${on_failure_line}$'\n'}StartLimitIntervalSec=${interval}
StartLimitBurst=${burst}

[Service]"

    # Only add User/Group directives if not skipping user creation
    if [ "$SKIP_USER" = false ]; then
        service_content+="
User=aiscatcher
Group=aiscatcher
SupplementaryGroups=plugdev dialout"
    fi

    service_content+="
ExecStart=/bin/bash -c '/usr/bin/AIS-catcher -G system on -o 0 -C ${CONFIG_FILE} \$\$(/bin/grep -v \"^#\" ${CMD_FILE} | /bin/grep -v \"^[[:space:]]*\$\$\" | tr -d \"\\r\" | /usr/bin/tr \"\\n\" \" \")'
Restart=${restart_policy}
RestartSec=10

[Install]
WantedBy=multi-user.target"

    echo "$service_content" | tee "$SERVICE_FILE" > /dev/null || error_exit "Failed to create service file"
    systemctl daemon-reload || error_exit "Failed to reload systemd"
}

# Function to stop the service and save its state
stop_disable_and_save_service_state() {
    log "INFO" "Checking AIS-catcher service state, stopping if active, and disabling"
    local was_active=false
    local was_enabled=false

    # Check if the service is active
    if systemctl is-active --quiet "$SERVICE_NAME"; then
        was_active=true
        log "INFO" "$SERVICE_NAME is active, attempting to stop"
        if systemctl stop "$SERVICE_NAME"; then
            log "INFO" "$SERVICE_NAME stopped successfully"
        else
            log "WARN" "Failed to stop $SERVICE_NAME"
        fi
    elif systemctl is-failed --quiet "$SERVICE_NAME" 2>/dev/null; then
        # Service is in a failed/crash-loop state — it was intended to be running
        was_active=true
        log "INFO" "$SERVICE_NAME is in failed state, resetting before update"
        systemctl reset-failed "$SERVICE_NAME" || log "WARN" "Failed to reset failed state"
    else
        log "INFO" "$SERVICE_NAME is not active, no need to stop"
    fi

    # Check if the service is enabled
    if systemctl is-enabled --quiet "$SERVICE_NAME"; then
        was_enabled=true
        log "INFO" "$SERVICE_NAME is enabled, attempting to disable"
        if systemctl disable "$SERVICE_NAME"; then
            log "INFO" "$SERVICE_NAME disabled successfully"
        else
            log "WARN" "Failed to disable $SERVICE_NAME"
        fi
    else
        log "INFO" "$SERVICE_NAME is not enabled, no need to disable"
    fi

    # Save the state to the global variable
    SERVICE_STATE="${was_active},${was_enabled}"
    log "INFO" "Service state saved: $SERVICE_STATE"
}

# FIXED: Improved service state restoration with better error handling
restore_service_state() {
    if [[ -z "$SERVICE_STATE" ]]; then
        log "INFO" "No service state to restore"
        return 0
    fi

    log "INFO" "Restoring AIS-catcher service state"
    
    IFS=',' read -ra STATE_ARRAY <<< "$SERVICE_STATE"
    local was_active=${STATE_ARRAY[0]:-false}
    local was_enabled=${STATE_ARRAY[1]:-false}
    
    # Enable first if it was enabled
    if [[ "$was_enabled" == true ]]; then
        if systemctl enable "$SERVICE_NAME"; then
            log "INFO" "$SERVICE_NAME enabled successfully"
        else
            log "ERROR" "Failed to enable $SERVICE_NAME"
            return 1
        fi
    fi
    
    # Then start if it was active
    if [[ "$was_active" == true ]]; then
        # Clear any stale failed state from before the update so start isn't blocked
        systemctl reset-failed "$SERVICE_NAME" 2>/dev/null || true
        if systemctl start "$SERVICE_NAME"; then
            log "INFO" "$SERVICE_NAME started successfully"
            local started=false
            for i in {1..10}; do
                systemctl is-active --quiet "$SERVICE_NAME" && { started=true; break; }
                sleep 1
            done
            if [[ "$started" == true ]]; then
                log "INFO" "$SERVICE_NAME is running properly"
            else
                log "WARN" "$SERVICE_NAME may have failed to start properly"
                systemctl status "$SERVICE_NAME" --no-pager || true
            fi
        else
            log "ERROR" "Failed to start $SERVICE_NAME"
            return 1
        fi
    fi
    
    # Clear the global state after restore
    SERVICE_STATE=""
    log "INFO" "Service state restored and cleared"
    return 0
}

################################################################################
# INSTALLATION METHOD FUNCTIONS
################################################################################

# Parse command-line arguments
parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -p|--package)
                log "INFO" "Package installation requested"
                INSTALL_PACKAGE=true
                INSTALL_REQUESTED=true
                shift
                ;;
            -b|--branch)
                if [[ -n "${2:-}" && "${2:0:1}" != "-" ]]; then
                    BRANCH="$2"
                    INSTALL_REQUESTED=true
                    log "INFO" "Will clone branch: ${BRANCH}"
                    shift 2
                else
                    error_exit "Option -b/--branch requires a branch name argument"
                fi
                ;;
            --no-systemd)
                log "INFO" "Skipping systemd service installation"
                SKIP_SYSTEMD=true
                shift
                ;;
            --no-user)
                log "INFO" "Skipping user creation, will run as root"
                SKIP_USER=true
                shift
                ;;
            --set-reboot-on-failure)
                SET_REBOOT="true"
                if [[ -n "${2:-}" && "${2}" =~ ^[0-9]+$ ]]; then
                    REBOOT_BURST="$2"
                    shift
                fi
                if [[ -n "${2:-}" && "${2}" =~ ^[0-9]+$ ]]; then
                    REBOOT_INTERVAL="$2"
                    shift
                fi
                shift
                ;;
            --unset-reboot-on-failure)
                SET_REBOOT="false"
                shift
                ;;
            --set-auto-restart)
                SET_AUTO_RESTART="true"
                shift
                ;;
            --unset-auto-restart)
                SET_AUTO_RESTART="false"
                shift
                ;;
            -h|--help)
                usage
                ;;
            *)
                log "ERROR" "Unknown option: $1"
                log "INFO" "Use --help for usage information"
                exit 1
                ;;
        esac
    done

    # Warn if --branch is ignored because --package was also specified
    if [[ "$INSTALL_PACKAGE" == true && "$BRANCH" != "main" ]]; then
        log "WARN" "--branch is ignored when --package is used (pre-built packages are always Edge)"
    fi
}



# Install from source
install_from_source_method() {
    log "INFO" "AIS-catcher installation from source"
    
    cd "$TMP_DIR" || error_exit "Failed to change to temporary directory"
    
    install_source_build_dependencies
    install_sdr_libraries_from_source
    build_ais_catcher_from_source
    install_source_additional_files
    
    log "INFO" "AIS-catcher source installation completed successfully"
}

# Install binary (delegates to package or source method, with source fallback)
install_binary() {
    if [ "$INSTALL_PACKAGE" = true ]; then
        if install_package_method; then
            return 0
        fi
        log "WARN" "Package installation failed, falling back to source build..."
    fi
    install_from_source_method
}

# Configure installation (user, config, permissions)
configure_installation() {
    setup_aiscatcher_user
    setup_configuration

    # Ensure correct ownership after config files are created (skip if running as root)
    if [ "$SKIP_USER" = false ]; then
        log "INFO" "Ensuring correct ownership of configuration directory"
        chown -R aiscatcher:aiscatcher "$CONFIG_DIR" || log "WARN" "Failed to set ownership of config directory"
    fi
}

# Main function
main() {
    # Phase 1: Logging and cleanup trap (TMP_DIR is created at script startup)
    log "INFO" "Starting AIS-catcher installation script"
    trap cleanup EXIT

    # Phase 2: Parse arguments first — ensures --help exits cleanly and standalone
    # toggle flags don't hit the installation pre-flight checks unnecessarily.
    parse_arguments "$@"

    # Phase 3: Cancel any pending watchdog — must happen before any service state change,
    # whether we're on the standalone-toggle path or the full-install path.
    if [ "$SKIP_SYSTEMD" = false ]; then
        cancel_reboot_watchdog
    fi

    # Phase 4: Handle standalone toggles (no reinstall needed).
    # If -p or -b was also passed (INSTALL_REQUESTED=true), skip early-exit and let
    # setup_systemd_service() apply the overrides inline during the installation instead.
    if [[ -n "$SET_REBOOT" && "$INSTALL_REQUESTED" == "false" ]]; then
        handle_reboot_setting
        exit 0
    fi

    if [[ -n "$SET_AUTO_RESTART" && "$INSTALL_REQUESTED" == "false" ]]; then
        handle_auto_restart_setting
        exit 0
    fi

    # Phase 5: Pre-flight checks (only reached for actual installations)
    check_existing_service

    # Phase 6: System preparation
    apt-get update || error_exit "Failed to update package lists"

    # Phase 7: Stop existing service if needed
    if [ "$SKIP_SYSTEMD" = false ]; then
        stop_disable_and_save_service_state
    fi

    # Phase 8: Install binary (package or source)
    install_binary
    [[ -x /usr/bin/AIS-catcher ]] || error_exit "AIS-catcher binary not found after installation"

    # Phase 9: Post-install configuration
    configure_installation

    # Phase 10: Setup and start service
    if [ "$SKIP_SYSTEMD" = false ]; then
        setup_systemd_service

        if ! restore_service_state; then
            log "ERROR" "Failed to restore service state"
            exit 1
        fi
    else
        log "INFO" "Skipping systemd service installation (--no-systemd flag specified)"
    fi

    # Phase 11: Installation complete
    log "INFO" "AIS-catcher installation completed successfully"
    display_installation_summary
}

# Display installation summary with useful information
display_installation_summary() {
    local divider="═══════════════════════════════════════════════════════════════════════════════"
    
    echo -e "\n${COLOR_BOLD_BLUE}${divider}${COLOR_RESET}"
    echo -e "${COLOR_BOLD_GREEN}                    AIS-CATCHER INSTALLATION SUMMARY${COLOR_RESET}"
    echo -e "${COLOR_BOLD_BLUE}${divider}${COLOR_RESET}\n"
    echo -e "${COLOR_BOLD_GREEN}✓ Installation completed successfully!${COLOR_RESET}\n"
    
    echo -e "${COLOR_BOLD}INSTALLED FILES:${COLOR_RESET}"
    echo "  • Executable:         /usr/bin/AIS-catcher"
    echo "  • Configuration:      ${CONFIG_DIR}/"
    echo "    - Main config:      ${CONFIG_FILE}"
    echo "    - Command options:  ${CMD_FILE}"
    echo "    - Plugins:          ${PLUGIN_DIR}/"
    echo "    - Web assets:       ${CONFIG_DIR}/webassets/"
    echo "    - Database schema:  ${DBMS_DIR}/create.sql"
    echo "    - Documentation:    ${README_DIR}/README.md"
    echo "    - License:          ${CONFIG_DIR}/LICENSE"
    if [ "$SKIP_SYSTEMD" = false ]; then
        echo "  • Service file:       ${SERVICE_FILE}"
        echo "  • Reboot watchdog:    ${REBOOT_SERVICE_FILE}"
        echo "  • Reboot script:      ${LIB_DIR}/wait-and-reboot.sh"
    fi
    echo "  • Log file:           ${LOG_FILE}"
    echo ""
    if [ "$SKIP_SYSTEMD" = false ]; then
        echo -e "${COLOR_BOLD}REBOOT WATCHDOG:${COLOR_RESET}"
        echo "  • Off by default — service restarts indefinitely without rebooting"
        echo "  • When armed: reboots after N rapid restarts within a time window"
        echo "  • Useful for recovering from unresolvable USB device failures"
        echo "  • Re-run this script to toggle (no reinstall needed):"
        echo -e "    ${COLOR_YELLOW}sudo $SCRIPT_NAME --set-reboot-on-failure${COLOR_RESET}           # 3 restarts in 30 min"
        echo -e "    ${COLOR_YELLOW}sudo $SCRIPT_NAME --set-reboot-on-failure 5 3600${COLOR_RESET}   # 5 restarts in 60 min"
        echo -e "    ${COLOR_YELLOW}sudo $SCRIPT_NAME --unset-reboot-on-failure${COLOR_RESET}        # disable"
        echo "  • Settings stored directly in the service file: ${SERVICE_FILE}"
        echo ""
        echo -e "${COLOR_BOLD}AUTO-RESTART:${COLOR_RESET}"
        echo "  • On by default — service restarts automatically after any crash (Restart=always)"
        echo "  • Disable to prevent restart loops during troubleshooting"
        echo "  • Re-run this script to toggle (no reinstall needed):"
        echo -e "    ${COLOR_YELLOW}sudo $SCRIPT_NAME --set-auto-restart${COLOR_RESET}                # enable Restart=always"
        echo -e "    ${COLOR_YELLOW}sudo $SCRIPT_NAME --unset-auto-restart${COLOR_RESET}              # disable Restart=no"
        echo "  • If service locks up after failures, recover with:"
        echo -e "    ${COLOR_YELLOW}systemctl reset-failed ${SERVICE_NAME} && systemctl start ${SERVICE_NAME}${COLOR_RESET}"
        echo "  • To abort a pending reboot (during the 5-minute wait):"
        echo -e "    ${COLOR_YELLOW}shutdown -c${COLOR_RESET}"
        echo ""
    fi
    
    if [ "$SKIP_SYSTEMD" = false ]; then
        echo -e "${COLOR_BOLD}SYSTEMD SERVICE MANAGEMENT:${COLOR_RESET}"
        echo "  • View service status:"
        echo -e "    ${COLOR_YELLOW}systemctl status ${SERVICE_NAME}${COLOR_RESET}"
        echo ""
        echo "  • Start the service:"
        echo -e "    ${COLOR_YELLOW}systemctl start ${SERVICE_NAME}${COLOR_RESET}"
        echo ""
        echo "  • Stop the service:"
        echo -e "    ${COLOR_YELLOW}systemctl stop ${SERVICE_NAME}${COLOR_RESET}"
        echo ""
        echo "  • Restart the service:"
        echo -e "    ${COLOR_YELLOW}systemctl restart ${SERVICE_NAME}${COLOR_RESET}"
        echo ""
        echo "  • Enable at boot:"
        echo -e "    ${COLOR_YELLOW}systemctl enable ${SERVICE_NAME}${COLOR_RESET}"
        echo ""
        echo "  • Disable at boot:"
        echo -e "    ${COLOR_YELLOW}systemctl disable ${SERVICE_NAME}${COLOR_RESET}"
        echo ""
        echo "  • View service logs:"
        echo -e "    ${COLOR_YELLOW}journalctl -u ${SERVICE_NAME} -f${COLOR_RESET}"
        echo ""
    fi
    
    echo -e "${COLOR_BOLD}CONFIGURATION:${COLOR_RESET}"
    echo "  • Edit main configuration:"
    echo -e "    ${COLOR_YELLOW}nano ${CONFIG_FILE}${COLOR_RESET}"
    echo ""
    echo "  • Edit command-line options:"
    echo -e "    ${COLOR_YELLOW}nano ${CMD_FILE}${COLOR_RESET}"
    echo ""
    echo "  • After editing, restart the service:"
    echo -e "    ${COLOR_YELLOW}systemctl restart ${SERVICE_NAME}${COLOR_RESET}"
    echo ""
    
    echo -e "${COLOR_BOLD}WEB INTERFACE:${COLOR_RESET}"
    echo "  • Once running, access the web interface at:"
    echo -e "    ${COLOR_CYAN}http://localhost:8100${COLOR_RESET}"
    echo "    (or use your server's IP address)"
    echo ""
    echo "  • Enable the web server by editing ${CONFIG_FILE}"
    echo -e "    and setting ${COLOR_YELLOW}\"active\": true${COLOR_RESET} in the \"server\" section"
    echo ""
    
    echo -e "${COLOR_BOLD}USEFUL COMMANDS:${COLOR_RESET}"
    echo "  • Run AIS-catcher directly (for testing):"
    echo -e "    ${COLOR_YELLOW}AIS-catcher -h${COLOR_RESET}"
    echo ""
    echo "  • View this installation log:"
    echo -e "    ${COLOR_YELLOW}cat ${LOG_FILE}${COLOR_RESET}"
    echo ""
    
    echo -e "${COLOR_BOLD}DOCUMENTATION:${COLOR_RESET}"
    echo -e "  • Online docs:  ${COLOR_CYAN}https://docs.aiscatcher.org/${COLOR_RESET}"
    echo -e "  • GitHub repo:  ${COLOR_CYAN}https://github.com/jvde-github/AIS-catcher${COLOR_RESET}"
    echo "  • Local README: ${README_DIR}/README.md"
    echo ""
    
    echo -e "${COLOR_BOLD_BLUE}${divider}${COLOR_RESET}\n"
}

# Run the script
main "$@"