#!/usr/bin/env bash
set -euo pipefail

# Color codes for output formatting
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
DIM='\033[2m'
NC='\033[0m' # No Color

# Icons
ICON_SUCCESS="✓"
ICON_ERROR="✗"
ICON_INFO="ℹ"
ICON_WARNING="⚠"
ICON_ARROW="→"
ICON_DOWNLOAD="⬇"
ICON_UPLOAD="⬆"
ICON_INSTALL="📦"
ICON_UPDATE="🔄"
ICON_DEPLOY="🚀"
ICON_TIMER="⏰"
ICON_CHECK="🔍"
ICON_CONFIG="⚙"
ICON_NETWORK="🌐"
ICON_VERSION="📋"
ICON_LOG="📝"
ICON_HOST="🖥"

SSH_HOST=""
HOSTS_FILE=""
INSTALL_TIMER=false
TIMER_TIME="daily"
DEFAULT_PASSWORD=""
REMOTE_ARGS=()
SSH_TIMEOUT=30

# Arrays to track deployment results
declare -a SUCCESS_HOSTS=()
declare -a FAILED_HOSTS=()
declare -a HOST_ACTIONS=()

# Helper functions for formatted output
print_header() {
    echo -e "\n${BLUE}${BOLD}╔═══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${BLUE}${BOLD}║  $1${NC}"
    echo -e "${BLUE}${BOLD}╚═══════════════════════════════════════════════════════════════╝${NC}\n"
}

print_info() {
    echo -e "  ${BLUE}${ICON_INFO}${NC}  $1"
}

print_success() {
    echo -e "  ${GREEN}${ICON_SUCCESS}${NC}  $1"
}

print_error() {
    echo -e "  ${RED}${ICON_ERROR}${NC}  $1"
}

print_warning() {
    echo -e "  ${YELLOW}${ICON_WARNING}${NC}  $1"
}

print_status() {
    echo -e "  ${CYAN}${ICON_ARROW}${NC}  $1"
}

print_step() {
    echo -e "  ${DIM}•${NC} $1"
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --ssh)
            SSH_HOST="$2"
            shift 2
            ;;
        --hosts-file|--file)
            HOSTS_FILE="$2"
            shift 2
            ;;
        --timer)
            INSTALL_TIMER=true
            REMOTE_ARGS+=("--timer")
            shift
            ;;
        --time)
            TIMER_TIME="$2"
            REMOTE_ARGS+=("--time" "$2")
            shift 2
            ;;
        --password|-p)
            DEFAULT_PASSWORD="$2"
            shift 2
            ;;
        *)
            echo "Unknown option: $1"
            exit 1
            ;;
    esac
done

deploy_to_host() {
    local host="$1"
    local sudo_pass="${2:-}"
    echo -e "${CYAN}${BOLD}┌─────────────────────────────────────────────────────────────┐${NC}"
    echo -e "${CYAN}${BOLD}│${NC} ${ICON_DEPLOY} ${YELLOW}${host}${NC}"
    echo -e "${CYAN}${BOLD}└─────────────────────────────────────────────────────────────┘${NC}"
    
    REMOTE_SCRIPT="/tmp/.netbird-install-update-$$"
    REMOTE_ACTION_FILE="/tmp/.netbird-action-$$"
    LOCAL_ACTION_FILE="/tmp/.netbird-deploy-action-$$"
    
    print_step "Copying script..."
    if ! scp -o ConnectTimeout=$SSH_TIMEOUT -o BatchMode=yes "$0" "$host:$REMOTE_SCRIPT" 2>&1; then
        print_error "Failed to copy script to $host (timeout or connection error)"
        return 1
    fi
    
    print_step "Executing script..."
    # Build SSH command based on whether we have a password
    if [ -n "$sudo_pass" ]; then
        # Write password to remote temp file using printf to handle special characters
        # Use base64 to safely transfer the password
        local encoded_pass=$(echo -n "$sudo_pass" | base64)
        if ! ssh -o ConnectTimeout=$SSH_TIMEOUT -o StrictHostKeyChecking=no "$host" "echo '$encoded_pass' | base64 -d > /tmp/.netbird-pass-$$ && chmod 600 /tmp/.netbird-pass-$$" 2>&1; then
            print_error "Failed to setup password file on $host"
            return 1
        fi
        
        # Execute script with sudo, reading password from file
        if ! ssh -o ConnectTimeout=$SSH_TIMEOUT -o StrictHostKeyChecking=no "$host" "sudo -S bash $REMOTE_SCRIPT ${REMOTE_ARGS[*]} < /tmp/.netbird-pass-$$" 2>&1; then
            print_error "Failed to execute script on $host"
            ssh -o ConnectTimeout=$SSH_TIMEOUT -o StrictHostKeyChecking=no "$host" "rm -f $REMOTE_SCRIPT $REMOTE_ACTION_FILE /tmp/.netbird-pass-$$" 2>/dev/null || true
            return 1
        fi
        
        # Clean up password file
        ssh -o ConnectTimeout=$SSH_TIMEOUT -o StrictHostKeyChecking=no "$host" "rm -f /tmp/.netbird-pass-$$" 2>/dev/null || true
    else
        # Interactive sudo with terminal
        if ! ssh -o ConnectTimeout=$SSH_TIMEOUT -o BatchMode=yes -t "$host" "sudo bash $REMOTE_SCRIPT ${REMOTE_ARGS[*]}" 2>&1; then
            print_error "Failed to execute script on $host"
            ssh -o ConnectTimeout=$SSH_TIMEOUT -o BatchMode=yes "$host" "rm -f $REMOTE_SCRIPT $REMOTE_ACTION_FILE" 2>/dev/null || true
            return 1
        fi
    fi
    
    # Read the action from the remote action file and write to local action file
    if action_output=$(ssh -o ConnectTimeout=$SSH_TIMEOUT -o StrictHostKeyChecking=no "$host" "cat $REMOTE_ACTION_FILE 2>/dev/null" < /dev/null 2>&1); then
        echo "$action_output" > "$LOCAL_ACTION_FILE"
    else
        echo "unknown" > "$LOCAL_ACTION_FILE"
    fi
    
    # Clean up remote files
    ssh -o ConnectTimeout=$SSH_TIMEOUT -o StrictHostKeyChecking=no "$host" "rm -f $REMOTE_SCRIPT $REMOTE_ACTION_FILE" < /dev/null 2>/dev/null || true
    print_success "Completed: $host"
    return 0
}

