#!/usr/bin/env bash
# test-vm.sh — Build, start, and diagnose the RadiantWave NixOS test VM.
#
# On macOS (ARM): builds and runs the VM on a remote Linux server.
# On Linux (x86_64): builds and runs locally.
#
# Usage:
#   ./scripts/test-vm.sh              Build + start + run diagnostics
#   ./scripts/test-vm.sh --build      Build only
#   ./scripts/test-vm.sh --start      Start only (must already be built)
#   ./scripts/test-vm.sh --stop       Stop the running VM
#   ./scripts/test-vm.sh --diagnose   Run diagnostics on a running VM
#   ./scripts/test-vm.sh --shell      Interactive SSH shell into the VM
#
# Environment variables:
#   RADIANTWAVE_REMOTE   Remote server for macOS builds (default: nclack@192.168.1.203)
#   RADIANTWAVE_REMOTE_DIR  Remote repo path (default: ~/radiantwave)

set -euo pipefail

# ── Paths & defaults ──────────────────────────────────────────────────────────
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
REMOTE="${RADIANTWAVE_REMOTE:-nclack@192.168.1.203}"
REMOTE_DIR="${RADIANTWAVE_REMOTE_DIR:-~/radiantwave}"
SSH_PORT=2222
VM_USER="kiosk"
VM_PASS="kiosk"

# ── Colors ────────────────────────────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; BLUE='\033[0;34m'
DIM='\033[2m'; BOLD='\033[1m'; NC='\033[0m'

log()  { echo -e "${BLUE}▸${NC} $*"; }
ok()   { echo -e "${GREEN}✓${NC} $*"; }
warn() { echo -e "${YELLOW}⚠${NC} $*"; }
err()  { echo -e "${RED}✗${NC} $*" >&2; }
hr()   { echo -e "${DIM}────────────────────────────────────────────────${NC}"; }

# ── Platform detection ────────────────────────────────────────────────────────
detect_platform() {
    local os arch
    os="$(uname -s)"
    arch="$(uname -m)"
    case "$os" in
        Darwin) echo "macos" ;;
        Linux)
            case "$arch" in
                x86_64)  echo "linux" ;;
                aarch64) echo "linux-arm" ;;
                *)       echo "linux-unknown" ;;
            esac
            ;;
        *) echo "unknown" ;;
    esac
}

PLATFORM="$(detect_platform)"

