#!/bin/bash

echo "Don't forget to allow TCP port 5901 in your instance's security group. You have to add it manually to the inbound rules."
echo "This script is tested on Ubuntu."
echo "You have to use a VNC client to connect to your server. You can use Vinagre for example."



# Detect the current user
USER=$(whoami)
echo "Detected user: $USER"

# Create aws directory and set the environment variable
mkdir -p ~/aws
export AWS_FPGA_REPO_DIR=~/aws/

# Update package list and install necessary utilities
sudo apt update
sudo apt install -y git htop awscli # dont install cmake, its too old to build xrt from src!

# Configure AWS CLI (will prompt for keys)
aws configure

# Install kernel headers and Xfce dependencies
sudo apt install -y linux-headers-$(uname -r)
sudo apt install -y xfce4 xfce4-goodies xorg dbus-x11 x11-xserver-utils

# Define the path to the xstartup file for VNC
XSTARTUP_FILE="$HOME/.vnc/xstartup"

# Create the .vnc directory if it doesn't exist
mkdir -p "$(dirname "$XSTARTUP_FILE")"

# Write the content to the xstartup file with Full HD resolution
cat <<EOL > "$XSTARTUP_FILE"
#!/bin/bash
export XDG_SESSION_TYPE=x11
export XDG_SESSION_DESKTOP=xfce
export XDG_CURRENT_DESKTOP=XFCE
# Set screen resolution to Full HD (1920x1080)
xrandr --output VNC-0 --mode 1920x1080
exec startxfce4 &
EOL

# Make the xstartup file executable
chmod +x "$XSTARTUP_FILE"
echo "xstartup file created and configured."

# Set graphical target as the default runlevel
sudo systemctl set-default graphical.target

# Install and configure TigerVNC and xRDP
sudo apt install -y xrdp tigervnc-standalone-server tigervnc-common

# Start and enable xRDP
sudo systemctl start xrdp
sudo systemctl enable xrdp

# Disable UFW firewall for VNC (be cautious)
sudo systemctl stop ufw
sudo systemctl disable ufw

# Set password for the detected user
sudo passwd $USER

# Start or restart VNC server with Full HD resolution
vncserver -kill :1
vncserver :1 -geometry 1920x1080 -localhost no

# Clone the AWS FPGA repository
git clone https://github.com/aws/aws-fpga.git $AWS_FPGA_REPO_DIR
cd $AWS_FPGA_REPO_DIR
git checkout f1_xdma_shell # ubuntu AMI is F1 - Vitis 2024.1, so we need to use this branch.
source vitis_setup.sh

# Check if the variable is already defined in ~/.bashrc
if grep -q "export AWS_PLATFORM=" ~/.bashrc; then
    echo "AWS_PLATFORM is already defined in ~/.bashrc. Updating it to the new value."
    # Update the existing value in ~/.bashrc
    sed -i "s|^export AWS_PLATFORM=.*|export AWS_PLATFORM=\"$AWS_PLATFORM\"|" ~/.bashrc
else
    echo "Adding AWS_PLATFORM to ~/.bashrc."
    # Append the export command to ~/.bashrc
    echo "export AWS_PLATFORM=\"$AWS_PLATFORM\"" >> ~/.bashrc
fi

PLATFORM_REPO_PATHS="$(dirname "$AWS_PLATFORM")"

# Check if PLATFORM_REPO_PATHS is already defined in ~/.bashrc
if grep -q "export PLATFORM_REPO_PATHS=" ~/.bashrc; then
    echo "PLATFORM_REPO_PATHS is already defined in ~/.bashrc. Updating it to the new value."
    # Update the existing value in ~/.bashrc
    sed -i "s|^export PLATFORM_REPO_PATHS=.*|export PLATFORM_REPO_PATHS=\"$PLATFORM_REPO_PATHS\"|" ~/.bashrc
else
    echo "Adding PLATFORM_REPO_PATHS to ~/.bashrc."
    # Append the export command to ~/.bashrc
    echo "export PLATFORM_REPO_PATHS=\"$PLATFORM_REPO_PATHS\"" >> ~/.bashrc
fi

source ~/.bashrc

# Print the path to the VNC log file
VNC_LOG_PATH="$HOME/.vnc/$(hostname):1.log"
echo "VNC log file can be found at: $VNC_LOG_PATH"

echo "To confirm that the setup is correct, please run: /opt/Xilinx/Vitis/2024.1/bin/platforminfo -l" 

echo "==========================================="

echo "** AWS_PLATFORM is set to: $AWS_PLATFORM"
echo "** PLATFORM_REPO_PATHS is set to: $PLATFORM_REPO_PATHS"

# Fetch the public IP address
PUBLIC_IP=$(curl -s ifconfig.me)

echo "** VNC is now available at: $PUBLIC_IP:5901"