#!/bin/bash
set -euo pipefail

# ARM64 Ubuntu QEMU setup with cloud-init for hdmistat testing

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

# QEMU settings
QEMU_MEM="2G"
QEMU_CPUS="4"
DISK_SIZE="8G"
VNC_PORT="5902"

# Ubuntu Cloud Image ARM64
UBUNTU_VERSION="22.04"
UBUNTU_IMAGE="ubuntu-${UBUNTU_VERSION}-server-cloudimg-arm64.img"
UBUNTU_URL="https://cloud-images.ubuntu.com/releases/${UBUNTU_VERSION}/release/${UBUNTU_IMAGE}"

# Colors
GREEN='\033[0;32m'
RED='\033[0;31m'
NC='\033[0m'

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

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

# Check dependencies
check_dependencies() {
    log "Checking dependencies..."
    
    local deps=("qemu-system-aarch64" "qemu-img" "wget" "genisoimage")
    for dep in "${deps[@]}"; do
        if ! command -v "$dep" &> /dev/null; then
            error "$dep not found"
        fi
    done
    
    log "Dependencies OK"
}

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

# Download Ubuntu cloud image
download_ubuntu() {
    cd "$WORK_DIR"
    
    if [ -f "$UBUNTU_IMAGE" ]; then
        log "Ubuntu cloud image already exists"
        return
    fi
    
    log "Downloading Ubuntu ${UBUNTU_VERSION} ARM64 cloud image..."
    wget -O "$UBUNTU_IMAGE" "$UBUNTU_URL" || error "Failed to download Ubuntu image"
}

# Create working disk from cloud image
create_disk() {
    cd "$WORK_DIR"
    
    if [ -f "hdmistat-test.qcow2" ]; then
        log "Working disk already exists"
        return
    fi
    
    log "Creating working disk from cloud image..."
    qemu-img create -f qcow2 -F qcow2 -b "$UBUNTU_IMAGE" hdmistat-test.qcow2 $DISK_SIZE
}

# 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: ubuntu
    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

packages:
  - git
  - build-essential
  - wget
  - fbset

# Enable SSH (already enabled in cloud images)
ssh_pwauth: true

# Update system
package_update: true
package_upgrade: false

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/setup-hdmistat.sh
    permissions: '0755'
    content: |
      #!/bin/bash
      set -e
      export PATH=$PATH:/usr/local/go/bin
      export GOPATH=$HOME/go
      export PATH=$PATH:$GOPATH/bin
      
      # Mount hdmistat source
      sudo mkdir -p /mnt/hdmistat
      sudo mount -t 9p -o trans=virtio,version=9p2000.L hdmistat /mnt/hdmistat || echo "Failed to mount hdmistat source"
      
      # Build hdmistat
      if [ -d "/mnt/hdmistat" ]; then
          echo "Building hdmistat from mounted source..."
          cd /mnt/hdmistat
          make build
          sudo cp hdmistat /usr/local/bin/
      fi
      
      # Create config
      sudo mkdir -p /etc/hdmistat
      sudo tee /etc/hdmistat/config.yaml << 'CONFIG_EOF'
      framebuffer_device: /dev/fb0
      rotation_interval: 10s
      update_interval: 1s
      log_level: info
      width: 1024
      height: 768
      screens:
        - overview
        - top_cpu
        - top_memory
      CONFIG_EOF
      
      # Enable framebuffer
      echo "fbcon=map:10" | sudo tee -a /etc/default/grub
      sudo update-grub || true

runcmd:
  - /tmp/install-go.sh
  - su - ubuntu -c /tmp/setup-hdmistat.sh
  - echo "hdmistat setup complete! Use 'hdmistat daemon' to start"

final_message: "hdmistat test environment ready!"
EOF

    # Create cloud-init ISO
    log "Creating cloud-init ISO..."
    cd "$WORK_DIR"
    genisoimage -output cloud-init.iso -volid cidata -joliet -rock "${CLOUD_INIT_DIR}/user-data" "${CLOUD_INIT_DIR}/meta-data"
}

