#!/bin/bash
set -e

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"

# Parse arguments
AUTO_MODE=false
TARGET_ARCH="amd64"

for arg in "$@"; do
    case $arg in
        --auto)
            AUTO_MODE=true
            shift
            ;;
        --arch=*)
            TARGET_ARCH="${arg#*=}"
            shift
            ;;
        -h|--help)
            echo "Usage: $0 [--arch=amd64|arm64|rpi] [--auto]"
            echo ""
            echo "Options:"
            echo "  --arch=ARCH  Target architecture (default: amd64)"
            echo "               amd64 - x86_64 with UEFI/OVMF"
            echo "               arm64 - Generic ARM64 with UEFI"
            echo "               rpi   - Raspberry Pi 3 emulation"
            echo "  --auto       Run in automatic mode for CI/testing"
            exit 0
            ;;
    esac
done

# Validate architecture
case "$TARGET_ARCH" in
    amd64|arm64|rpi)
        ;;
    *)
        echo "ERROR: Invalid architecture '$TARGET_ARCH'"
        echo "Valid options: amd64, arm64, rpi"
        exit 1
        ;;
esac

PROJECT_ROOT="$SCRIPT_DIR/.."
VM_DIR="$PROJECT_ROOT/.nogit/vm"
MONITOR_SOCK="$VM_DIR/qemu-monitor.sock"
SERIAL_SOCK="$VM_DIR/serial.sock"
SERIAL_LOG="$VM_DIR/serial.log"
PID_FILE="$VM_DIR/qemu.pid"

# Architecture-specific settings
case "$TARGET_ARCH" in
    amd64)
        ISO_PATH="$PROJECT_ROOT/.nogit/iso/ecoos.iso"
        DISK_PATH="$VM_DIR/test-disk.qcow2"
        QEMU_CMD="qemu-system-x86_64"
        QEMU_MACHINE=""
        QEMU_BIOS="-bios /usr/share/qemu/OVMF.fd"
        KVM_CHECK="/dev/kvm"
        DISK_IF="virtio"
        ;;
    arm64)
        ISO_PATH="$PROJECT_ROOT/.nogit/iso/ecoos-arm64.iso"
        DISK_PATH="$VM_DIR/test-disk-arm64.qcow2"
        QEMU_CMD="qemu-system-aarch64"
        QEMU_MACHINE="-M virt -cpu cortex-a72"
        QEMU_BIOS="-bios /usr/share/qemu-efi-aarch64/QEMU_EFI.fd"
        KVM_CHECK=""  # ARM KVM only works on ARM hosts
        DISK_IF="virtio"
        ;;
    rpi)
        IMG_PATH="$PROJECT_ROOT/.nogit/iso/ecoos-rpi.img"
        DISK_PATH=""  # RPi uses the image directly
        QEMU_CMD="qemu-system-aarch64"
        QEMU_MACHINE="-M raspi3b -cpu cortex-a53"
        QEMU_BIOS=""  # RPi uses direct kernel boot
        KVM_CHECK=""
        DISK_IF="sd"
        ;;
esac

# Create VM directory if not exists
mkdir -p "$VM_DIR"

# Check if image exists
if [ "$TARGET_ARCH" = "rpi" ]; then
    if [ ! -f "$IMG_PATH" ]; then
        echo "ERROR: RPi image not found at $IMG_PATH"
        echo "Run 'pnpm run build:rpi' first to create the image"
        exit 1
    fi
else
    if [ ! -f "$ISO_PATH" ]; then
        echo "ERROR: ISO not found at $ISO_PATH"
        echo "Run 'pnpm run build:$TARGET_ARCH' first to create the ISO"
        exit 1
    fi
fi

# Create test disk if not exists (not needed for RPi)
if [ -n "$DISK_PATH" ] && [ ! -f "$DISK_PATH" ]; then
    echo "Creating test disk (20GB)..."
    qemu-img create -f qcow2 "$DISK_PATH" 20G
fi

# Check if already running
if [ -f "$PID_FILE" ]; then
    PID=$(cat "$PID_FILE")
    if kill -0 "$PID" 2>/dev/null; then
        echo "QEMU already running (PID: $PID)"
        echo "Run 'pnpm run test:stop' to stop it first"
        exit 1
    fi
