#!/bin/bash
# =============================================================================
# Drone Simulation - Ubuntu/Debian Installation Script
# =============================================================================
# Installs ROS 2 Humble/Jazzy, PyBullet, and all required dependencies
# Uses a Python virtual environment for pip packages (PEP 668 compliant)
#
# Usage:
#   chmod +x install_ubuntu.sh
#   ./install_ubuntu.sh
#
# Tested on: Ubuntu 22.04 LTS, Ubuntu 24.04 LTS
# =============================================================================

set -e  # Exit on error

echo "=============================================="
echo "  Drone Simulation - Ubuntu Installation"
echo "=============================================="

# Get the script directory and project root
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
VENV_DIR="$PROJECT_ROOT/venv"

echo "[INFO] Project root: $PROJECT_ROOT"
echo "[INFO] Virtual environment: $VENV_DIR"

# Detect Ubuntu version
UBUNTU_VERSION=$(lsb_release -rs)
echo "[INFO] Detected Ubuntu version: $UBUNTU_VERSION"

# Determine ROS 2 distribution based on Ubuntu version
if [[ "$UBUNTU_VERSION" == "22.04" ]]; then
    ROS_DISTRO="humble"
elif [[ "$UBUNTU_VERSION" == "24.04" ]]; then
    ROS_DISTRO="jazzy"
else
    echo "[WARN] Ubuntu $UBUNTU_VERSION may not be officially supported."
    echo "       Attempting to install ROS 2 Humble..."
    ROS_DISTRO="humble"
fi

echo "[INFO] Will install ROS 2 $ROS_DISTRO"

# -----------------------------------------------------------------------------
# Step 1: Set Locale
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 1/7] Setting locale..."
sudo apt update && sudo apt install -y locales
sudo locale-gen en_US en_US.UTF-8
sudo update-locale LC_ALL=en_US.UTF-8 LANG=en_US.UTF-8
export LANG=en_US.UTF-8

# -----------------------------------------------------------------------------
# Step 2: Add ROS 2 Repository
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 2/7] Adding ROS 2 apt repository..."
sudo apt install -y software-properties-common
sudo add-apt-repository -y universe
sudo apt update && sudo apt install -y curl

sudo curl -sSL https://raw.githubusercontent.com/ros/rosdistro/master/ros.key \
    -o /usr/share/keyrings/ros-archive-keyring.gpg

echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/ros-archive-keyring.gpg] http://packages.ros.org/ros2/ubuntu $(. /etc/os-release && echo $UBUNTU_CODENAME) main" \
    | sudo tee /etc/apt/sources.list.d/ros2.list > /dev/null

# -----------------------------------------------------------------------------
# Step 3: Install ROS 2
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 3/8] Installing ROS 2 $ROS_DISTRO..."
sudo apt update
sudo apt install -y ros-${ROS_DISTRO}-desktop

# Install development tools
sudo apt install -y python3-colcon-common-extensions python3-rosdep

# Install Gazebo and ROS-Gazebo bridge
echo "[INFO] Installing Gazebo and ros_gz_bridge..."
sudo apt install -y ros-${ROS_DISTRO}-ros-gz ros-${ROS_DISTRO}-ros-gz-bridge

# -----------------------------------------------------------------------------
# Step 4: Initialize rosdep
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 4/8] Initializing rosdep..."
if [ ! -f /etc/ros/rosdep/sources.list.d/20-default.list ]; then
    sudo rosdep init
fi
rosdep update

# -----------------------------------------------------------------------------
# Step 5: Create Python Virtual Environment
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 5/8] Creating Python virtual environment..."

# Install venv package if not present
sudo apt install -y python3-venv python3-full

# Create virtual environment with access to system site-packages
# This allows access to ROS 2 packages installed via apt
python3 -m venv "$VENV_DIR" --system-site-packages

echo "[INFO] Virtual environment created at: $VENV_DIR"

# -----------------------------------------------------------------------------
# Step 6: Install Python Dependencies in venv
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 6/8] Installing Python dependencies in virtual environment..."

# Activate venv and install packages
source "$VENV_DIR/bin/activate"

# Upgrade pip
pip install --upgrade pip

# Install from requirements.txt if it exists
if [ -f "$PROJECT_ROOT/requirements.txt" ]; then
    pip install -r "$PROJECT_ROOT/requirements.txt"
else
    pip install pybullet pyinstaller
fi

echo "[INFO] Python packages installed in venv"

# -----------------------------------------------------------------------------
# Step 7: Create Activation Script
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 7/8] Creating activation script..."

ACTIVATE_SCRIPT="$PROJECT_ROOT/activate.sh"
cat > "$ACTIVATE_SCRIPT" << EOF
#!/bin/bash
# =============================================================================
# Drone Competition - Environment Activation Script
# =============================================================================
# This script activates both ROS 2 and the Python virtual environment.
#
# Usage:
#   source activate.sh
# =============================================================================

# Get the directory where this script is located
SCRIPT_DIR="\$(cd "\$(dirname "\${BASH_SOURCE[0]}")" && pwd)"

# Source ROS 2
source /opt/ros/${ROS_DISTRO}/setup.bash
echo "[OK] ROS 2 ${ROS_DISTRO} sourced"

# Activate Python virtual environment
source "\$SCRIPT_DIR/venv/bin/activate"
echo "[OK] Python venv activated"

echo ""
echo "Environment ready! You can now run:"
echo "  python simulation_host.py    # PyBullet"
echo "  python gazebo_bridge.py       # Gazebo"
echo "  python ros_bridge.py"
echo ""
EOF

chmod +x "$ACTIVATE_SCRIPT"
echo "[INFO] Created activation script: $ACTIVATE_SCRIPT"

# -----------------------------------------------------------------------------
# Verification
# -----------------------------------------------------------------------------
echo ""
echo "=============================================="
echo "  Installation Complete!"
echo "=============================================="
echo ""
echo "Verifying installation..."
echo ""

# Ensure we're in venv
source "$VENV_DIR/bin/activate"
source /opt/ros/${ROS_DISTRO}/setup.bash

echo -n "  ROS 2: "
ros2 --version 2>/dev/null && echo "" || echo "FAILED"

echo -n "  PyBullet: "
python3 -c "import pybullet; print('OK')" 2>/dev/null || echo "FAILED"

echo -n "  rclpy: "
python3 -c "import rclpy; print('OK')" 2>/dev/null || echo "FAILED"

echo -n "  geometry_msgs: "
python3 -c "from geometry_msgs.msg import Twist; print('OK')" 2>/dev/null || echo "FAILED"

echo -n "  std_msgs: "
python3 -c "from std_msgs.msg import String; print('OK')" 2>/dev/null || echo "FAILED"

echo -n "  PyInstaller: "
python3 -c "import PyInstaller; print('OK')" 2>/dev/null || echo "FAILED"

echo ""
echo "=============================================="
echo "  IMPORTANT: Activate the environment first!"
echo "=============================================="
echo ""
echo "Before running any scripts, activate the environment:"
echo "  source $ACTIVATE_SCRIPT"
echo ""
echo "Then run the simulation:"
echo "  python simulation_host.py"
echo ""