#!/bin/bash

# Rust Installation Script
# Installs Rust programming language using rustup
# Compatible with Ubuntu/Debian systems

set -e

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

log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] $1${NC}"
}

warn() {
    echo -e "${YELLOW}[$(date +'%Y-%m-%d %H:%M:%S')] WARNING: $1${NC}"
}

error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $1${NC}"
    exit 1
}

# Check if running as root
if [ "$EUID" -eq 0 ]; then
    warn "Running as root. Rust will be installed system-wide."
    INSTALL_DIR="/usr/local"
    CARGO_HOME="/usr/local/cargo"
    RUSTUP_HOME="/usr/local/rustup"
else
    log "Installing Rust for current user"
    INSTALL_DIR="$HOME"
    CARGO_HOME="$HOME/.cargo"
    RUSTUP_HOME="$HOME/.rustup"
fi

# Check if Rust is already installed
if command -v rustc &> /dev/null; then
    RUST_VERSION=$(rustc --version)
    warn "Rust is already installed: $RUST_VERSION"
    read -p "Do you want to update/reinstall? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        log "Installation cancelled"
        exit 0
    fi
fi

# Check for required dependencies
log "Checking system dependencies..."
MISSING_DEPS=()

# Check for curl
if ! command -v curl &> /dev/null; then
    MISSING_DEPS+=("curl")
fi

# Check for build essentials
if ! dpkg -l | grep -q build-essential; then
    MISSING_DEPS+=("build-essential")
fi

# Install missing dependencies
if [ ${#MISSING_DEPS[@]} -gt 0 ]; then
    log "Installing missing dependencies: ${MISSING_DEPS[*]}"
    if [ "$EUID" -eq 0 ]; then
        apt update
        apt install -y "${MISSING_DEPS[@]}"
    else
        log "Please install the following packages and run this script again:"
        log "sudo apt update && sudo apt install -y ${MISSING_DEPS[*]}"
        exit 1
    fi
fi

# Set environment variables for installation
export CARGO_HOME="$CARGO_HOME"
export RUSTUP_HOME="$RUSTUP_HOME"

log "Starting Rust installation..."
log "CARGO_HOME: $CARGO_HOME"
log "RUSTUP_HOME: $RUSTUP_HOME"

# Download and run rustup installer
log "Downloading rustup installer..."
if [ "$EUID" -eq 0 ]; then
    # For root installation, we need to handle it differently
    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --no-modify-path
    
    # Create symlinks for system-wide access
    ln -sf "$CARGO_HOME/bin/cargo" /usr/local/bin/cargo
    ln -sf "$CARGO_HOME/bin/rustc" /usr/local/bin/rustc
    ln -sf "$CARGO_HOME/bin/rustup" /usr/local/bin/rustup
    ln -sf "$CARGO_HOME/bin/rustdoc" /usr/local/bin/rustdoc
    ln -sf "$CARGO_HOME/bin/rust-gdb" /usr/local/bin/rust-gdb
    ln -sf "$CARGO_HOME/bin/rust-lldb" /usr/local/bin/rust-lldb
    
    # Set up environment for all users
    cat > /etc/profile.d/rust.sh << 'EOF'
export CARGO_HOME="/usr/local/cargo"
export RUSTUP_HOME="/usr/local/rustup"
export PATH="/usr/local/cargo/bin:$PATH"
EOF
    
    log "Rust installed system-wide. All users can now use Rust."
    log "Environment variables set in /etc/profile.d/rust.sh"
else
    # Standard user installation
    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
    
    # Source the cargo environment
    source "$CARGO_HOME/env"
    
    log "Rust installed for current user."
    log "To use Rust in new shells, run: source ~/.cargo/env"
    log "Or add ~/.cargo/bin to your PATH in your shell profile"
fi

# Verify installation
log "Verifying Rust installation..."
if [ "$EUID" -eq 0 ]; then
    RUST_VERSION=$(/usr/local/bin/rustc --version 2>/dev/null || echo "")
    CARGO_VERSION=$(/usr/local/bin/cargo --version 2>/dev/null || echo "")
else
    RUST_VERSION=$("$CARGO_HOME/bin/rustc" --version 2>/dev/null || echo "")
    CARGO_VERSION=$("$CARGO_HOME/bin/cargo" --version 2>/dev/null || echo "")
fi

if [ -n "$RUST_VERSION" ] && [ -n "$CARGO_VERSION" ]; then
    log "✅ Rust installation successful!"
    log "Rust compiler: $RUST_VERSION"
    log "Cargo package manager: $CARGO_VERSION"
    
    # Install common components
    log "Installing common Rust components..."
    if [ "$EUID" -eq 0 ]; then
        /usr/local/bin/rustup component add clippy rustfmt
    else
        "$CARGO_HOME/bin/rustup" component add clippy rustfmt
    fi
    
    log "✅ Rust installation completed successfully!"
    log ""
    log "Next steps:"
    if [ "$EUID" -eq 0 ]; then
        log "- Rust is available system-wide"
        log "- Users may need to log out and back in for PATH changes to take effect"
    else
        log "- Run 'source ~/.cargo/env' to use Rust in this session"
        log "- Restart your terminal or add ~/.cargo/bin to your PATH"
    fi
    log "- Create a new Rust project: cargo new my_project"
    log "- Build and run: cd my_project && cargo run"
else
    error "Rust installation failed. Please check the output above for errors."
fi