fi

echo "Starting QEMU with EcoOS for $TARGET_ARCH..."

# Check if KVM is available (only for amd64 on x86 hosts)
KVM_OPTS=""
if [ -n "$KVM_CHECK" ] && [ -e "$KVM_CHECK" ] && [ -r "$KVM_CHECK" ] && [ -w "$KVM_CHECK" ]; then
    KVM_OPTS="-enable-kvm -cpu host"
    echo "Using KVM acceleration"
else
    if [ "$TARGET_ARCH" = "amd64" ]; then
        echo "KVM not available, using software emulation (slower)"
    else
        echo "Running ARM emulation on x86 host (slower)"
    fi
fi

# Cleanup function
cleanup() {
    echo ""
    echo "Shutting down..."
    if [ -n "$SCREENSHOT_LOOP_PID" ] && kill -0 "$SCREENSHOT_LOOP_PID" 2>/dev/null; then
        kill "$SCREENSHOT_LOOP_PID" 2>/dev/null || true
    fi
    if [ -n "$ENABLE_PID" ] && kill -0 "$ENABLE_PID" 2>/dev/null; then
        kill "$ENABLE_PID" 2>/dev/null || true
    fi
    if [ -n "$VIEWER_PID" ] && kill -0 "$VIEWER_PID" 2>/dev/null; then
        kill "$VIEWER_PID" 2>/dev/null || true
    fi
    if [ -n "$TWM_PID" ] && kill -0 "$TWM_PID" 2>/dev/null; then
        kill "$TWM_PID" 2>/dev/null || true
    fi
    if [ -n "$XORG_PID" ] && kill -0 "$XORG_PID" 2>/dev/null; then
        kill "$XORG_PID" 2>/dev/null || true
    fi
    if [ -f "$PID_FILE" ]; then
        PID=$(cat "$PID_FILE")
        if kill -0 "$PID" 2>/dev/null; then
            kill "$PID" 2>/dev/null || true
        fi
        rm -f "$PID_FILE"
    fi
    echo "Done"
}
trap cleanup EXIT INT TERM

# Start QEMU based on architecture
> "$SERIAL_LOG"  # Clear old log

if [ "$TARGET_ARCH" = "amd64" ]; then
    # AMD64 with multi-display support
    $QEMU_CMD \
      $KVM_OPTS \
      -m 4G \
      -smp 4 \
      $QEMU_BIOS \
      -drive file="$ISO_PATH",media=cdrom \
      -drive file="$DISK_PATH",format=qcow2,if=virtio \
      -device qxl-vga,id=video0,ram_size=67108864,vram_size=67108864,vgamem_mb=64 \
      -device qxl,id=video1,ram_size=67108864,vram_size=67108864,vgamem_mb=64 \
      -device qxl,id=video2,ram_size=67108864,vram_size=67108864,vgamem_mb=64 \
      -display none \
      -spice port=5930,disable-ticketing=on \
      -device virtio-serial-pci \
      -chardev spicevmc,id=vdagent,name=vdagent \
      -device virtserialport,chardev=vdagent,name=com.redhat.spice.0 \
      -serial unix:"$SERIAL_SOCK",server,nowait \
      -monitor unix:"$MONITOR_SOCK",server,nowait \
      -nic user,model=virtio-net-pci,hostfwd=tcp::3006-:3006,hostfwd=tcp::2222-:22 \
      -pidfile "$PID_FILE" &

elif [ "$TARGET_ARCH" = "arm64" ]; then
    # ARM64 with UEFI
    $QEMU_CMD \
      $QEMU_MACHINE \
      -m 4G \
      -smp 4 \
      $QEMU_BIOS \
      -drive file="$ISO_PATH",media=cdrom,if=none,id=cdrom \
      -device virtio-blk-device,drive=cdrom \
      -drive file="$DISK_PATH",format=qcow2,if=none,id=hd0 \
      -device virtio-blk-device,drive=hd0 \
      -device virtio-gpu-pci \
      -display none \
      -serial unix:"$SERIAL_SOCK",server,nowait \
      -monitor unix:"$MONITOR_SOCK",server,nowait \
      -device virtio-net-device,netdev=net0 \
      -netdev user,id=net0,hostfwd=tcp::3006-:3006,hostfwd=tcp::2222-:22 \
      -pidfile "$PID_FILE" &

