#!/usr/bin/env bash

################################################################################
# Remote Homelab Collection Wrapper
# Purpose: Executes the collection script on a remote Proxmox host via SSH
#          and retrieves the results back to your local machine (WSL/Linux)
#
# Usage: ./collect-remote.sh [PROXMOX_HOST] [OPTIONS]
################################################################################

set -euo pipefail

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

# Script configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
COLLECTION_SCRIPT="${SCRIPT_DIR}/collect-homelab-config.sh"
REMOTE_SCRIPT_PATH="/tmp/collect-homelab-config.sh"
LOCAL_OUTPUT_DIR="${SCRIPT_DIR}"

# SSH configuration
SSH_USER="${SSH_USER:-root}"
SSH_PORT="${SSH_PORT:-22}"
SSH_OPTS="-o ConnectTimeout=10 -o StrictHostKeyChecking=no"

################################################################################
# Functions
################################################################################

log() {
    local level="$1"
    shift
    local message="$*"

    case "${level}" in
        INFO)
            echo -e "${BLUE}[INFO]${NC} ${message}"
            ;;
        SUCCESS)
            echo -e "${GREEN}[✓]${NC} ${message}"
            ;;
        WARN)
            echo -e "${YELLOW}[WARN]${NC} ${message}"
            ;;
        ERROR)
            echo -e "${RED}[ERROR]${NC} ${message}" >&2
            ;;
    esac
}

banner() {
    echo ""
    echo -e "${BOLD}${CYAN}======================================================================${NC}"
    echo -e "${BOLD}${CYAN}  $1${NC}"
    echo -e "${BOLD}${CYAN}======================================================================${NC}"
    echo ""
}

usage() {
    cat <<EOF
${BOLD}Remote Homelab Collection Wrapper${NC}

${BOLD}USAGE:${NC}
    $0 PROXMOX_HOST [OPTIONS]

${BOLD}DESCRIPTION:${NC}
    Executes the homelab collection script on a remote Proxmox host via SSH,
    then retrieves the results back to your local machine.

${BOLD}ARGUMENTS:${NC}
    PROXMOX_HOST        IP address or hostname of your Proxmox server

${BOLD}OPTIONS:${NC}
    -u, --user USER     SSH username (default: root)
    -p, --port PORT     SSH port (default: 22)
    -l, --level LEVEL   Collection level: basic, standard, full, paranoid
                        (default: standard)
    -s, --sanitize OPT  Sanitization: all, ips, none (default: passwords/tokens only)
    -o, --output DIR    Local directory to store results (default: current directory)
    -k, --keep-remote   Keep the export on the remote host (default: remove after download)
    -v, --verbose       Verbose output
    -h, --help          Show this help message

${BOLD}ENVIRONMENT VARIABLES:${NC}
    SSH_USER            Default SSH username (default: root)
    SSH_PORT            Default SSH port (default: 22)

${BOLD}EXAMPLES:${NC}
    # Basic usage
    $0 192.168.1.100

    # Full collection with complete sanitization
    $0 192.168.1.100 --level full --sanitize all

    # Custom SSH user and port
    $0 proxmox.local --user admin --port 2222

    # Keep results on remote host
    $0 192.168.1.100 --keep-remote

    # Verbose output with custom output directory
    $0 192.168.1.100 -v -o ~/backups/homelab

${BOLD}PREREQUISITES:${NC}
    1. SSH access to the Proxmox host
    2. collect-homelab-config.sh in the same directory as this script
    3. Sufficient disk space on both remote and local machines

${BOLD}WORKFLOW:${NC}
    1. Copies collection script to remote Proxmox host
    2. Executes the script remotely
    3. Downloads the compressed archive to local machine
    4. Optionally removes the remote copy
    5. Extracts the archive locally

${BOLD}NOTES:${NC}
    - Requires passwordless SSH or SSH key authentication (recommended)
    - The script will be run as the specified SSH user (typically root)
    - Remote execution output is displayed in real-time

EOF
}

check_prerequisites() {
    # Check if collection script exists
    if [[ ! -f "${COLLECTION_SCRIPT}" ]]; then
        log ERROR "Collection script not found: ${COLLECTION_SCRIPT}"
        log ERROR "Ensure collect-homelab-config.sh is in the same directory as this script"
        exit 1
    fi

    # Check if ssh is available
    if ! command -v ssh &> /dev/null; then
        log ERROR "SSH client not found. Please install openssh-client"
        exit 1
    fi

    # Check if scp is available
    if ! command -v scp &> /dev/null; then
        log ERROR "SCP not found. Please install openssh-client"
        exit 1
    fi
}

