#!/bin/bash

# ROS 2 Auto-Install Script for Ubuntu
# Supports Ubuntu 22.04 (Humble) and 24.04 (Jazzy)

# Check if running with bash
if [ -z "$BASH_VERSION" ]; then
    echo "Error: This script must be run with bash"
    echo "Usage: bash autoinstall_ros2.sh"
    echo "   or: ./autoinstall_ros2.sh"
    exit 1
fi

set -e  # Exit on any error

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

echo -e "${GREEN}================================${NC}"
echo -e "${GREEN}ROS 2 Auto-Installer${NC}"
echo -e "${GREEN}================================${NC}\n"

# Check if running as root
if [ "$EUID" -eq 0 ]; then 
    echo -e "${YELLOW}Warning: Running as root. This is OK but not recommended for ROS development.${NC}\n"
    SUDO=""
else
    SUDO="sudo"
fi

# Detect Ubuntu version
if [ -f /etc/os-release ]; then
    . /etc/os-release
    UBUNTU_VERSION=$VERSION_ID
    UBUNTU_CODENAME=$UBUNTU_CODENAME
else
    echo -e "${RED}Cannot detect Ubuntu version!${NC}"
    exit 1
fi

# Determine ROS 2 distro based on Ubuntu version
case $UBUNTU_VERSION in
    "22.04")
        ROS_DISTRO="humble"
        echo -e "${GREEN}Detected Ubuntu 22.04 - Installing ROS 2 Humble${NC}\n"
        ;;
    "24.04")
        ROS_DISTRO="jazzy"
        echo -e "${GREEN}Detected Ubuntu 24.04 - Installing ROS 2 Jazzy${NC}\n"
        ;;
    "20.04")
        ROS_DISTRO="foxy"
        echo -e "${YELLOW}Ubuntu 20.04 detected - ROS 2 Foxy is EOL. Consider upgrading Ubuntu.${NC}\n"
        ;;
    *)
        echo -e "${RED}Unsupported Ubuntu version: $UBUNTU_VERSION${NC}"
        echo "This script supports Ubuntu 22.04 (Humble) and 24.04 (Jazzy)"
        exit 1
        ;;
esac

# Ask for installation type
echo "Select installation type:"
echo "1) ros-base (minimal, no GUI tools - recommended for headless)"
echo "2) ros-core (bare minimum)"
echo "3) desktop (full with GUI tools like RViz)"
read -p "Enter choice [1-3] (default: 1): " INSTALL_TYPE
INSTALL_TYPE=${INSTALL_TYPE:-1}

case $INSTALL_TYPE in
    1)
        ROS_PACKAGE="ros-${ROS_DISTRO}-ros-base"
        echo -e "${GREEN}Installing ros-base (headless)${NC}\n"
        ;;
    2)
        ROS_PACKAGE="ros-${ROS_DISTRO}-ros-core"
        echo -e "${GREEN}Installing ros-core (minimal)${NC}\n"
        ;;
    3)
        ROS_PACKAGE="ros-${ROS_DISTRO}-desktop"
        echo -e "${GREEN}Installing desktop (full with GUI)${NC}\n"
        ;;
    *)
        echo -e "${RED}Invalid choice. Exiting.${NC}"
        exit 1
        ;;
esac

echo -e "${YELLOW}Starting installation...${NC}\n"

# Step 1: Set locale
echo -e "${GREEN}[1/8] Setting up locale...${NC}"
$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: Setup sources
echo -e "${GREEN}[2/8] Adding universe repository...${NC}"
$SUDO apt install -y software-properties-common
$SUDO add-apt-repository universe -y

# Step 3: Add ROS 2 GPG key
echo -e "${GREEN}[3/8] Adding ROS 2 GPG key...${NC}"
$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

# Step 4: Add repository
echo -e "${GREEN}[4/8] Adding ROS 2 repository...${NC}"
echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/ros-archive-keyring.gpg] http://packages.ros.org/ros2/ubuntu $UBUNTU_CODENAME main" | $SUDO tee /etc/apt/sources.list.d/ros2.list > /dev/null

# Step 5: Update package cache
echo -e "${GREEN}[5/8] Updating package cache...${NC}"
$SUDO apt update

# Step 6: Install ROS 2
echo -e "${GREEN}[6/8] Installing ROS 2 ${ROS_DISTRO}... (this may take a while)${NC}"
$SUDO apt install -y $ROS_PACKAGE

# Step 7: Install development tools
echo -e "${GREEN}[7/8] Installing development tools...${NC}"
$SUDO apt install -y ros-dev-tools python3-colcon-common-extensions

# Step 8: Setup environment
echo -e "${GREEN}[8/8] Setting up environment...${NC}"

# Determine the correct shell config file
if [ -n "$BASH_VERSION" ]; then
    SHELL_CONFIG="$HOME/.bashrc"
elif [ -n "$ZSH_VERSION" ]; then
    SHELL_CONFIG="$HOME/.zshrc"
else
    SHELL_CONFIG="$HOME/.bashrc"
fi

# Add ROS 2 sourcing to shell config if not already present
if ! grep -q "source /opt/ros/${ROS_DISTRO}/setup.bash" "$SHELL_CONFIG" 2>/dev/null; then
    echo "" >> "$SHELL_CONFIG"
    echo "# ROS 2 ${ROS_DISTRO}" >> "$SHELL_CONFIG"
    echo "source /opt/ros/${ROS_DISTRO}/setup.bash" >> "$SHELL_CONFIG"
    echo -e "${GREEN}Added ROS 2 sourcing to $SHELL_CONFIG${NC}"
else
    echo -e "${YELLOW}ROS 2 already sourced in $SHELL_CONFIG${NC}"
fi

# Source for current session (use . for better compatibility)
. /opt/ros/${ROS_DISTRO}/setup.bash

# Verification
echo -e "\n${GREEN}================================${NC}"
echo -e "${GREEN}Installation Complete!${NC}"
echo -e "${GREEN}================================${NC}\n"

echo "Verifying installation..."
if command -v ros2 &> /dev/null; then
    echo -e "${GREEN}✓ ROS 2 command found${NC}"
    echo -e "${GREEN}✓ ROS 2 Distro: ${ROS_DISTRO}${NC}"
    
    # Check if environment is properly set
    if [ -n "$ROS_DISTRO" ]; then
        echo -e "${GREEN}✓ ROS environment configured${NC}"
    fi
    
    echo -e "\n${GREEN}Installation successful!${NC}"
    echo -e "\nTo start using ROS 2:"
    echo "  1. Open a new terminal (or run: source $SHELL_CONFIG)"
    echo "  2. Try: ros2 topic list"
    echo "  3. Install turtlesim: sudo apt install ros-${ROS_DISTRO}-turtlesim"
    echo "  4. Run turtlesim: ros2 run turtlesim turtlesim_node"
    echo -e "\nUseful commands:"
    echo "  ros2 --help          - Show all ROS 2 commands"
    echo "  ros2 node list       - List running nodes"
    echo "  ros2 topic list      - List active topics"
    echo -e "\nHappy robot building! 🤖"
else
    echo -e "${RED}✗ ROS 2 command not found. Something went wrong.${NC}"
    exit 1
fi