#!/bin/bash

# Ubuntu User Setup Script
# Creates user account, sets up SSH keys from GitHub, and optionally grants sudo privileges
# Must be run as root

set -e  # Exit on any error

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

# Function to print colored output
print_status() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

print_question() {
    echo -e "${BLUE}[QUESTION]${NC} $1"
}

# Check if running as root
if [[ $EUID -ne 0 ]]; then
   print_error "This script must be run as root"
   exit 1
fi

# Check if required tools are available
if ! command -v curl &> /dev/null; then
    print_error "curl is required but not installed. Please install it first: apt update && apt install curl"
    exit 1
fi

print_status "Starting Ubuntu User Setup Script"
echo "=================================="
echo

# Step 1: Ask for username
while true; do
    print_question "Enter the username for the new user account:"
    read -r USERNAME < /dev/tty
    
    echo "DEBUG: Received input: '$USERNAME'" >&2
    
    if [[ -z "$USERNAME" ]]; then
        print_warning "Username cannot be empty. Please try again."
        continue
    fi
    
    if [[ ! "$USERNAME" =~ ^[a-z_][a-z0-9_-]*$ ]]; then
        print_warning "Invalid username format. Please use lowercase letters, numbers, hyphens and underscores only."
        continue
    fi
    
    break
done

# Step 2: Create user account if it doesn't exist
if id "$USERNAME" &>/dev/null; then
    print_warning "User '$USERNAME' already exists. Skipping user creation."
else
    print_status "Creating user account '$USERNAME'..."
    useradd -m -s /bin/bash "$USERNAME"
    print_status "User '$USERNAME' created successfully with home directory."
fi

# Get user's home directory
USER_HOME=$(eval echo "~$USERNAME")

# Step 3: Ask for GitHub username
while true; do
    print_question "Enter the GitHub username to fetch SSH keys from:"
    read -r GITHUB_USER < /dev/tty
    
    if [[ -z "$GITHUB_USER" ]]; then
        print_warning "GitHub username cannot be empty. Please try again."
        continue
    fi
    
    # Validate GitHub user exists
    print_status "Validating GitHub user '$GITHUB_USER'..."
    HTTP_STATUS=$(curl -s -o /dev/null -w "%{http_code}" "https://api.github.com/users/$GITHUB_USER")
    
    if [[ "$HTTP_STATUS" == "200" ]]; then
        print_status "GitHub user '$GITHUB_USER' found."
        break
    elif [[ "$HTTP_STATUS" == "404" ]]; then
        print_warning "GitHub user '$GITHUB_USER' not found. Please check the username and try again."
        continue
    else
        print_warning "Unable to validate GitHub user (HTTP $HTTP_STATUS). Please check your internet connection and try again."
        continue
    fi
done

# Step 4: Download and setup SSH keys
print_status "Downloading SSH keys for GitHub user '$GITHUB_USER'..."

# Create .ssh directory with correct permissions
SSH_DIR="$USER_HOME/.ssh"
mkdir -p "$SSH_DIR"
chmod 700 "$SSH_DIR"

# Download SSH keys from GitHub
KEYS_URL="https://github.com/$GITHUB_USER.keys"
AUTHORIZED_KEYS_FILE="$SSH_DIR/authorized_keys"

# Download keys and check if any were found
KEYS_CONTENT=$(curl -s "$KEYS_URL")

if [[ -z "$KEYS_CONTENT" ]] || [[ "$KEYS_CONTENT" == "Not Found" ]]; then
    print_warning "No public SSH keys found for GitHub user '$GITHUB_USER'."
    print_warning "The user will need to add SSH keys manually or use password authentication."
    # Create empty authorized_keys file
    touch "$AUTHORIZED_KEYS_FILE"
else
    # Write keys to authorized_keys file (overwrites existing content)
    echo "$KEYS_CONTENT" > "$AUTHORIZED_KEYS_FILE"
    KEY_COUNT=$(echo "$KEYS_CONTENT" | wc -l)
    print_status "Successfully downloaded and installed $KEY_COUNT SSH key(s) for user '$USERNAME'."
fi

# Set correct permissions for SSH files
chmod 600 "$AUTHORIZED_KEYS_FILE"
chown -R "$USERNAME:$USERNAME" "$SSH_DIR"

# Step 5: Ask about sudo privileges
echo
while true; do
    print_question "Should user '$USERNAME' get sudo privileges? (y/N):"
    read -r GRANT_SUDO < /dev/tty
    
    case $GRANT_SUDO in
        [Yy]|[Yy][Ee][Ss])
            # Add user to sudo group
            print_status "Adding user '$USERNAME' to sudo group..."
            usermod -aG sudo "$USERNAME"
            
            # Create sudoers file
            SUDOERS_FILE="/etc/sudoers.d/$USERNAME"
            print_status "Creating sudoers configuration file..."
            echo "$USERNAME ALL=(ALL:ALL) NOPASSWD: ALL" > "$SUDOERS_FILE"
            chmod 440 "$SUDOERS_FILE"
            
            # Validate sudoers file
            if visudo -c -f "$SUDOERS_FILE"; then
                print_status "User '$USERNAME' granted sudo privileges with NOPASSWD."
            else
                print_error "Error in sudoers file configuration. Removing file..."
                rm -f "$SUDOERS_FILE"
                exit 1
            fi
            break
            ;;
        [Nn]|[Nn][Oo]|"")
            print_status "User '$USERNAME' will not have sudo privileges."
            break
            ;;
        *)
            print_warning "Please answer 'y' for yes or 'n' for no."
            ;;
    esac
done

# Step 6: Summary
echo
print_status "Setup completed successfully!"
echo "=================================="
echo "Summary:"
echo "- User: $USERNAME"
echo "- Home directory: $USER_HOME"
echo "- SSH directory: $SSH_DIR (permissions: 700)"
echo "- SSH keys source: GitHub user '$GITHUB_USER'"
if [[ -n "$KEYS_CONTENT" ]] && [[ "$KEYS_CONTENT" != "Not Found" ]]; then
    echo "- SSH keys: $(echo "$KEYS_CONTENT" | wc -l) key(s) installed"
else
    echo "- SSH keys: No keys found (manual setup required)"
fi
echo "- Sudo privileges: $(if [[ "$GRANT_SUDO" =~ ^[Yy] ]]; then echo "Yes (NOPASSWD)"; else echo "No"; fi)"
echo
print_status "The user '$USERNAME' can now connect via SSH using their GitHub SSH keys."

if [[ "$GRANT_SUDO" =~ ^[Yy] ]]; then
    print_status "The user has passwordless sudo access."
fi

echo
print_status "Setup script finished."