elif [ "$TARGET_ARCH" = "rpi" ]; then
    # Raspberry Pi 3B emulation
    # Note: raspi3b machine has limited support, uses direct kernel boot
    echo "NOTE: Raspberry Pi emulation is limited."
    echo "      For full testing, use real hardware."

    # Extract kernel and initrd from image for direct boot
    TEMP_MNT=$(mktemp -d)
    LOOP_DEV=$(sudo losetup --find --show --partscan "$IMG_PATH")
    sudo mount "${LOOP_DEV}p1" "$TEMP_MNT"

    KERNEL="$TEMP_MNT/vmlinuz"
    INITRD="$TEMP_MNT/initrd.img"
    DTB="$TEMP_MNT/bcm2710-rpi-3-b.dtb"

    if [ ! -f "$KERNEL" ]; then
        echo "ERROR: Kernel not found in RPi image"
        sudo umount "$TEMP_MNT"
        sudo losetup -d "$LOOP_DEV"
        rm -rf "$TEMP_MNT"
        exit 1
    fi

    # Copy kernel/initrd to temp location for QEMU
    cp "$KERNEL" "$VM_DIR/rpi-kernel"
    cp "$INITRD" "$VM_DIR/rpi-initrd" 2>/dev/null || true
    cp "$DTB" "$VM_DIR/rpi-dtb" 2>/dev/null || true

    sudo umount "$TEMP_MNT"
    sudo losetup -d "$LOOP_DEV"
    rm -rf "$TEMP_MNT"

    $QEMU_CMD \
      $QEMU_MACHINE \
      -m 1G \
      -kernel "$VM_DIR/rpi-kernel" \
      -initrd "$VM_DIR/rpi-initrd" \
      -dtb "$VM_DIR/rpi-dtb" \
      -append "console=ttyAMA0,115200 root=LABEL=EcoOS rootfstype=ext4 rootwait" \
      -drive file="$IMG_PATH",format=raw,if=sd \
      -serial unix:"$SERIAL_SOCK",server,nowait \
      -display none \
      -pidfile "$PID_FILE" &
fi

QEMU_PID=$!

echo ""
echo "=== EcoOS Test VM Started ($TARGET_ARCH) ==="
echo "QEMU PID: $QEMU_PID"
if [ "$TARGET_ARCH" != "rpi" ]; then
    echo "Management UI: http://localhost:3006"
fi
echo ""