# Create run script
create_run_script() {
    cat > "${WORK_DIR}/run-qemu.sh" << EOF
#!/bin/bash
cd "$WORK_DIR"

echo "Starting QEMU ARM64 Ubuntu with cloud-init..."
echo ""
echo "Login: ubuntu / ubuntu"
echo "SSH available on port 2223: ssh ubuntu@localhost -p 2223"
echo "VNC available on port $VNC_PORT"
echo ""
echo "hdmistat will be installed automatically via cloud-init"
echo "Check cloud-init status with: cloud-init status"
echo ""
echo "Press Ctrl+A, X to quit"
echo ""

# Find EFI firmware and vars
EFI_CODE=""
EFI_VARS=""
for dir in /nix/store/*/share/qemu /run/current-system/sw/share/qemu; do
    if [ -f "\$dir/edk2-aarch64-code.fd" ]; then
        EFI_CODE="\$dir/edk2-aarch64-code.fd"
        if [ -f "\$dir/edk2-arm-vars.fd" ]; then
            EFI_VARS="\$dir/edk2-arm-vars.fd"
        fi
        break
    fi
done

# Create a copy of EFI vars if found
if [ -n "\$EFI_VARS" ] && [ ! -f "efivars.fd" ]; then
    cp "\$EFI_VARS" efivars.fd
fi

BIOS_ARGS=""
if [ -n "\$EFI_CODE" ]; then
    BIOS_ARGS="-bios \$EFI_CODE"
    if [ -f "efivars.fd" ]; then
        BIOS_ARGS="\$BIOS_ARGS -drive if=pflash,format=raw,file=efivars.fd"
    fi
fi

qemu-system-aarch64 \\
    -M virt \\
    -accel hvf \\
    -cpu host \\
    -smp $QEMU_CPUS \\
    -m $QEMU_MEM \\
    -nographic \\
    -drive file=hdmistat-test.qcow2,format=qcow2,if=virtio \\
    -drive file=cloud-init.iso,format=raw,if=virtio \\
    -netdev user,id=net0,hostfwd=tcp::2223-:22 \\
    -device virtio-net-pci,netdev=net0 \\
    -device virtio-gpu-pci \\
    -display default,show-cursor=on \\
    -vnc :2 \\
    -serial mon:stdio \\
    -virtfs local,path="${PROJECT_DIR}",mount_tag=hdmistat,security_model=none,id=hdmistat \\
    \$BIOS_ARGS \\
    \${@}
EOF
    chmod +x "${WORK_DIR}/run-qemu.sh"
}

# Create helper scripts
create_helpers() {
    # VNC viewer
    cat > "${WORK_DIR}/view-vnc.sh" << EOF
#!/bin/bash
echo "Opening VNC viewer on localhost:$VNC_PORT"
vncviewer localhost:$VNC_PORT || open vnc://localhost:$VNC_PORT || echo "Connect to VNC at localhost:$VNC_PORT"
EOF
    chmod +x "${WORK_DIR}/view-vnc.sh"
    
    # SSH helper
    cat > "${WORK_DIR}/ssh-ubuntu.sh" << EOF
#!/bin/bash
echo "Connecting to Ubuntu VM..."
echo "Password: ubuntu"
ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null ubuntu@localhost -p 2223
EOF
    chmod +x "${WORK_DIR}/ssh-ubuntu.sh"
}

# Print instructions
print_instructions() {
    echo ""
    echo "======================================="
    echo "QEMU ARM64 Ubuntu Setup Complete!"
    echo "======================================="
    echo ""
    echo "To start testing:"
    echo "  cd $WORK_DIR"
    echo "  ./run-qemu.sh"
    echo ""
    echo "Features:"
    echo "  - Ubuntu ${UBUNTU_VERSION} ARM64 with cloud-init"
    echo "  - Native ARM64 virtualization (hvf)"
    echo "  - Automatic hdmistat installation"
    echo "  - Framebuffer support"
    echo "  - SSH access on port 2223"
    echo "  - VNC access on port $VNC_PORT"
    echo ""
    echo "The cloud-init config is in: $CLOUD_INIT_DIR/"
    echo "You can modify it and regenerate the ISO with:"
    echo "  genisoimage -output cloud-init.iso -volid cidata -joliet -rock user-data meta-data"
    echo ""
}

# Main
main() {
    log "Setting up QEMU ARM64 Ubuntu environment with cloud-init..."
    
    check_dependencies
    setup_directories
    download_ubuntu
    create_disk
    create_cloud_init
    create_run_script
    create_helpers
    print_instructions
    
    log "Setup complete!"
}

main "$@"