if [ -n "$HOSTS_FILE" ]; then
    if [ ! -f "$HOSTS_FILE" ]; then
        print_error "Hosts file not found: $HOSTS_FILE"
        exit 1
    fi
    
    print_header "${ICON_NETWORK} Multi-Host Deployment"
    print_info "Reading hosts from: ${YELLOW}$HOSTS_FILE${NC}"
    print_info "SSH timeout: ${YELLOW}${SSH_TIMEOUT}s${NC}"
    echo ""
    
    SUCCESS_COUNT=0
    FAIL_COUNT=0
    ACTION_FILE="/tmp/.netbird-deploy-action-$$"
    
    # Use file descriptor 3 to read hosts file (prevents SSH from consuming stdin)
    exec 3< "$HOSTS_FILE"
    while IFS= read -r line <&3 || [ -n "$line" ]; do
        # Skip empty lines and comments
        [[ -z "$line" || "$line" =~ ^[[:space:]]*# ]] && continue
        # Trim whitespace
        line=$(echo "$line" | xargs)
        
        # Parse hostname and optional password
        host=$(echo "$line" | awk '{print $1}')
        sudo_pass=$(echo "$line" | awk '{print $2}')
        
        # Use default password if no host-specific password is provided
        if [ -z "$sudo_pass" ] && [ -n "$DEFAULT_PASSWORD" ]; then
            sudo_pass="$DEFAULT_PASSWORD"
        fi
        
        # Run deployment (output streams to terminal)
        if deploy_to_host "$host" "$sudo_pass"; then
            # Read the action from the temp file
            action="unknown"
            if [ -f "$ACTION_FILE" ]; then
                action=$(cat "$ACTION_FILE")
                rm -f "$ACTION_FILE"
            fi
            SUCCESS_HOSTS+=("$host")
            HOST_ACTIONS+=("$action")
            ((SUCCESS_COUNT++)) || true
        else
            FAILED_HOSTS+=("$host")
            ((FAIL_COUNT++)) || true
        fi
        echo ""
    done
    exec 3<&-
    
    # Print detailed summary
    echo -e "\n${GREEN}${BOLD}╔═══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${GREEN}${BOLD}║${NC} ${ICON_SUCCESS} ${GREEN}${BOLD}Deployment Summary${NC}"
    echo -e "${GREEN}${BOLD}╠═══════════════════════════════════════════════════════════════╣${NC}"
    echo -e "${GREEN}${BOLD}║${NC}  ${GREEN}${ICON_SUCCESS}${NC} Successful: ${BOLD}$SUCCESS_COUNT${NC}"
    echo -e "${GREEN}${BOLD}║${NC}  ${RED}${ICON_ERROR}${NC} Failed:     ${BOLD}$FAIL_COUNT${NC}"
    echo -e "${GREEN}${BOLD}╚═══════════════════════════════════════════════════════════════╝${NC}"
    
    # Show successful hosts with actions
    if [ ${#SUCCESS_HOSTS[@]} -gt 0 ]; then
        echo -e "\n${GREEN}${BOLD}Successful Hosts:${NC}"
        for i in "${!SUCCESS_HOSTS[@]}"; do
            host="${SUCCESS_HOSTS[$i]}"
            action="${HOST_ACTIONS[$i]}"
            echo -e "  ${GREEN}${ICON_SUCCESS}${NC} ${YELLOW}$host${NC} ${DIM}($action)${NC}"
        done
    fi
    
    # Show failed hosts
    if [ ${#FAILED_HOSTS[@]} -gt 0 ]; then
        echo -e "\n${RED}${BOLD}Failed Hosts:${NC}"
        for host in "${FAILED_HOSTS[@]}"; do
            echo -e "  ${RED}${ICON_ERROR}${NC} ${YELLOW}$host${NC}"
        done
    fi
    
    echo ""
    exit 0
fi

if [ -n "$SSH_HOST" ]; then
    print_header "${ICON_DEPLOY} Single Host Deployment"
    ACTION_FILE="/tmp/.netbird-deploy-action-$$"
    if deploy_to_host "$SSH_HOST" "$DEFAULT_PASSWORD"; then
        # Read and display the action
        if [ -f "$ACTION_FILE" ]; then
            action=$(cat "$ACTION_FILE")
            rm -f "$ACTION_FILE"
            print_info "Action performed: ${YELLOW}$action${NC}"
        fi
    fi
    echo ""
    exit 0
fi

if [ "$EUID" -ne 0 ]; then
    print_warning "Elevating to root..."
    exec sudo "$0" "$@"
fi

WORK_DIR=$(mktemp -d)
trap 'rm -rf "$WORK_DIR"' EXIT

detect_install() {
    if ! command -v netbird &>/dev/null; then
        echo "none"
        return
    fi

    if [ -f /etc/netbird/install.conf ]; then
        local pm
        pm=$(grep -oP 'package_manager=\K.*' /etc/netbird/install.conf 2>/dev/null || true)
        if [ "$pm" = "bin" ]; then
            echo "bin"
            return
        elif [ -n "$pm" ]; then
            echo "pkg:$pm"
            return
        fi
    fi

    if dpkg -s netbird &>/dev/null; then
        echo "pkg:apt"
        return
    fi
    if rpm -q netbird &>/dev/null; then
        if command -v dnf &>/dev/null; then
            echo "pkg:dnf"
        else
            echo "pkg:yum"
        fi
        return
    fi
    if pacman -Q netbird &>/dev/null; then
        echo "pkg:pacman"
        return
    fi

    echo "bin"
}

install_timer() {
    print_header "${ICON_TIMER} Installing Auto-Update Timer"
    
    if systemctl is-active --quiet netbird-update.timer; then
        print_warning "Updating existing timer..."
        systemctl stop netbird-update.timer
        systemctl disable netbird-update.timer
    fi

    cat > /etc/systemd/system/netbird-update.service <<'UNIT'
[Unit]
Description=Netbird Auto-Update

[Service]
Type=oneshot
ExecStart=/bin/bash -c 'curl -fsSL https://pkgs.netbird.io/install.sh | UPDATE_NETBIRD=true sh && chown root:root /usr/bin/netbird 2>/dev/null || true && chmod +x /usr/bin/netbird 2>/dev/null || true && restorecon -v /usr/bin/netbird 2>/dev/null || true && systemctl restart netbird 2>/dev/null || true'
StandardOutput=append:/var/log/netbird-update.log
StandardError=append:/var/log/netbird-update.log
UNIT

    cat > /etc/systemd/system/netbird-update.timer <<UNIT
[Unit]
Description=Run Netbird update daily

[Timer]
OnCalendar=$TIMER_TIME
Persistent=true
RandomizedDelaySec=600

[Install]
WantedBy=timers.target
UNIT

    systemctl daemon-reload
    systemctl enable --now netbird-update.timer
    
    print_success "Timer installed successfully"
    print_info "Schedule: ${YELLOW}$TIMER_TIME${NC} (with random 0-10 minute delay)"
    print_info "Logs: ${YELLOW}/var/log/netbird-update.log${NC}"
}

print_header "${ICON_INSTALL} NetBird Install/Update"

print_status "${ICON_DOWNLOAD} Downloading Netbird install script..."
curl -fsSL -o "$WORK_DIR/install.sh" https://pkgs.netbird.io/install.sh
chmod +x "$WORK_DIR/install.sh"
print_success "Download complete"

print_status "${ICON_CHECK} Detecting installation state..."
state=$(detect_install)
print_info "Detected: ${YELLOW}$state${NC}"
echo ""

case "$state" in
    none)
        print_header "${ICON_INSTALL} Fresh Installation"
        print_info "Netbird not found. Installing..."
        USE_BIN_INSTALL=true SKIP_UI_APP=true "$WORK_DIR/install.sh"
        chown root:root /usr/bin/netbird 2>/dev/null || true
        restorecon -v /usr/bin/netbird 2>/dev/null || true
        systemctl restart netbird 2>/dev/null || true
        print_success "Netbird installed successfully"
        echo "fresh install" > /tmp/.netbird-action-$$
        ;;
    pkg:*)
        print_header "${ICON_UPDATE} Package Manager Replacement"
        pm="${state#pkg:}"
        print_info "Netbird installed via ${YELLOW}$pm${NC}"
        print_info "Replacing with binary install..."
        cat > /tmp/.netbird_replace.sh <<'EOF'
#!/usr/bin/env bash
set -euo pipefail
sleep 3
systemctl stop netbird 2>/dev/null || true
case "$1" in
    apt)        apt-get remove -y netbird netbird-ui 2>/dev/null || true ;;
    yum)        yum remove -y netbird netbird-ui 2>/dev/null || true ;;
    dnf)        dnf remove -y netbird netbird-ui 2>/dev/null || true ;;
    zypper)     zypper remove -y netbird netbird-ui 2>/dev/null || true ;;
    rpm-ostree) rpm-ostree uninstall -y netbird netbird-ui 2>/dev/null || true ;;
    pacman)     pacman -Rns --noconfirm netbird netbird-ui 2>/dev/null || true ;;