# ── Dependency checks ─────────────────────────────────────────────────────────
check_deps() {
    local missing=()

    if ! command -v nix &>/dev/null; then
        missing+=("nix")
    fi

    if [[ "$PLATFORM" == "macos" ]]; then
        for cmd in ssh rsync; do
            command -v "$cmd" &>/dev/null || missing+=("$cmd")
        done
    fi

    if (( ${#missing[@]} > 0 )); then
        err "Missing required commands: ${missing[*]}"
        exit 1
    fi

    # sshpass — try direct, then nix-shell wrapper
    if command -v sshpass &>/dev/null; then
        SSPASS="sshpass"
    elif nix-shell -p sshpass --run 'which sshpass' &>/dev/null; then
        SSPASS="nix-shell -p sshpass --run"
    else
        err "sshpass not found. Install via: brew install sshpass (macOS) or nix profile install nixpkgs#sshpass"
        exit 1
    fi
}

# ── VM path helpers ───────────────────────────────────────────────────────────
# The nix build output is a directory containing bin/run-radiantwave-test-vm.
VM_RESULT_FILE="/tmp/radiantwave-test-vm-path"
VM_BIN_DIR=""

resolve_vm_bin() {
    if [[ -n "$VM_BIN_DIR" ]]; then
        return
    fi

    local cached=""
    if [[ "$PLATFORM" == "linux" ]]; then
        if [[ -f "$VM_RESULT_FILE" ]]; then
            cached="$(cat "$VM_RESULT_FILE")"
            if [[ -x "$cached/bin/run-radiantwave-test-vm" ]]; then
                VM_BIN_DIR="$cached"
                return
            fi
        fi
    else
        cached="$(ssh "$REMOTE" "cat '$VM_RESULT_FILE' 2>/dev/null")" || true
        if [[ -n "$cached" ]]; then
            # Check if it exists on the remote server
            if ssh "$REMOTE" "test -x '$cached/bin/run-radiantwave-test-vm'" 2>/dev/null; then
                VM_BIN_DIR="$cached"
                return
            fi
        fi
    fi

    return 1
}

# ── Platform-specific execution ───────────────────────────────────────────────
# On Linux: run commands locally.
# On macOS: sync repo to remote, then run commands there via SSH.
run_cmd() {
    if [[ "$PLATFORM" == "linux" ]]; then
        "$@"
    else
        # Run on remote server
        ssh "$REMOTE" "cd '$REMOTE_DIR' && $(printf '%q ' "$@")"
    fi
}

run_cmd_capture() {
    if [[ "$PLATFORM" == "linux" ]]; then
        "$@"
    else
        ssh "$REMOTE" "cd '$REMOTE_DIR' && $(printf '%q ' "$@")"
    fi
}

sync_to_remote() {
    if [[ "$PLATFORM" == "macos" ]]; then
        log "Syncing repo to remote server ($REMOTE)..."
        # Ensure remote dir exists
        ssh "$REMOTE" "mkdir -p '$REMOTE_DIR'"
        rsync -avz --delete \
            --exclude='.git/' \
            --exclude='result' \
            --exclude='result-*' \
            --exclude='bin/' \
            --exclude='build/' \
            "$REPO_ROOT/" "$REMOTE:$REMOTE_DIR/"
    fi
}

# ── Commands ──────────────────────────────────────────────────────────────────

cmd_build() {
    log "Building test VM ($PLATFORM)..."
    sync_to_remote

    local out_path
    if [[ "$PLATFORM" == "linux" ]]; then
        out_path="$(nix build ".#packages.x86_64-linux.test-vm" --print-out-paths 2>&1)"
        echo "$out_path" > "$VM_RESULT_FILE"
        VM_BIN_DIR="$out_path"
    else
        # Build on remote, save path there
        out_path="$(ssh "$REMOTE" "cd '$REMOTE_DIR' && nix build .#packages.x86_64-linux.test-vm --print-out-paths 2>&1")"
        ssh "$REMOTE" "echo '$out_path' > $VM_RESULT_FILE"
        VM_BIN_DIR="$out_path"
    fi

    ok "VM built: $VM_BIN_DIR"
}

cmd_start() {
    if ! resolve_vm_bin 2>/dev/null; then
        err "VM not built yet. Run: $0 --build"
        exit 1
    fi

    if vm_is_running; then
        warn "VM already running (PID $(cat /tmp/radiantwave-test-vm.pid 2>/dev/null || echo '?'))"
        return
    fi

    local vm_script="$VM_BIN_DIR/bin/run-radiantwave-test-vm"
    log "Starting VM in background..."

    if [[ "$PLATFORM" == "linux" ]]; then
        # Run locally with serial to a log file
        nohup "$vm_script" \
            -nographic \
            -serial file:/tmp/radiantwave-test-vm-serial.log \
            > /tmp/radiantwave-test-vm.log 2>&1 &
        echo $! > /tmp/radiantwave-test-vm.pid
        ok "VM started (PID $!)"
    else
        # Run on remote server — use nohup so the VM survives SSH disconnect
        ssh "$REMOTE" "cd '$REMOTE_DIR' && \
            nohup $vm_script -nographic \
                -serial file:/tmp/radiantwave-test-vm-serial.log \
                > /tmp/radiantwave-test-vm.log 2>&1 & \
            echo \$! > /tmp/radiantwave-test-vm.pid && \
            echo \"VM started on remote (PID \$(cat /tmp/radiantwave-test-vm.pid))\""
        ok "VM started on remote server"
    fi
}

cmd_stop() {
    if [[ "$PLATFORM" == "linux" ]]; then
        if [[ -f /tmp/radiantwave-test-vm.pid ]]; then
            local pid
            pid="$(cat /tmp/radiantwave-test-vm.pid)"
            if kill -0 "$pid" 2>/dev/null; then
                log "Stopping VM (PID $pid)..."
                kill "$pid"
                rm -f /tmp/radiantwave-test-vm.pid
                ok "VM stopped"
            else
                warn "VM process $pid not running"
                rm -f /tmp/radiantwave-test-vm.pid
            fi
        else
            warn "No PID file found. VM may not be running."
        fi
    else
        log "Stopping VM on remote server..."
        ssh "$REMOTE" "
            if [[ -f /tmp/radiantwave-test-vm.pid ]]; then
                pid=\$(cat /tmp/radiantwave-test-vm.pid)
                kill \$pid 2>/dev/null && echo 'VM stopped' || echo 'VM not running'
                rm -f /tmp/radiantwave-test-vm.pid
            else
                echo 'No PID file found'
            fi
        "
    fi
}

vm_is_running() {
    if [[ "$PLATFORM" == "linux" ]]; then
        [[ -f /tmp/radiantwave-test-vm.pid ]] && \
            kill -0 "$(cat /tmp/radiantwave-test-vm.pid)" 2>/dev/null
    else
        ssh "$REMOTE" "
            [[ -f /tmp/radiantwave-test-vm.pid ]] && \
                kill -0 \$(cat /tmp/radiantwave-test-vm.pid) 2>/dev/null
        " 2>/dev/null
    fi
}

vm_ssh() {
    # Run a command in the VM via SSH. Usage: vm_ssh "command"
    local cmd="$1"
    local ssh_opts="-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null"

    if [[ "$PLATFORM" == "linux" ]]; then
        eval "$SSPASS -p '$VM_PASS' ssh $ssh_opts -p '$SSH_PORT' '${VM_USER}@localhost' '$cmd'"
    else
        # On macOS: find sshpass on the remote server, then SSH to the VM
        ssh "$REMOTE" "
            SSPASS=\$(command -v sshpass 2>/dev/null || nix-shell -p sshpass --run 'which sshpass' 2>/dev/null)
            if [[ -z \"\$SSPASS\" ]]; then
                echo 'sshpass not found on remote server' >&2
                exit 1
            fi
            \$SSPASS -p '$VM_PASS' ssh $ssh_opts -p $SSH_PORT ${VM_USER}@localhost '$cmd'
        "
    fi
}

vm_wait_for_boot() {
    local max_wait=90
    local waited=0

    log "Waiting for VM to boot (up to ${max_wait}s)..."

    while (( waited < max_wait )); do
        if vm_ssh "echo ok" &>/dev/null; then
            ok "VM ready after ${waited}s"
            return 0
        fi
        sleep 2
        waited=$((waited + 2))
        echo -ne "\r  ${DIM}${waited}s / ${max_wait}s${NC}"
    done

    echo
    err "VM did not become ready within ${max_wait}s"
    return 1
}

cmd_diagnose() {
    hr
    echo -e "${BOLD}  RadiantWave Test VM — Diagnostic Report${NC}"
    hr

    log "Connecting to VM..."

    # Build and run a diagnostic script on the server (avoids nested quoting issues)
    local diag_script
    diag_script="$(mktemp)"
    cat > "$diag_script" << 'DIAGEOF'
#!/usr/bin/env bash
set -uo pipefail

run_in_vm() {
    sshpass -p kiosk ssh \
        -o StrictHostKeyChecking=no \
        -o UserKnownHostsFile=/dev/null \
        -o ConnectTimeout=5 \
        -p 2222 kiosk@localhost "$@"
}

echo "── System Services ──"
run_in_vm systemctl list-units --type=service --state=running --no-pager | \
    grep -E 'NetworkManager|sshd|seatd|getty|cage|systemd-logind' || true

echo ""
echo "── User Services (PipeWire) ──"
run_in_vm systemctl --user list-units --type=service --state=running --no-pager | \
    grep -E 'pipewire|wireplumber|pulse' || echo "No user services running"

echo ""
echo "── PipeWire Devices (wpctl status) ──"
run_in_vm wpctl status 2>/dev/null || echo "wpctl not available"

echo ""
echo "── ALSA Devices (aplay -l) ──"
run_in_vm aplay -l 2>/dev/null || echo "No ALSA devices"

echo ""
echo "── PulseAudio Sinks (pactl) ──"
run_in_vm pactl list sinks short 2>/dev/null || echo "pactl not available"

echo ""
echo "── Network Devices ──"
run_in_vm nmcli device status 2>/dev/null || echo "nmcli not available"

echo ""
echo "── Key Processes ──"
run_in_vm ps aux | \
    grep -E 'radiantwave|cage|pipewire|wireplumber|NetworkManager' | \
    grep -v grep || echo "No matching processes"

echo ""
echo "── Login Sessions ──"
run_in_vm loginctl list-sessions --no-pager 2>/dev/null || echo "loginctl not available"

echo ""
echo "── App Log (last 30 lines) ──"
run_in_vm tail -30 /tmp/kiosk.log 2>/dev/null || echo "No log file yet"
DIAGEOF

    if [[ "$PLATFORM" == "linux" ]]; then
        bash "$diag_script"
    else
        # Upload and run on remote server
        local remote_diag="/tmp/radiantwave-diagnose-$$.sh"
        scp -q "$diag_script" "$REMOTE:$remote_diag"
        ssh "$REMOTE" "nix-shell -p sshpass --run 'bash $remote_diag'; rm -f $remote_diag"
    fi

    rm -f "$diag_script"
    hr
}

cmd_shell() {
    log "Connecting interactive shell to VM..."
    local ssh_opts="-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null"

    if [[ "$PLATFORM" == "linux" ]]; then
        eval "$SSPASS -p '$VM_PASS' ssh $ssh_opts -p '$SSH_PORT' '${VM_USER}@localhost'"
    else
        log "Connecting via remote server ($REMOTE)..."
        ssh -t "$REMOTE" "
            SSPASS=\$(command -v sshpass 2>/dev/null || nix-shell -p sshpass --run 'which sshpass' 2>/dev/null)
            if [[ -z \"\$SSPASS\" ]]; then
                echo 'sshpass not found on remote server' >&2
                exit 1
            fi
            \$SSPASS -p '$VM_PASS' ssh $ssh_opts -t -p $SSH_PORT ${VM_USER}@localhost
        "
    fi
}

# ── Main ──────────────────────────────────────────────────────────────────────

usage() {
    cat <<EOF
Usage: $(basename "$0") [COMMAND]

Commands:
  (none)        Build + start + run diagnostics
  --build       Build the test VM only
  --start       Start the VM (must already be built)
  --stop        Stop the running VM
  --diagnose    Run diagnostics on a running VM
  --shell       Interactive SSH shell into the VM

Platform:
  macOS (ARM)   Builds/runs on remote server: ${REMOTE}
  Linux (x86_64)  Builds/runs locally

Environment:
  RADIANTWAVE_REMOTE        Remote server (default: nclack@192.168.1.203)
  RADIANTWAVE_REMOTE_DIR    Remote repo path (default: ~/radiantwave)

Requirements:
  nix         Nix package manager
  sshpass     Non-interactive SSH (for diagnostics)
              Install: brew install sshpass (macOS) or nix-shell -p sshpass
EOF
}

main() {
    check_deps

    case "${1:-}" in
        -h|--help)
            usage
            ;;
        --build)
            cmd_build
            ;;
        --start)
            cmd_start
            ;;
        --stop)
            cmd_stop
            ;;
        --diagnose)
            if ! vm_is_running; then
                err "VM is not running. Start it first: $0 --start"
                exit 1
            fi
            cmd_diagnose
            ;;
        --shell)
            if ! vm_is_running; then
                err "VM is not running. Start it first: $0 --start"
                exit 1
            fi
            cmd_shell
            ;;
        "")
            cmd_build
            cmd_start
            sleep 5
            if vm_wait_for_boot; then
                cmd_diagnose
            else
                warn "VM boot timed out. Check logs:"
                if [[ "$PLATFORM" == "linux" ]]; then
                    echo "  tail -f /tmp/radiantwave-test-vm.log"
                    echo "  tail -f /tmp/radiantwave-test-vm-serial.log"
                else
                    echo "  ssh $REMOTE 'tail -f /tmp/radiantwave-test-vm.log'"
                fi
            fi
            ;;
        *)
            err "Unknown command: $1"
            usage
            exit 1
            ;;
    esac
}

main "$@"