#!/bin/bash
set -euo pipefail

# Script to run Raspberry Pi OS in QEMU with hdmistat installation

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
WORK_DIR="${PROJECT_DIR}/qemu-test"
CLOUD_INIT_DIR="${WORK_DIR}/cloud-init"

# QEMU and image settings
QEMU_ARCH="aarch64"
RPI_IMAGE_URL="https://downloads.raspberrypi.org/raspios_lite_arm64/images/raspios_lite_arm64-2024-03-15/2024-03-15-raspios-bookworm-arm64-lite.img.xz"
RPI_IMAGE_NAME="raspios-lite-arm64.img"
QEMU_MEM="2G"
QEMU_CPUS="4"
VNC_PORT="5901"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
NC='\033[0m' # No Color

log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $*"
}

error() {
    echo -e "${RED}[ERROR]${NC} $*" >&2
    exit 1
}

warn() {
    echo -e "${YELLOW}[WARN]${NC} $*"
}

# Check dependencies
check_dependencies() {
    log "Checking dependencies..."
    
    local deps=("qemu-system-aarch64" "cloud-localds" "xz" "wget")
    local missing=()
    
    for dep in "${deps[@]}"; do
        if ! command -v "$dep" &> /dev/null; then
            missing+=("$dep")
        fi
    done
    
    if [ ${#missing[@]} -ne 0 ]; then
        error "Missing dependencies: ${missing[*]}"
    fi
    
    log "All dependencies found"
}

# Create work directory
setup_directories() {
    log "Setting up directories..."
    mkdir -p "$WORK_DIR"
    mkdir -p "$CLOUD_INIT_DIR"
}

# Download Raspberry Pi OS image
download_image() {
    if [ -f "${WORK_DIR}/${RPI_IMAGE_NAME}" ]; then
        log "Raspberry Pi OS image already exists"
        return
    fi
    
    log "Downloading Raspberry Pi OS image..."
    cd "$WORK_DIR"
    
    if [ ! -f "${RPI_IMAGE_NAME}.xz" ]; then
        wget -O "${RPI_IMAGE_NAME}.xz" "$RPI_IMAGE_URL" || error "Failed to download image"
    fi
    
    log "Extracting image..."
    xz -d -k "${RPI_IMAGE_NAME}.xz" || error "Failed to extract image"
    
    # Resize image to have more space
    log "Resizing image to 8GB..."
    qemu-img resize "${RPI_IMAGE_NAME}" 8G
}

# Create cloud-init configuration
create_cloud_init() {
    log "Creating cloud-init configuration..."
    
    # Create meta-data
    cat > "${CLOUD_INIT_DIR}/meta-data" << 'EOF'
instance-id: hdmistat-test-01
local-hostname: hdmistat-test
EOF

    # Create user-data with hdmistat installation
    cat > "${CLOUD_INIT_DIR}/user-data" << 'EOF'
#cloud-config
hostname: hdmistat-test
manage_etc_hosts: true

users:
  - name: pi
    groups: [adm, audio, cdrom, dialout, dip, floppy, lxd, netdev, plugdev, sudo, video]
    sudo: ALL=(ALL) NOPASSWD:ALL
    shell: /bin/bash
    lock_passwd: false
    passwd: $6$rounds=4096$8VNRmFT9yR$7TVlOepTJjMW0CzBkpMrdWA7aH4rZ94pZng8XjqaY8d8qqBmFvO/hYL5L2gqJ5nKLhDR8Jz9Z9nqfHBl5kVZHe1
    ssh_authorized_keys:
      - ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDf0q4PyG0doiBQYV7OlOxbRjle026hJPBWbZe test@hdmistat

packages:
  - git
  - golang
  - build-essential
  - systemd
  - htop
  - tmux

# Enable SSH
ssh_pwauth: true

# Update system
package_update: true
package_upgrade: true

# Install Go 1.24.4
write_files:
  - path: /tmp/install-go.sh
    permissions: '0755'
    content: |
      #!/bin/bash
      set -e
      GO_VERSION="1.24.4"
      wget -q -O /tmp/go.tar.gz "https://go.dev/dl/go${GO_VERSION}.linux-arm64.tar.gz"
      sudo rm -rf /usr/local/go
      sudo tar -C /usr/local -xzf /tmp/go.tar.gz
      echo 'export PATH=$PATH:/usr/local/go/bin' | sudo tee -a /etc/profile
      echo 'export GOPATH=$HOME/go' | sudo tee -a /etc/profile
      echo 'export PATH=$PATH:$GOPATH/bin' | sudo tee -a /etc/profile

  - path: /tmp/build-hdmistat.sh
    permissions: '0755'
    content: |
      #!/bin/bash
      set -e
      export PATH=$PATH:/usr/local/go/bin
      export GOPATH=/home/pi/go
      export PATH=$PATH:$GOPATH/bin
      
      # Clone and build hdmistat from local source
      mkdir -p /home/pi/hdmistat-src
      cd /home/pi/hdmistat-src
      
      # Note: In real usage, you'd clone from git repo
      # For testing, we'll mount the source directory
      
      # Build hdmistat
      if [ -d "/mnt/hdmistat" ]; then
        cd /mnt/hdmistat
        make build
        sudo cp hdmistat /usr/local/bin/
      else
        # Fallback: try to install from git
        go install git.eeqj.de/sneak/hdmistat/cmd/hdmistat@latest
        sudo cp $GOPATH/bin/hdmistat /usr/local/bin/
      fi
      
      # Install as systemd service
      sudo /usr/local/bin/hdmistat install
      
      # Enable framebuffer console
      sudo sed -i 's/console=serial0,115200 //' /boot/cmdline.txt
      echo "hdmi_force_hotplug=1" | sudo tee -a /boot/config.txt
      echo "hdmi_group=2" | sudo tee -a /boot/config.txt
      echo "hdmi_mode=82" | sudo tee -a /boot/config.txt
      echo "framebuffer_width=1920" | sudo tee -a /boot/config.txt
      echo "framebuffer_height=1080" | sudo tee -a /boot/config.txt

# Run installation scripts
runcmd:
  - /tmp/install-go.sh
  - sleep 5
  - /tmp/build-hdmistat.sh
  - sudo systemctl daemon-reload
  - sudo systemctl enable hdmistat
  - sudo systemctl start hdmistat

# Final message
final_message: "hdmistat installation complete! System ready for testing."
EOF

    # Create cloud-init ISO
    log "Creating cloud-init ISO..."
    cd "$WORK_DIR"
    cloud-localds cloud-init.iso "${CLOUD_INIT_DIR}/user-data" "${CLOUD_INIT_DIR}/meta-data"
}

# Create QEMU run script
create_run_script() {
    log "Creating QEMU run script..."
    
    cat > "${WORK_DIR}/run-qemu.sh" << EOF
#!/bin/bash
# Run QEMU with Raspberry Pi emulation

cd "$WORK_DIR"

echo "Starting QEMU Raspberry Pi emulation..."
echo "VNC server will be available on port $VNC_PORT"
echo "SSH: ssh pi@localhost -p 2222 (password: raspberry)"
echo "Framebuffer: The emulated display supports framebuffer at /dev/fb0"
echo ""
echo "Press Ctrl+A, X to quit QEMU"

qemu-system-aarch64 \\
    -M raspi3b \\
    -cpu cortex-a72 \\
    -smp $QEMU_CPUS \\
    -m $QEMU_MEM \\
    -kernel kernel8.img \\
    -dtb bcm2710-rpi-3-b-plus.dtb \\
    -append "rw earlyprintk loglevel=8 console=ttyAMA0,115200 dwc_otg.lpm_enable=0 root=/dev/mmcblk0p2 rootdelay=1" \\
    -drive file=${RPI_IMAGE_NAME},format=raw,if=sd \\
    -drive file=cloud-init.iso,format=raw,if=none,id=cloud-init \\
    -device usb-storage,drive=cloud-init \\
    -netdev user,id=net0,hostfwd=tcp::2222-:22 \\
    -device usb-net,netdev=net0 \\
    -vnc :1 \\
    -serial stdio \\
    -display none \\
    -virtfs local,path="${PROJECT_DIR}",mount_tag=hdmistat,security_model=passthrough,id=hdmistat \\
    \${@}
EOF

    chmod +x "${WORK_DIR}/run-qemu.sh"
}

# Download kernel and DTB files
download_kernel() {
    log "Downloading kernel and DTB files..."
    
    cd "$WORK_DIR"
    
    # These are required for booting Raspberry Pi OS in QEMU
    if [ ! -f "kernel8.img" ]; then
        wget -q https://github.com/dhruvvyas90/qemu-rpi-kernel/raw/master/kernel8.img || \
            error "Failed to download kernel"
    fi
    
    if [ ! -f "bcm2710-rpi-3-b-plus.dtb" ]; then
        wget -q https://github.com/dhruvvyas90/qemu-rpi-kernel/raw/master/bcm2710-rpi-3-b-plus.dtb || \
            error "Failed to download DTB"
    fi
}

# Create helper scripts
create_helpers() {
    log "Creating helper scripts..."
    
    # VNC viewer script
    cat > "${WORK_DIR}/view-vnc.sh" << EOF
#!/bin/bash
echo "Opening VNC viewer on localhost:$VNC_PORT"
vncviewer localhost:$VNC_PORT || echo "VNC viewer not found. Connect manually to localhost:$VNC_PORT"
EOF
    chmod +x "${WORK_DIR}/view-vnc.sh"
    
    # SSH script
    cat > "${WORK_DIR}/ssh-pi.sh" << EOF
#!/bin/bash
echo "Connecting to Raspberry Pi via SSH..."
echo "Default password is: raspberry"
ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null pi@localhost -p 2222
EOF
    chmod +x "${WORK_DIR}/ssh-pi.sh"
    
    # Mount script for 9p filesystem
    cat > "${WORK_DIR}/mount-hdmistat.sh" << EOF
#!/bin/bash
# Run this inside the VM to mount the hdmistat source
echo "Run this inside the Raspberry Pi VM:"
echo "sudo mkdir -p /mnt/hdmistat"
echo "sudo mount -t 9p -o trans=virtio hdmistat /mnt/hdmistat -oversion=9p2000.L"
EOF
    chmod +x "${WORK_DIR}/mount-hdmistat.sh"
}

# Print usage instructions
print_instructions() {
    log "Setup complete!"
    echo ""
    echo "========================================="
    echo "QEMU Raspberry Pi Test Environment Ready"
    echo "========================================="
    echo ""
    echo "To start the emulated Raspberry Pi:"
    echo "  cd $WORK_DIR"
    echo "  ./run-qemu.sh"
    echo ""
    echo "To view the display (framebuffer):"
    echo "  ./view-vnc.sh"
    echo ""
    echo "To SSH into the system:"
    echo "  ./ssh-pi.sh"
    echo ""
    echo "To mount hdmistat source inside VM:"
    echo "  1. SSH into the VM"
    echo "  2. Run: sudo mount -t 9p -o trans=virtio hdmistat /mnt/hdmistat -oversion=9p2000.L"
    echo ""
    echo "Default credentials:"
    echo "  Username: pi"
    echo "  Password: raspberry"
    echo ""
    echo "The framebuffer device will be available at /dev/fb0"
    echo "hdmistat will be installed automatically via cloud-init on first boot"
    echo ""
}

# Main execution
main() {
    log "Starting QEMU Raspberry Pi setup for hdmistat testing..."
    
    check_dependencies
    setup_directories
    download_image
    download_kernel
    create_cloud_init
    create_run_script
    create_helpers
    print_instructions
    
    log "Setup completed successfully!"
}

# Run main function
main "$@"