#!/bin/bash
# =============================================================================
# Drone Simulation - Ubuntu/Debian Installation Script
# =============================================================================
# Installs ROS 2, Gazebo, PyBullet, and all required dependencies
# Includes optional ArduPilot SITL setup for realistic flight controller
#
# Usage: ./install_ubuntu.sh [--with-ardupilot]
# =============================================================================

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 \
    cmake \
    build-essential

echo "[INFO] System dependencies installed"

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

# -----------------------------------------------------------------------------
# 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"
    GZ_PKG="gz-harmonic"
else
    GZ_VERSION="fortress"
    GZ_PKG="gz-fortress"
fi

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

# Install Gazebo itself (provides gz or ign command)
sudo apt-get install -y $GZ_PKG || {
    echo "[WARN] Could not install $GZ_PKG"
    echo "[INFO] Trying ros-ign-gazebo..."
    sudo apt-get install -y ros-${ROS_DISTRO}-ros-ign-gazebo || true
}

# 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 PyBullet instead"
fi

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"

# Set ArduPilot paths if installed
if [ -d "$HOME/ardupilot" ]; then
    export ARDUPILOT_HOME="$HOME/ardupilot"
    export PATH="$PATH:$ARDUPILOT_HOME/Tools/autotest"
    echo "[OK] ArduPilot SITL available"
fi

# Set ArduPilot Gazebo plugin path if installed
if [ -d "$HOME/ardupilot_gazebo/build" ]; then
    export GZ_SIM_SYSTEM_PLUGIN_PATH="$HOME/ardupilot_gazebo/build:$GZ_SIM_SYSTEM_PLUGIN_PATH"
    export GZ_SIM_RESOURCE_PATH="$HOME/ardupilot_gazebo/models:$HOME/ardupilot_gazebo/worlds:$GZ_SIM_RESOURCE_PATH"
    echo "[OK] ArduPilot Gazebo plugin available"
fi

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 "  python run_ardupilot.py          (With ArduPilot SITL)"
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"
python3 -c "from pymavlink import mavutil; print('  pymavlink: OK')" || echo "  pymavlink: FAILED"

# -----------------------------------------------------------------------------
# Step 9: Optional ArduPilot SITL Installation
# -----------------------------------------------------------------------------
if [ "$INSTALL_ARDUPILOT" = true ]; then
    echo ""
    echo "[STEP 9] Installing ArduPilot SITL..."
    
    # Install ArduPilot dependencies
    sudo apt-get install -y \
        python3-dev \
        python3-opencv \
        python3-wxgtk4.0 \
        python3-matplotlib \
        python3-lxml \
        libxml2-dev \
        libxslt1-dev || true
    
    # Clone ArduPilot if not exists
    if [ ! -d "$HOME/ardupilot" ]; then
        echo "[INFO] Cloning ArduPilot..."
        git clone --recurse-submodules https://github.com/ArduPilot/ardupilot.git "$HOME/ardupilot"
        cd "$HOME/ardupilot"
        Tools/environment_install/install-prereqs-ubuntu.sh -y
        cd "$PROJECT_ROOT"
    else
        echo "[INFO] ArduPilot already exists at $HOME/ardupilot"
    fi
    
    # Clone ArduPilot Gazebo plugin if not exists
    if [ ! -d "$HOME/ardupilot_gazebo" ]; then
        echo "[INFO] Cloning ArduPilot Gazebo plugin..."
        git clone https://github.com/ArduPilot/ardupilot_gazebo.git "$HOME/ardupilot_gazebo"
        cd "$HOME/ardupilot_gazebo"
        mkdir -p build && cd build
        cmake .. -DCMAKE_BUILD_TYPE=Release
        make -j$(nproc)
        cd "$PROJECT_ROOT"
    else
        echo "[INFO] ArduPilot Gazebo plugin already exists at $HOME/ardupilot_gazebo"
    fi
    
    # Add to bashrc
    if ! grep -q "ARDUPILOT_HOME" ~/.bashrc; then
        echo '' >> ~/.bashrc
        echo '# ArduPilot SITL' >> ~/.bashrc
        echo 'export ARDUPILOT_HOME=$HOME/ardupilot' >> ~/.bashrc
        echo 'export PATH=$PATH:$ARDUPILOT_HOME/Tools/autotest' >> ~/.bashrc
    fi
    
    if ! grep -q "ardupilot_gazebo" ~/.bashrc; then
        echo '' >> ~/.bashrc
        echo '# ArduPilot Gazebo Plugin' >> ~/.bashrc
        echo 'export GZ_SIM_SYSTEM_PLUGIN_PATH=$HOME/ardupilot_gazebo/build:$GZ_SIM_SYSTEM_PLUGIN_PATH' >> ~/.bashrc
        echo 'export GZ_SIM_RESOURCE_PATH=$HOME/ardupilot_gazebo/models:$HOME/ardupilot_gazebo/worlds:$GZ_SIM_RESOURCE_PATH' >> ~/.bashrc
    fi
    
    echo "[INFO] ArduPilot SITL installed"
fi

echo ""
echo "=============================================="
echo "  Installation Complete!"
echo "=============================================="
echo ""
echo "Quick start:"
echo "  source activate.sh"
echo "  python standalone_simulation.py"
echo ""
echo "With ROS 2 + Gazebo:"
echo "  python simulation_host.py     # Terminal 1"
echo "  python run_bridge.py          # Terminal 2"
echo ""
echo "With ArduPilot SITL (realistic flight controller):"
echo "  ros2 launch gazebo/launch/ardupilot_drone.launch.py  # Terminal 1"
echo "  sim_vehicle.py -v ArduCopter -f gazebo-iris --model JSON  # Terminal 2"
echo "  python run_ardupilot.py --no-sitl                    # Terminal 3"
echo ""
if [ "$INSTALL_ARDUPILOT" != true ]; then
    echo "To install ArduPilot SITL, run:"
    echo "  ./setup/install_ubuntu.sh --with-ardupilot"
    echo ""
fi