#!/bin/bash

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

this_script_url="https://example.com/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 dependencies
install_dependencies() {
    local install_output
    local update_output
    
    # Update package lists with error handling
    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

    # Check if packages are already installed
    local missing_packages=()
    local packages=("ca-certificates" "curl" "gnupg" "apt-transport-https" "software-properties-common")
    
    for pkg in "${packages[@]}"; do
        if ! dpkg -l | grep -q "^ii  $pkg "; then
            missing_packages+=("$pkg")
        fi
    done

    # If all packages are installed, return success
    if [ ${#missing_packages[@]} -eq 0 ]; then
        echo -e "All required packages are already installed $CHECK_MARK"
        return 0
    fi

    # Install missing packages
    echo "Installing missing packages: ${missing_packages[*]}"
    install_output=$(sudo apt-get install -y "${missing_packages[@]}" 2>&1)
    if [ $? -ne 0 ]; then
        echo -e "Failed to install dependencies $CROSS_MARK"
        echo "Error output: $install_output"
        return 1
    fi

    # Verify installation
    local failed_packages=()
    for pkg in "${missing_packages[@]}"; do
        if ! dpkg -l | grep -q "^ii  $pkg "; then
            failed_packages+=("$pkg")
        fi
    done

    if [ ${#failed_packages[@]} -gt 0 ]; then
        echo -e "Failed to install some packages: ${failed_packages[*]} $CROSS_MARK"
        return 1
    fi

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

# Function to setup Docker repository
setup_docker_repo() {
    # Create directory for keyrings if it doesn't exist
    sudo install -m 0755 -d /etc/apt/keyrings

    # Download Docker's official GPG key with error handling
    if ! curl -fsSL "https://download.docker.com/linux/$ID/gpg" -o /tmp/docker.gpg; then
        echo -e "Failed to download Docker GPG key $CROSS_MARK"
        return 1
    fi

    # Dearmor and store the GPG key with error handling
    if ! sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg /tmp/docker.gpg; then
        echo -e "Failed to dearmor GPG key $CROSS_MARK"
        rm -f /tmp/docker.gpg
        return 1
    fi
    rm -f /tmp/docker.gpg

    # Set proper permissions
    sudo chmod a+r /etc/apt/keyrings/docker.gpg

    # Remove any existing Docker repository file
    sudo rm -f /etc/apt/sources.list.d/docker.list

    # Create the repository configuration with error handling
    local repo_config="deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/${ID} ${VERSION_CODENAME} stable"
    
    if ! echo "$repo_config" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null; then
        echo -e "Failed to create repository configuration $CROSS_MARK"
        return 1
    fi

    # Display the repository configuration for verification
    echo "Created repository configuration:"
    cat /etc/apt/sources.list.d/docker.list

    # Verify the repository file exists and has content
    if [ ! -s /etc/apt/sources.list.d/docker.list ]; then
        echo -e "Repository configuration file is empty or missing $CROSS_MARK"
        return 1
    fi

    echo -e "Docker repository setup completed $CHECK_MARK"
    return 0
}

# Function to install Docker
install_docker() {
    sudo apt-get update > /dev/null 2>&1
    sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin > /dev/null 2>&1
    
    if [ $? -eq 0 ]; then
        echo -e "Docker installation completed $CHECK_MARK"
        return 0
    else
        echo -e "Docker installation failed $CROSS_MARK"
        return 1
    fi
}

# Function to verify Docker installation
verify_docker() {
    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"
        return 0
    else
        echo -e "Failed to add user to docker group $CROSS_MARK"
        return 1
    fi
}

# Function to clean up in case of failure
cleanup_repo() {
    echo "Cleaning up repository files..."
    sudo rm -f /etc/apt/sources.list.d/docker.list
    sudo rm -f /etc/apt/keyrings/docker.gpg
    echo -e "Repository cleanup completed $CHECK_MARK"
}

# 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_dependencies "Installing Dependencies"; then
    success=1
    exit $success
fi

if ! format_output setup_docker_repo "Setting up Docker Repository"; then
    format_output cleanup_repo "Cleaning up failed repository setup"
    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