#!/bin/bash

#-----------------------------------#
#             VARIABLES             #
#-----------------------------------#

this_script_url="https://git.hhf.technology/hhf/script-management-cloudpanel/raw/branch/main/install/docker-install.sh"
this_script_name="Docker Installation Script"
formatter_url="https://git.hhf.technology/hhf/TaskFormatter/raw/branch/main/bash_task_formatter/task_formatter.sh"
scriptname=$0

# Initialize success flag
success=0

# Determine the user (use the first argument if provided, otherwise fallback)
USER_TO_RUN_AS="${1:-$SUDO_USER}"
USER_HOME=$(eval echo ~$USER_TO_RUN_AS)

#-----------------------------------#
#             FORMATTER             #
#-----------------------------------#

# Download and source the formatter with error handling
download_formatter() {
    if [ ! -f "task_formatter.sh" ]; then
        if ! wget "$formatter_url" --no-check-certificate -O task_formatter.sh > /dev/null 2>&1; then
            echo "Error: Failed to download task_formatter.sh"
            exit 1
        fi
    fi

    if [ ! -f "task_formatter.sh" ]; then
        echo "Error: task_formatter.sh not found after download attempt"
        exit 1
    fi

    # Make the formatter executable
    chmod +x task_formatter.sh

    # Source the formatter with error checking
    source ./task_formatter.sh || { echo "Error: Failed to source task_formatter.sh"; exit 1; }
    
    # Check if print_header is available after sourcing
    if ! declare -f print_header > /dev/null; then
        echo "Error: print_header function not found after sourcing."
        exit 1
    fi
}

# Call the download_formatter function
download_formatter

#-----------------------------------#
#             FUNCTIONS             #
#-----------------------------------#

# Function to check system architecture
check_architecture() {
    local arch=$(uname -m)
    if [[ "$arch" != "x86_64" && "$arch" != "aarch64" ]]; then
        echo -e "Unsupported architecture: $arch $CROSS_MARK"
        return 1
    fi
    echo -e "Architecture $arch is supported $CHECK_MARK"
    return 0
}

# Function to check OS version
check_os_version() {
    if [ ! -f /etc/os-release ]; then
        echo -e "Cannot determine OS version $CROSS_MARK"
        return 1
    fi

    source /etc/os-release
    
    case "$ID $VERSION_ID" in
        "debian 12"|"ubuntu 22.04"|"ubuntu 24.04")
            echo -e "OS version $ID $VERSION_ID is supported $CHECK_MARK"
            return 0
            ;;
        *)
            echo -e "Unsupported OS version: $ID $VERSION_ID $CROSS_MARK"
            return 1
            ;;
    esac
}

# Function to install Docker
install_docker() {
    local update_output
    local install_output
    
    # Update package lists
    echo "Updating package lists..."
    update_output=$(sudo apt-get update 2>&1)
    if [ $? -ne 0 ]; then
        echo -e "Failed to update package lists $CROSS_MARK"
        echo "Error output: $update_output"
        return 1
    fi

    # Install Docker
    echo "Installing Docker..."
    install_output=$(sudo apt-get install -y docker.io 2>&1)
    if [ $? -ne 0 ]; then
        echo -e "Failed to install Docker $CROSS_MARK"
        echo "Error output: $install_output"
        return 1
    fi

    echo -e "Docker installed successfully $CHECK_MARK"
    return 0
}

# Function to verify Docker installation
verify_docker() {
    # Check Docker service status
    if ! systemctl is-active --quiet docker; then
        echo "Starting Docker service..."
        sudo systemctl start docker
    fi

    if sudo docker run hello-world > /dev/null 2>&1; then
        echo -e "Docker verification successful $CHECK_MARK"
        return 0
    else
        echo -e "Docker verification failed $CROSS_MARK"
        return 1
    fi
}

# Function to add user to docker group
setup_docker_user() {
    sudo usermod -aG docker $USER_TO_RUN_AS
    if [ $? -eq 0 ]; then
        echo -e "User added to docker group $CHECK_MARK"
        echo "Note: You may need to log out and back in for group changes to take effect"
        return 0
    else
        echo -e "Failed to add user to docker group $CROSS_MARK"
        return 1
    fi
}

# Remove the script itself
remove_script() {
    if [ -f "$0" ]; then
        echo "Deleted docker installation script..."
        rm -- "$0"
    fi
    if [ -f "task_formatter.sh" ]; then
        rm task_formatter.sh
    fi
    echo -e "Cleaned up $CHECK_MARK"
    return 0
}

#-----------------------------------#
#             MAIN LOGIC            #
#-----------------------------------#

# Check if print_header function exists
if ! command -v print_header >/dev/null 2>&1; then
    echo "Error: print_header function not found. Formatter may not be properly sourced."
    exit 1
fi

# Print header
print_header "$this_script_name" "$this_script_url"

echo -e "Running as User: $USER_TO_RUN_AS\nUser Home: $USER_HOME\n"

# Run the installation functions with formatted output
if ! format_output check_architecture "Checking System Architecture"; then
    success=1
    exit $success
fi

if ! format_output check_os_version "Checking OS Version"; then
    success=1
    exit $success
fi

if ! format_output install_docker "Installing Docker"; then
    success=1
    exit $success
fi

if ! format_output verify_docker "Verifying Docker Installation"; then
    success=1
    exit $success
fi

if ! format_output setup_docker_user "Setting up Docker User"; then
    success=1
    exit $success
fi

format_output remove_script "Removing script"

# Print final message
final_message "$this_script_name" "$success"

# Exit with appropriate status
exit $success