# AMD64-specific display setup
if [ "$TARGET_ARCH" = "amd64" ]; then
    # Wait for QEMU to start and SPICE to be ready
    echo "Waiting for SPICE server..."
    sleep 3

    # Check if remote-viewer is available
    if ! command -v remote-viewer &> /dev/null; then
        echo "WARNING: remote-viewer not installed"
        echo "Install with: sudo apt install virt-viewer"
        echo ""
        echo "Running without display viewer. Press Ctrl-C to stop."
        wait $QEMU_PID
        exit 0
    fi

    # Set up virt-viewer settings for multi-display
    VIRT_VIEWER_CONFIG_DIR="${XDG_CONFIG_HOME:-$HOME/.config}/virt-viewer"
    mkdir -p "$VIRT_VIEWER_CONFIG_DIR"
    if [ -f "$SCRIPT_DIR/virt-viewer-settings" ]; then
        cp "$SCRIPT_DIR/virt-viewer-settings" "$VIRT_VIEWER_CONFIG_DIR/settings"
        echo "Configured virt-viewer for 3 displays"
    fi

    # Detect DISPLAY if not set
    if [ -z "$DISPLAY" ]; then
        if [ -S /tmp/.X11-unix/X0 ]; then
            export DISPLAY=:0
        elif [ -S /tmp/.X11-unix/X1 ]; then
            export DISPLAY=:1
        fi
    fi

    # Detect WAYLAND_DISPLAY if not set
    if [ -z "$WAYLAND_DISPLAY" ] && [ -z "$DISPLAY" ]; then
        if [ -S "$XDG_RUNTIME_DIR/wayland-0" ]; then
            export WAYLAND_DISPLAY=wayland-0
        elif [ -S "/run/user/$(id -u)/wayland-0" ]; then
            export XDG_RUNTIME_DIR="/run/user/$(id -u)"
            export WAYLAND_DISPLAY=wayland-0
        fi
    fi

    # Launch remote-viewer
    if [ -z "$DISPLAY" ] && [ -z "$WAYLAND_DISPLAY" ]; then
        echo "No display found, starting headless X server with 3 virtual monitors..."

        XDISPLAY=99
        while [ -S "/tmp/.X11-unix/X$XDISPLAY" ]; do
            XDISPLAY=$((XDISPLAY + 1))
        done

        XORG_CONFIG="$SCRIPT_DIR/xorg-dummy.conf"
        Xorg :$XDISPLAY -config "$XORG_CONFIG" -noreset +extension GLX +extension RANDR +extension RENDER &
        XORG_PID=$!
        sleep 2

        export DISPLAY=:$XDISPLAY

        xrandr --newmode "1920x1080" 173.00 1920 2048 2248 2576 1080 1083 1088 1120 -hsync +vsync 2>/dev/null || true
        xrandr --addmode DUMMY1 "1920x1080" 2>/dev/null || true
        xrandr --addmode DUMMY2 "1920x1080" 2>/dev/null || true
        xrandr --output DUMMY0 --mode 1920x1080 --pos 0x0 --primary
        xrandr --output DUMMY1 --mode 1920x1080 --pos 1920x0 2>/dev/null || true
        xrandr --output DUMMY2 --mode 1920x1080 --pos 3840x0 2>/dev/null || true

        echo "Headless X server started on :$XDISPLAY"

        remote-viewer --full-screen spice://localhost:5930 &
        VIEWER_PID=$!
        echo "remote-viewer running headlessly (PID: $VIEWER_PID)"
    else
        echo "Launching remote-viewer with fullscreen for multi-display..."
        remote-viewer --full-screen spice://localhost:5930 &
        VIEWER_PID=$!
    fi
    echo ""

    # Enable all 3 displays via SPICE protocol
    if [ -f "$SCRIPT_DIR/enable-displays.py" ]; then
        echo "Enabling displays (waiting for SPICE agent, up to 5 minutes)..."
        python3 "$SCRIPT_DIR/enable-displays.py" --timeout 300 2>&1 &
        ENABLE_PID=$!
    fi

    # Start screenshot loop in background
    echo "Starting screenshot loop..."
    (while true; do "$SCRIPT_DIR/screenshot.sh" 2>/dev/null; sleep 5; done) &
    SCREENSHOT_LOOP_PID=$!
fi

echo "Tips:"
echo "  - socat - UNIX-CONNECT:.nogit/vm/serial.sock  - Serial console (login: ecouser/ecouser)"
if [ "$TARGET_ARCH" != "rpi" ]; then
    echo "  - http://localhost:3006     - Management UI"
fi
echo ""

if [ "$AUTO_MODE" = true ] && [ "$TARGET_ARCH" = "amd64" ]; then
    echo "=== Auto mode: waiting for display setup ==="

    if [ -n "$ENABLE_PID" ]; then
        wait $ENABLE_PID
        ENABLE_EXIT=$?
        if [ $ENABLE_EXIT -ne 0 ]; then
            echo "FAIL: Could not enable displays (exit code: $ENABLE_EXIT)"
            exit 1
        fi
    fi

    echo "Taking screenshot..."
    "$SCRIPT_DIR/screenshot.sh"

    SCREENSHOT="$PROJECT_ROOT/.nogit/screenshots/latest.png"
    if [ -f "$SCREENSHOT" ]; then
        WIDTH=$(identify -format "%w" "$SCREENSHOT" 2>/dev/null || echo "0")
        if [ "$WIDTH" -ge 5760 ]; then
            echo "SUCCESS: Multi-display test passed (width: ${WIDTH}px)"
            exit 0
        else
            echo "FAIL: Screenshot width is ${WIDTH}px, expected >= 5760px"
            exit 1
        fi
    else
        echo "FAIL: Screenshot not found"
        exit 1
    fi
else
    echo "=== Press Ctrl-C to stop ==="
    echo ""
    wait $QEMU_PID 2>/dev/null || true
fi