test_ssh_connection() {
    local host="$1"

    log INFO "Testing SSH connection to ${SSH_USER}@${host}:${SSH_PORT}..."

    if ssh ${SSH_OPTS} -p "${SSH_PORT}" "${SSH_USER}@${host}" "exit 0" 2>/dev/null; then
        log SUCCESS "SSH connection successful"
        return 0
    else
        log ERROR "Cannot connect to ${SSH_USER}@${host}:${SSH_PORT}"
        log ERROR "Possible issues:"
        log ERROR "  - Host is unreachable"
        log ERROR "  - SSH service is not running"
        log ERROR "  - Incorrect credentials"
        log ERROR "  - Firewall blocking connection"
        log ERROR ""
        log ERROR "Try manually: ssh -p ${SSH_PORT} ${SSH_USER}@${host}"
        return 1
    fi
}

verify_proxmox_host() {
    local host="$1"

    log INFO "Verifying Proxmox installation on remote host..."

    if ssh ${SSH_OPTS} -p "${SSH_PORT}" "${SSH_USER}@${host}" "test -f /etc/pve/.version" 2>/dev/null; then
        local pve_version=$(ssh ${SSH_OPTS} -p "${SSH_PORT}" "${SSH_USER}@${host}" "cat /etc/pve/.version" 2>/dev/null)
        log SUCCESS "Confirmed Proxmox VE installation (version: ${pve_version})"
        return 0
    else
        log WARN "Remote host does not appear to be a Proxmox VE server"
        log WARN "Proceeding anyway, but collection may fail..."
        return 0
    fi
}

upload_script() {
    local host="$1"

    banner "Uploading Collection Script"

    log INFO "Copying collection script to ${host}..."

    if scp ${SSH_OPTS} -P "${SSH_PORT}" "${COLLECTION_SCRIPT}" "${SSH_USER}@${host}:${REMOTE_SCRIPT_PATH}"; then
        log SUCCESS "Script uploaded successfully"

        # Make executable
        ssh ${SSH_OPTS} -p "${SSH_PORT}" "${SSH_USER}@${host}" "chmod +x ${REMOTE_SCRIPT_PATH}"
        log SUCCESS "Script permissions set"
        return 0
    else
        log ERROR "Failed to upload script"
        return 1
    fi
}

execute_remote_collection() {
    local host="$1"
    shift
    local collection_args=("$@")

    banner "Executing Collection on Remote Host"

    log INFO "Running collection script on ${host}..."
    log INFO "Arguments: ${collection_args[*]}"

    # Build the remote command
    local remote_cmd="${REMOTE_SCRIPT_PATH} ${collection_args[*]}"

    # Execute remotely and stream output
    if ssh ${SSH_OPTS} -p "${SSH_PORT}" "${SSH_USER}@${host}" "${remote_cmd}"; then
        log SUCCESS "Collection completed successfully on remote host"
        return 0
    else
        log ERROR "Collection failed on remote host"
        return 1
    fi
}

download_results() {
    local host="$1"
    local output_dir="$2"

    banner "Downloading Results"

    log INFO "Finding remote export archive..."

    # Find the most recent export archive
    local remote_archive=$(ssh ${SSH_OPTS} -p "${SSH_PORT}" "${SSH_USER}@${host}" \
        "ls -t /root/homelab-export-*.tar.gz 2>/dev/null | head -1" 2>/dev/null)

    if [[ -z "${remote_archive}" ]]; then
        log ERROR "No export archive found on remote host"
        log ERROR "Collection may have failed or compression was disabled"
        return 1
    fi

    log INFO "Found archive: ${remote_archive}"

    # Create output directory
    mkdir -p "${output_dir}"

    # Download the archive
    local local_archive="${output_dir}/$(basename "${remote_archive}")"

    log INFO "Downloading to: ${local_archive}"

    if scp ${SSH_OPTS} -P "${SSH_PORT}" "${SSH_USER}@${host}:${remote_archive}" "${local_archive}"; then
        log SUCCESS "Archive downloaded successfully"

        # Extract the archive
        log INFO "Extracting archive..."
        if tar -xzf "${local_archive}" -C "${output_dir}"; then
            log SUCCESS "Archive extracted to: ${output_dir}/$(basename "${local_archive}" .tar.gz)"

            # Show summary
            local extracted_dir="${output_dir}/$(basename "${local_archive}" .tar.gz)"
            if [[ -f "${extracted_dir}/SUMMARY.md" ]]; then
                echo ""
                log INFO "Collection Summary:"
                echo ""
                head -30 "${extracted_dir}/SUMMARY.md"
                echo ""
                log INFO "Full summary: ${extracted_dir}/SUMMARY.md"
            fi

            return 0
        else
            log ERROR "Failed to extract archive"
            return 1
        fi
    else
        log ERROR "Failed to download archive"
        return 1
    fi
}