esac
USE_BIN_INSTALL=true SKIP_UI_APP=true bash /tmp/.netbird_install.sh
chown root:root /usr/bin/netbird 2>/dev/null || true
chmod +x /usr/bin/netbird 2>/dev/null || true
restorecon -v /usr/bin/netbird 2>/dev/null || true
systemctl restart netbird 2>/dev/null || true
rm -f /tmp/.netbird_replace.sh /tmp/.netbird_install.sh
EOF
        chmod +x /tmp/.netbird_replace.sh
        cp "$WORK_DIR/install.sh" /tmp/.netbird_install.sh
        setsid bash /tmp/.netbird_replace.sh "$pm" </dev/null >/var/log/netbird-update.log 2>&1 &
        print_warning "Replacement scheduled in background"
        print_info "Connection may drop momentarily"
        print_info "Check ${YELLOW}/var/log/netbird-update.log${NC} for progress"
        echo "pkg replacement" > /tmp/.netbird-action-$$
        ;;
    bin)
        print_header "${ICON_VERSION} Binary Update Check"
        if [ ! -f /etc/netbird/install.conf ]; then
            mkdir -p /etc/netbird
            echo "package_manager=bin" > /etc/netbird/install.conf
        fi
        current=$(netbird version 2>/dev/null || echo "unknown")
        latest=$(curl -fsSL https://pkgs.netbird.io/releases/latest 2>/dev/null | grep -oP '"tag_name":\s*"v?\K[^"]+' || echo "unknown")

        print_info "Current version: ${YELLOW}$current${NC}"
        print_info "Latest version:  ${YELLOW}$latest${NC}"
        echo ""

        if [ "$current" = "$latest" ] && [ "$current" != "unknown" ]; then
            print_success "Netbird is up to date"
            echo "up to date" > /tmp/.netbird-action-$$
        else
            print_info "Updating Netbird: ${YELLOW}$current${NC} → ${GREEN}$latest${NC}"
            cp "$WORK_DIR/install.sh" /tmp/.netbird_update.sh
            setsid bash -c 'sleep 3 && bash /tmp/.netbird_update.sh --update && chown root:root /usr/bin/netbird 2>/dev/null || true && chmod +x /usr/bin/netbird 2>/dev/null || true && restorecon -v /usr/bin/netbird 2>/dev/null || true && systemctl restart netbird 2>/dev/null || true && rm -f /tmp/.netbird_update.sh' \
                </dev/null >/var/log/netbird-update.log 2>&1 &
            print_warning "Update scheduled in background"
            print_info "Connection may drop momentarily"
            print_info "Check ${YELLOW}/var/log/netbird-update.log${NC} for progress"
            echo "updated" > /tmp/.netbird-action-$$
        fi
        ;;
esac

if $INSTALL_TIMER; then
    install_timer
fi

echo ""