#!/bin/bash
# =============================================================================
# Drone Simulation - Ubuntu/Debian Installation Script
# =============================================================================
# Installs ROS 2, Gazebo, PyBullet, and all required dependencies
#
# Usage: ./install_ubuntu.sh
# =============================================================================

set -e

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

# Get 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
if [ -f /etc/os-release ]; then
    . /etc/os-release
    UBUNTU_VERSION=$VERSION_ID
    echo "[INFO] Detected: $NAME $VERSION_ID"
else
    echo "[WARN] Could not detect Ubuntu version, assuming 22.04"
    UBUNTU_VERSION="22.04"
fi

# -----------------------------------------------------------------------------
# Step 1: System Dependencies
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 1/8] Installing system dependencies..."

sudo apt-get update
sudo apt-get install -y \
    curl \
    gnupg \
    lsb-release \
    software-properties-common \
    python3 \
    python3-pip \
    python3-venv \
    git

echo "[INFO] System dependencies installed"

# -----------------------------------------------------------------------------
# Step 2: ROS 2 Repository Setup
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 2/8] Setting up ROS 2 repository..."

# Add ROS 2 GPG key
sudo curl -sSL https://raw.githubusercontent.com/ros/rosdistro/master/ros.key -o /usr/share/keyrings/ros-archive-keyring.gpg

# Determine ROS 2 distro based on Ubuntu version
if [ "$UBUNTU_VERSION" = "24.04" ]; then
    ROS_DISTRO="jazzy"
elif [ "$UBUNTU_VERSION" = "22.04" ]; then
    ROS_DISTRO="humble"
else
    echo "[WARN] Ubuntu $UBUNTU_VERSION not officially supported, trying humble"
    ROS_DISTRO="humble"
fi

echo "[INFO] Using ROS 2 $ROS_DISTRO"

# Add repository
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-get update
sudo apt-get install -y ros-${ROS_DISTRO}-ros-base ros-${ROS_DISTRO}-geometry-msgs ros-${ROS_DISTRO}-std-msgs ros-${ROS_DISTRO}-nav-msgs ros-${ROS_DISTRO}-sensor-msgs

echo "[INFO] ROS 2 $ROS_DISTRO installed"

# -----------------------------------------------------------------------------
# Step 4: Install Gazebo (optional)
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 4/8] Installing Gazebo..."

if [ "$ROS_DISTRO" = "jazzy" ]; then
    GZ_VERSION="harmonic"
else
    GZ_VERSION="fortress"
fi

sudo apt-get install -y ros-${ROS_DISTRO}-ros-gz || {
    echo "[WARN] Could not install ros-gz, Gazebo simulation will not be available"
    echo "[INFO] PyBullet simulation will still work"
}

echo "[INFO] Gazebo installation complete"

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

# Remove existing venv if present
if [ -d "$VENV_DIR" ]; then
    rm -rf "$VENV_DIR"
fi

# Create virtual environment
python3 -m venv "$VENV_DIR"

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

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

source "$VENV_DIR/bin/activate"
pip install --upgrade pip

if [ -f "$PROJECT_ROOT/requirements.txt" ]; then
    echo "[INFO] Installing from requirements.txt..."
    pip install -r "$PROJECT_ROOT/requirements.txt"
else
    echo "[INFO] Installing packages manually..."
    pip install pybullet numpy pillow pyinstaller
fi

echo "[INFO] Python packages installed"

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

cat > "$PROJECT_ROOT/activate.sh" << 'EOF'
#!/bin/bash
# Drone Simulation - Environment Activation

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Source ROS 2
if [ -f "/opt/ros/jazzy/setup.bash" ]; then
    source /opt/ros/jazzy/setup.bash
    echo "[OK] ROS 2 jazzy sourced"
elif [ -f "/opt/ros/humble/setup.bash" ]; then
    source /opt/ros/humble/setup.bash
    echo "[OK] ROS 2 humble sourced"
else
    echo "[WARN] ROS 2 not found - standalone_simulation.py will work"
fi

# Activate Python venv
if [ -f "$SCRIPT_DIR/venv/bin/activate" ]; then
    source "$SCRIPT_DIR/venv/bin/activate"
    echo "[OK] Python venv activated"
fi

# Set Gazebo model path
export GZ_SIM_RESOURCE_PATH="$SCRIPT_DIR/gazebo/models:$GZ_SIM_RESOURCE_PATH"

echo ""
echo "Environment ready! Run one of:"
echo "  python standalone_simulation.py  (No ROS 2 required)"
echo "  python simulation_host.py        (With ROS 2)"
echo ""
EOF

chmod +x "$PROJECT_ROOT/activate.sh"
echo "[INFO] Created: $PROJECT_ROOT/activate.sh"

# -----------------------------------------------------------------------------
# Step 8: Verification
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 8/8] Verifying installation..."

source "$PROJECT_ROOT/activate.sh"

echo ""
echo "Checking Python packages:"
python3 -c "import pybullet; print('  PyBullet: OK')" || echo "  PyBullet: FAILED"
python3 -c "import numpy; print('  NumPy: OK')" || echo "  NumPy: FAILED"
python3 -c "from PIL import Image; print('  Pillow: OK')" || echo "  Pillow: FAILED"

echo ""
echo "=============================================="
echo "  Installation Complete!"
echo "=============================================="
echo ""
echo "Quick start:"
echo "  source activate.sh"
echo "  python standalone_simulation.py"
echo ""
echo "Or with ROS 2:"
echo "  python simulation_host.py     # Terminal 1"
echo "  python ros_bridge.py          # Terminal 2"
echo "  python controllers.py         # Terminal 3"
echo ""