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

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

# Check for ArduPilot option
INSTALL_ARDUPILOT=false
for arg in "$@"; do
    if [ "$arg" = "--with-ardupilot" ]; then
        INSTALL_ARDUPILOT=true
    fi
done

# -----------------------------------------------------------------------------
# Step 1: System Dependencies
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 1/7] 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 \
    cmake \
    build-essential \
    wget

echo "[INFO] System dependencies installed"

# -----------------------------------------------------------------------------
# Step 2: ROS 2 Repository Setup
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 2/7] 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/7] 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
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 4/7] Installing Gazebo..."

# Install ros-gz bridge
sudo apt-get install -y ros-${ROS_DISTRO}-ros-gz || {
    echo "[WARN] Could not install ros-gz"
}

# Install Gazebo
if [ "$ROS_DISTRO" = "jazzy" ]; then
    sudo apt-get install -y gz-harmonic || true
else
    sudo apt-get install -y gz-fortress || sudo apt-get install -y ros-${ROS_DISTRO}-ros-ign-gazebo || true
fi

# Verify installation
if command -v gz &> /dev/null; then
    echo "[INFO] Gazebo installed (gz command available)"
elif command -v ign &> /dev/null; then
    echo "[INFO] Gazebo Fortress installed (ign command available)"
else
    echo "[WARN] Gazebo command not found - use standalone mode"
fi

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

if [ -d "$VENV_DIR" ]; then
    rm -rf "$VENV_DIR"
fi

python3 -m venv "$VENV_DIR"
echo "[INFO] Virtual environment created at: $VENV_DIR"

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

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

if [ -f "$PROJECT_ROOT/requirements.txt" ]; then
    pip install -r "$PROJECT_ROOT/requirements.txt"
else
    pip install pybullet numpy pillow opencv-python pymavlink
fi

echo "[INFO] Python packages installed"

# -----------------------------------------------------------------------------
# Step 7: Create Activation Script
# -----------------------------------------------------------------------------
echo ""
echo "[STEP 7/7] 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"
elif [ -f "/opt/ros/humble/setup.bash" ]; then
    source /opt/ros/humble/setup.bash
    echo "[OK] ROS 2 humble"
fi

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

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

# Add ~/.local/bin to PATH (for mavproxy.py)
export PATH="$PATH:$HOME/.local/bin"

# ArduPilot workspace
if [ -f "$HOME/ardu_ws/install/setup.bash" ]; then
    source "$HOME/ardu_ws/install/setup.bash"
    echo "[OK] ArduPilot workspace"
fi

echo ""
echo "Ready! Run: python standalone_simulation.py"
EOF

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

# -----------------------------------------------------------------------------
# Verification
# -----------------------------------------------------------------------------
echo ""
echo "=============================================="
echo "  Verifying Installation"
echo "=============================================="

source "$PROJECT_ROOT/activate.sh"

python3 -c "import pybullet; print('[OK] PyBullet')" || echo "[FAIL] PyBullet"
python3 -c "import numpy; print('[OK] NumPy')" || echo "[FAIL] NumPy"
python3 -c "import cv2; print('[OK] OpenCV')" || echo "[WARN] OpenCV not installed"
python3 -c "from pymavlink import mavutil; print('[OK] pymavlink')" || echo "[WARN] pymavlink not installed"

# -----------------------------------------------------------------------------
# ArduPilot Installation (if requested)
# -----------------------------------------------------------------------------
if [ "$INSTALL_ARDUPILOT" = true ]; then
    echo ""
    echo "[INFO] Installing ArduPilot SITL..."
    
    if [ -f "$SCRIPT_DIR/install_ardupilot.sh" ]; then
        bash "$SCRIPT_DIR/install_ardupilot.sh"
    else
        echo "[ERROR] install_ardupilot.sh not found!"
    fi
fi

echo ""
echo "=============================================="
echo "  Installation Complete!"
echo "=============================================="
echo ""
echo "Quick start:"
echo "  source activate.sh"
echo "  python standalone_simulation.py"
echo ""
echo "With Gazebo (2 terminals):"
echo "  Terminal 1: ros2 launch gazebo/launch/drone_landing.launch.py"
echo "  Terminal 2: python run_gazebo.py --pattern circular"
echo ""
if [ "$INSTALL_ARDUPILOT" != true ]; then
    echo "For ArduPilot SITL:"
    echo "  ./setup/install_ardupilot.sh"
    echo ""
fi