cleanup_remote() {
    local host="$1"
    local keep_remote="$2"

    if [[ "${keep_remote}" == "true" ]]; then
        log INFO "Keeping export on remote host (--keep-remote specified)"
        return 0
    fi

    banner "Cleaning Up Remote Host"

    log INFO "Removing export files from remote host..."

    # Remove the script
    ssh ${SSH_OPTS} -p "${SSH_PORT}" "${SSH_USER}@${host}" "rm -f ${REMOTE_SCRIPT_PATH}" 2>/dev/null || true

    # Remove export directories and archives
    ssh ${SSH_OPTS} -p "${SSH_PORT}" "${SSH_USER}@${host}" \
        "rm -rf /root/homelab-export-* 2>/dev/null" 2>/dev/null || true

    log SUCCESS "Remote cleanup completed"
}

################################################################################
# Main Execution
################################################################################

main() {
    # Parse arguments
    if [[ $# -eq 0 ]]; then
        usage
        exit 1
    fi

    local proxmox_host=""
    local collection_level="standard"
    local sanitize_option=""
    local keep_remote="false"
    local verbose="false"

    # First argument is the host
    proxmox_host="$1"
    shift

    # Parse remaining options
    local collection_args=()

    while [[ $# -gt 0 ]]; do
        case "$1" in
            -u|--user)
                SSH_USER="$2"
                shift 2
                ;;
            -p|--port)
                SSH_PORT="$2"
                shift 2
                ;;
            -l|--level)
                collection_level="$2"
                collection_args+=("--level" "$2")
                shift 2
                ;;
            -s|--sanitize)
                sanitize_option="$2"
                collection_args+=("--sanitize" "$2")
                shift 2
                ;;
            -o|--output)
                LOCAL_OUTPUT_DIR="$2"
                shift 2
                ;;
            -k|--keep-remote)
                keep_remote="true"
                shift
                ;;
            -v|--verbose)
                verbose="true"
                collection_args+=("--verbose")
                shift
                ;;
            -h|--help)
                usage
                exit 0
                ;;
            *)
                log ERROR "Unknown option: $1"
                usage
                exit 1
                ;;
        esac
    done

    # Validate host
    if [[ -z "${proxmox_host}" ]]; then
        log ERROR "Proxmox host not specified"
        usage
        exit 1
    fi

    # Display configuration
    banner "Remote Homelab Collection"
    echo -e "${BOLD}Target Host:${NC}     ${proxmox_host}"
    echo -e "${BOLD}SSH User:${NC}        ${SSH_USER}"
    echo -e "${BOLD}SSH Port:${NC}        ${SSH_PORT}"
    echo -e "${BOLD}Collection Level:${NC} ${collection_level}"
    echo -e "${BOLD}Output Directory:${NC} ${LOCAL_OUTPUT_DIR}"
    echo -e "${BOLD}Keep Remote:${NC}     ${keep_remote}"
    echo ""

    # Execute workflow
    check_prerequisites
    test_ssh_connection "${proxmox_host}" || exit 1
    verify_proxmox_host "${proxmox_host}"
    upload_script "${proxmox_host}" || exit 1
    execute_remote_collection "${proxmox_host}" "${collection_args[@]}" || exit 1
    download_results "${proxmox_host}" "${LOCAL_OUTPUT_DIR}" || exit 1
    cleanup_remote "${proxmox_host}" "${keep_remote}"

    banner "Collection Complete"

    log SUCCESS "Homelab infrastructure export completed successfully"
    log INFO "Results are available in: ${LOCAL_OUTPUT_DIR}"
    echo ""
}

# Run main function
main "$@"