#!/bin/bash

# This script is a initialization helper for container instances
# on AutoDL based on my flavour. It does:
#    - Set up locale to avoid garbled text display.
#    - Set up APT mirrors and install packages, clean cache when done.
#    - Install `uv`, a powerful replacement for `pip`.
#    - Clean up bash history.
#
# Some behavior can be suppressed by set/unset variables in script.
#
# Also, there are some tips for your customization and simplifying server configuring.
#    - Add your SSH public keys. (Also can be done in AutoDL portal.)
#    - Add your custom hook function in this bash, do anything you like.
#    - Write files to your destination path.
#
# Usage:
#    Download the script, then:
#        1. Review the custom options below.
#        2. Add your function hooks.
#        3. Write files to your destinations.
#        4. Run the script by `bash autodl-initializer.sh`
#        5. Restart/relogin to your shell, and have a rock!
# 




###############################################################
# Custom options
#
# Make your customization here!
###############################################################

# Set the desired system locale.
#    Example: 'en_US.UTF-8', 'zh_CN.UTF-8'
SYSTEM_LOCALE="en_US.UTF-8"

# Set the APT mirror URL. Scheme (http/https) is needed.
#    Example: 'http://mirrors.kernel.org'
# Note: AutoDL instances have default fast mirrors.
# Disable this function by setting the var to empty.
APT_MIRROR=""

APT_SOURCE_FILE="/etc/apt/sources.list"
APT_SOURCES_LIST_D="/etc/apt/sources.list.d"

# List of packages to install, with spaces separated names.
#    Example: "tmux git curl htop"
# No package will be installed when this var is empty.
PACKAGES_TO_INSTALL="tmux git curl nano htop nvtop"

# SSH public key for authorization.
#    Example: "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAII7hPh6ucD76+bdbf4Tyb0uq0lsDbFYVMLFZqmxc8zjP YourKey"
SSH_PUBLIC_KEY=""

# uv install
# uv won't be installed when either of following vars is empty.
UV_DOWNLOAD_URL="https://github.com/astral-sh/uv/releases/latest/download/uv-x86_64-unknown-linux-gnu.tar.gz"
UV_INSTALL_PATH="${HOME}/.local/bin"




###############################################################
# Custom function hooks
#
# Here you can add your custom functions to the array,
# which will be executed later one by one.
###############################################################

# The array to hold the function names.
HOOK_FUNCTIONS=()

# Function to add a function to the hook list.
# Usage: add_hook_function <function_name>
add_hook_function() {
    if [[ -n "$1" ]]; then
        HOOK_FUNCTIONS+=("$1")
    fi
}

# You can define your own functions here and add them.
# Example custom function:
# my_custom_setup() {
#     echo "This is my custom setup function."
#     # Add your custom commands here.
#     echo "Setting up a custom user..."
#     useradd myuser
# }
# # Add the function to the hook list.
# add_hook_function my_custom_setup

# >>> Example of function hook.
# If you don't like this, feel free to delete it.
# Prevent conda activate base env automatically.
disable_conda_auto_base() {
    local condarc="${HOME}/.condarc"
    if ! grep -qF "auto_activate_base" $condarc; then
        echo -e "\nauto_activate_base: false" >> "$condarc"
        echo "Conda auto base activation is disabled."
    fi
}
add_hook_function disable_conda_auto_base
# <<< Delete the hook above.




###############################################################
# Custom files
#
# Specify the file contents and their destination. Here you
# can make some configuration so you won't need to do it later.
###############################################################

# File creation helper to create parent directories automatically.
#    Example:
#      write_file "$HOME/my/file.txt" <<EOF
#      file line 1
#      file line 2
#      EOF
write_file() {
    local dest="$1"
    mkdir -p -- "$(dirname -- "$dest")" || return 1
    cat > "$dest"
}

# Tmux config
write_file "${HOME}/.tmux.conf" <<EOF
# tmux configuration
set -g mouse on
set -g default-terminal "screen-256color"
EOF

# uv config
write_file "${HOME}/.config/uv/uv.toml" <<EOF
[[index]]
url = "https://mirrors.bfsu.edu.cn/pypi/web/simple"
default = true
EOF

# awesome vimrc from https://github.com/amix/vimrc/blob/master/vimrcs/basic.vim
# licensed under MIT and encoded with gzip-base64
_vimrc="H4sICAN/zWgAA2Jhc2ljLnZpbXJjAK1abXfaSLL+zq9oa+IBHBmS7O7sPVzDhGAcc64BL+D45oxnPEJqkIJeWL1AyO6d375PVbeEsD135hzbJ8Gou7q6urr6qadaNsTMlWI4mIkrz5ZhIkUND/WKIfCvF613sbd0U1Gz6+Ldm7c/iG7gxWJq+Z4rFzJ0PFtJXss48JLEi0LhJcKVsZzvxDK2wlQ6pljEUopoIWzXipfSFGkkrHAn1jJOMCCap5YXeuFSWMLGjFAH2dSFoiRapFsrlhB3hJUkke1Z0CicyM4CGaZWSjMuPF8mopZiIcZUjzDqPI0jLR/6vFBQb94ptl7qRlkqYpmksWeTFhNCtp85ZEfe7XuBp+eg4eyKBOqgOEuwDrLWFEHkeAv6LXlx62zue4lrCscj5fMsRWNCjexfk9bSjGKRSJ9Mgw4P1vOK9xayFM2zJsem2lUJtWzdKDhcjUc2LbI4xLSSRzkRXMezfpF2Si00YBH5frSlBdoRto7WlbTU/lEQWPNoI3lNas/DKIXJyhDajPV+j3VX4lq+L+ZSuw5zw9HWwbJiMiJJEQie5Yt1FPOsD5fb0FZc9sV0fDG77U76YjAV15Pxp8F5/1wY3SmeDVPcDmaX45uZgMSkO5p9FuML0R19Fv8zGJ2bov+/15P+dCrGEygbDK+vBn20Dka9q5vzweij+ICRozFCfYAYh9rZmKfUygb9Kakb9ie9Szx2PwyuBrPPJlRdDGYj0nsxnoiuuO5OZoPezVV3Iq5vJtfjaR8mnEPxaDC6mGCe/rA/mjUwL9pE/xMexPSye3VFk0Fb9wZrmJCVoje+/jwZfLycicvx1XkfjR/6sK774aqvJsPSelfdwdAU591h92OfR42hh1ZIgspGcXvZp0aas4t/vdlgPKLF9Maj2QSPJtY6mRWDbwfTvim6k8GU3HIxGQ9pmeRYjBmzGowc9ZUecvrh3kCEnm+m/UKlOO93r6BtSoPVQnPxRqViPO8HyobACMIJGVPIqp9HYCROxXsr8L7+ZVWhMd2tTKJA3m9wehB2+4EfZapiGgdh4S1NEVI425FP8Wq7MkDoUtz7UcqxuvazpRcmR6xV/QxCRDWin8LYUvMIPQ/wLgpaJVn6cdN0nbSazSWwJZs37ChokqXNjRfENstOpV0cSfycdmAmlssApp4/DYYEPDGOWSrjhWXLoqtHtiujL6KQUUp1XDA4UvvcslfZet8zk19TgKQ1514vdACoQESfEHY/pZdkOLjAOPmobxhtGLXjKAsd1qQnyhYLeKKQmwJDswRoGu7t7RP+YHBgrdf4vReGP5axXAMerdh2WT9UAk4kAerat3Zo26teA0WRV6S9KjcPvcQuHi6lD+gCQIbKv+TsZ4bjSwR0u7S/z1Y2lQhUN9rCn0is5GmOW4R4GsU7DhzX4vQRI7iDuYwrCc6A7m//7c0bLEn0Q2vuS86n6Q4O12FfedAgonDfpOMGTcoMTu4ZT2Q5SFYSGYE1EjEAAQiXSBN0QPjkIMsmniPZGBpFgypWhhi2s+QjHXfH/JAt+hTv4kQk0BOmR2rDUy+QNOktThTmQCQJH8PpeKRVpKsIuYqWo/Ihgj22xEruEEvB3AstHQtw1kqKMzWysxWJtYHvyDY7i2NaGhlfwWppBj1BWximQXNfWElKQyj4QrJgr6i1PTqz4w5JtW6RBSknF7pZpyFqOM2LzEdmjrE/oeNTvFP/Pteewr0eXCbjOIrrFRiPLXaOxC1WJG2QC1HdiiNWn4JlHYuOaDpy0wwz36+Ks7kVd4TEYTt6qaB9CoReIn7F33XcarYC9ydwy6mKoUBhDRAWzAO4u4MN1PClueLgSaL238nV3U3kOWJpxdh5hwmnZcPMhJhJzyX9UvgIwsxaEg0MnWibiPGU9/fVVXf0sV2VYZVVkhhIZtaWYSWJshgZ4hUWP7kBVRj24WSfehsArKe6iz7iNeBmQpPIW893BHXyHFs88QMRlmWI+KfwXHtkPNPaQsrj3vZJIzJPfjNPGuudXfEWdKprBtbx9gejLv797+L5L++MeoUA8HD863YDCejuxGy4S/pMNuHdSUX6iXxS+KRJ4s0TE1/cpfqNIerL+fR+Sie2Qql3wc73t9YOrJ2AKD89OIbMNHkhceYDeSB5KZliahqogxoZSsGSHTguC7TfslqdUcA0IYkIcT0Hp0IFBsgxgMWaY3wEtNCw5lS4fKHknsGnlPqStUWUNaIRiAmMSegrjM2wI2jhoYVkW0a+iRQfp6aCOLUTrme729hav26fmR3TNf3yzlkILjaqSF48SHmTehmuDgRECnhGxINEJwEmIxaOyoOE1d5zaz72Ek7x2XNKARUwmZ8qSddXjRVmSyv4SQthbQrl9LMXcj5HTM5jhD+lajYztPcKzqOwSmzAia0trRo4qvBGJW07jlBxLaPIIagCfGH/mEGRx+v6+HzbqfEMlDjKsVxmvhVD0Rp2J1wIpHQ0AmsJ7qZjhB84mah8lio/zXGQV5Tl2MEp4JxSD7AAIUgRgNolpmAKlNcwloey1/K8mGIbXU6MFnxBRRC7HrCz0ijzcDZtUdp+R4pGVLeGETEQBBJIDx1qBmYlGEb8MAcnyRs2zJ+ohRvS+828rb4Feda9jkFuYiAa+A3n3weq6dvQsj8BS/IDv8y8e2wC4EWfcsqdODXi482gyJWYZDPXM+6PqINYF3MEvkqIiC7K6BpzfblQcY5S0QEdzoKQjuDLJI7HFPXZWjVhSXYoDb7i2OvTwQdPtUmWyLnJ4wiEANYv2YHYZ5VbFFbJJcRUd7tMj9797QdVKoBjWMgbKZEg8TGkCmBGeTsEpcNOveqNr1B09SdD0W6L6pIETlMtUC3wNr3vRW3ozPcIgMB9agouRoSDtBWnFZsjGnIkU8DVkvl327HiVb5OtbfRWhXbHNpxFvIdC5l6M2AAOAgn3V/KGmjVGk7bsydb5eNFlMRQD/gWIr99PLwTx+k+BsnCLF38FyEwXQ3AcvCb0I6cnO7L8P5mSt0UlIVInrd55bl8mzSR1huA703ofX00jMkncVUV2YuknUHMdKLExBl6qejmOsvMiyxeBXYlegHlzB6UXkDXwgQLJqwNIqKeKTIjwR+2dfppZIolnezUbtAFG5Jxo9HQQKQU6IftXH9JttaauegLeeH/rymfPQPtMWdnWnCSUoVBFAKFp4qJr2tMi0eKhw95RuXwV2yRS9T/ru8Tq5Z9yybjlP5Vq9dJxFukW89J3fZfFWZjeBqt8US3L8CFOeqVFeEzoLxENlX+mysyk24V0HOJ4wmjS8VRiVIkaJuyoWWeAYohjFv6ZFL8h/ujuflT1fofjSsPUrAIT53Q1d13mjHA3VShlIuiRPrquoLOc0b1tSpljsQF1XYWiJkjkWh2qMXhGOmLkZUhCYeVDZEhLpVUNdfBXK2z3mnWIezV1kxz9bVq1RTVav2sN+k0ITXptN836YH+P6Hruz+p68dHul7oBDxxKWIWhcbB7cizZxvSwqcUrx1K4E1RU9tV53l6aeyflrp/FDXCgC0gMWeG9YrynZJpqqfeqX7+kXmg5iRFblXHKS9zqcIl+OG4QbAd7MWBUCuM3LweVuEOeCK7Arp1nst0K4nSK0fllnzp0Odt50vesNINq7zB1Q1u3uDrBibmPT8CYpTjVnm/Ui7V545ofbBJkuxrYceKh+VsryW/7Mv370CFJVpoRw0+d9Qiy70+OkMAI/eUO1x0wHMbL8qS3Dc36kYgvxfjg4dzA1Ks8etAQxoKMjiU20fK4VrqikJ/97jPFgfrPOwMuJP25bA9/6Kn1AsiKARiVVO/iu1cLv39bvIVq3p/QFV3QoWXzYFCK+Gae9midkZf8fbwCiWF31BzCINmE0ajEOXjamViZs2vpIXoOREPVeFzDaYQxrU6HWwxBpGDKQKOYqinNzxPxAVdG1mp+wDTVMhLfWOp3jbpN0qJRezMQ3VH92iH7pLsJxpGERl32jKxrbWsqSxVM45b65ZL76kMYdTJlU0+GzANNVrv9jwn5YX6vFxGuRA+Fcg2Ahn/WS+H8nrrFEduLW1vsVMRLF0LMRfrWpUn5KJHb1se4YqJco5iGbS34RGa3sRvLM6EP2kntVRq+SiUDlmqISaSizyshkOAHZLfCigTSOPetbXPUQZwCNUV/VGdtvpDtphgkdfEeU4EeCtlxZpRvTOMuuggf3///UHTWVs/vqJrEcFhFFKVigS1JBFqVHz0TyXWgyvsP5IvX4NoJloMVpVxkqo2VVZekGHpk6LqmR7b4ND/urSSawyWtfr/HV8cB8fxsYvm7Z24o4Bp4Xt8/K8l+N/WIRGXOu6YqVCfz3Jc1dHji1Hfx3f5z9Y64axOV41vKHAWXpwQXQ1PUVKgWC+oFsf/G/EL33lQIrHYd3xS6IpA0b7u1ez1T19WPxOn6am7Jv1MJb+tYWdI6Sb41gpe06H59VvRvFLNp+90+2Yv3gqqHSV+Fux+7QTflptfd1FJZsUyZ3psB0Jne6GiAIMRpes7PFGFwqWYsuGcMyHNWGpaddQEaNo8lto8lCrKr3PwIYATKkSP75u3rpdqWk0OoUtqU5QupvnVVn7fz/R5kYVHyIrSCvtUaDLHSGqqciQkJhX3+v62LRCPOO41o2HsJSLfuf9nJoFp3I9au1ZtVnXtqS/5j5PmXXL3+lWzqepMpnSJUlZtgMqVpqkfSGh15n6aOi0fZu9drm9MHtyfAGduY7jjOkZWaaQoZk4aXxK+bsXH1lt5+JW4+LBRiklAPM/5hC/+HLr8yRP28DXXs3Ve5zzfTBK68fXz1K3qVv2QHE5bzjUY1oKn/QjrV3LF246pG8WpnaWJPn35kMPxofj5kMoka/HTgxZLfHsw64/iW/vFmDq/JHy2JkAVAY+621c0/5dh/tqCWvObioQiPRGBIkCoxYvi5UqtLhBBcBmkLQQ+CMMn5s3N5pIZWjX9llb5ZcI/Ms9e+TvFAqz8apwPqh17c7rtKrvsnyBR4remEsu36IEO0PEVDA//UNnXsrJGULCLmYqXNeUmVUgCSCiUcEgOFKzXpbBhcR02L7Olj1/1vsD2EntJAJeZJOJRWiPIrbpedCr5jEdin6IZWDDiex5Syf8QIFZ0qHrdnc76Yjg+7wuhbgMVZJRlqhq2VKVdXMszb9e1kqkCzSFQV39DpZlh8aZQ1TYKHM+mg/MOQI6buK4ZpLV6yfzfEchx229pwgyRURYAvTEbWLZxbJSFLB+5ObRS+VDsO4jlfkGT78EzTu3RgHrhLR2R3ykH5e+puIcqkJLfSmq1QXSh88DgYjSI69NjC5MOxu0N0m9da8acXS6PUJk8Kaw0H+zf3s29wCFWVkvScupCQnFWcofs1DJEQ1gt7v4dFQ8vN1SZwH/Jxte/90jZcGpdxPTOvbQ5lDade+RIbMt7o1Ja1J4jb5abnVEpDUJZRDuEIbqKeW+gcLm7azZOqq9++e2nnw/2fy/Nf/3lpeSxopkGhq9oPP/fR4TVKhahrsw31cLvKs9qtxldeyWq5KP9XA1hVFFK1YtIeUJhLNc+8vPvaT1OSCWYw6HegpmUgoXW+b4p2nvBfbPBzYWbDzbwP8YOtGRwKQAA"
printf "%s" "$_vimrc" | base64 -d | gzip -cd | write_file "${HOME}/.vimrc"




###############################################################
# Script functions
#
# These function are for script itself, you can modify them
# as long as you know what you are doing.
###############################################################

# Function to execute all functions in the hook list.
execute_hooks() {
    if [[ ${#HOOK_FUNCTIONS[@]} -ne 0 ]]; then
        echo "--- Executing custom hook functions ---"
        for func_name in "${HOOK_FUNCTIONS[@]}"; do
            if type "$func_name" &>/dev/null; then
                echo "--> Running user hook: $func_name"
                "$func_name"
            else
                echo "Warning: Hook function '$func_name' does not exist."
            fi
        done
    fi
}

# Function to check for root privileges.
check_root() {
    if [[ $EUID -ne 0 ]]; then
        echo "This script must be run as root."
        exit 1
    fi
}

# Download helper, pass the URL and Destination as parameter.
download_file() {
    local url=$1
    local output=$2

    if command -v curl &> /dev/null; then
        curl -sSL -o "$output" "$url" || return 1
    elif command -v wget &> /dev/null; then
        wget -qO "$output" "$url" || return 1
    else
        echo "Error: Neither curl nor wget found. Install one to proceed." >&2
        return 1
    fi
}

# Set AutoDL network turbo to speed up installation.
turbo_autodl_network() {
    local turbo_script="/etc/network_turbo"
    if [[ -f "$turbo_script" ]]; then
        echo "--- Using AutoDL Network Turbo ---"
        . "$turbo_script"
    fi
}

# Set the system locale.
set_locale() {
    echo "--- Setting system locale to '$SYSTEM_LOCALE' ---"
    # Generate the locale if it doesn't exist
    if ! locale -a | grep -q "$SYSTEM_LOCALE"; then
        echo "Generating locale: $SYSTEM_LOCALE"
        locale-gen "$SYSTEM_LOCALE"
    fi
    # Set the system default locale
    update-locale LANG="$SYSTEM_LOCALE" LANGUAGE="$SYSTEM_LOCALE" LC_ALL="$SYSTEM_LOCALE"
    # echo "Locale setting complete."
}

# Function to check for required variables and distro
_check_apt_prerequisites() {
    if [[ -z "$APT_MIRROR" ]]; then
        echo "Error: The APT_MIRROR variable is not set. Please set it to your desired mirror, e.g., 'mirrors.kernel.org'."
        return 1
    fi

    if [[ ! -f "/etc/os-release" ]]; then
        echo "Error: Could not detect OS. This script is intended for Debian/Ubuntu based systems."
        return 1
    fi

    . /etc/os-release
    if [[ "$ID" != "debian" && "$ID" != "ubuntu" ]]; then
        echo "Error: This script is intended for Debian or Ubuntu, but detected '$ID'."
        return 1
    fi

    return 0
}

# Function to detect DEB-822 or legacy sources.list style
_get_source_style() {
    if [[ -d "$APT_SOURCES_LIST_D" ]]; then
        # Check if there are any .sources files in the sources.list.d directory
        if find "$APT_SOURCES_LIST_D" -name "*.sources" -print -quit | grep -q '.*'; then
            echo "deb-822"
            return
        fi
    fi
    echo "legacy"
}

# Function to create the new sources.list file
_create_new_sources_list() {
    local distro_id="$1"
    local distro_version_codename="$2"
    local components=""
    local security_url_suffix=""

    case "$distro_id" in
        "debian")
        components="main contrib non-free"
        security_url_suffix="-security"
        ;;
        "ubuntu")
        components="main restricted universe multiverse"
        security_url_suffix=""
        ;;
        *)
        echo "Error: Unsupported distribution '$distro_id'."
        return 1
        ;;
    esac

    local mirror_url="${APT_MIRROR}/${distro_id}/"
    local security_mirror_url="${APT_MIRROR}/${distro_id}${security_url_suffix}/"

    echo "Creating new $APT_SOURCE_FILE with mirror: $mirror_url"
    
    cat << EOF > "$APT_SOURCE_FILE"
#
# This file was automatically generated.
#
deb ${mirror_url} ${distro_version_codename} ${components}
deb ${mirror_url} ${distro_version_codename}-updates ${components}
deb ${mirror_url} ${distro_version_codename}-backports ${components}

# Security updates
deb ${security_mirror_url} ${distro_version_codename}-security ${components}
EOF

    return 0
}

# Function to replace APT source for legacy sources.list style
_replace_legacy_sources() {
    local distro_id="$1"
    local distro_version_codename="$2"

    echo "Detected legacy sources.list format."
    echo "Backing up old sources.list to ${APT_SOURCE_FILE}.bak..."
    mv "$APT_SOURCE_FILE" "${APT_SOURCE_FILE}.bak"

    _create_new_sources_list "$distro_id" "$distro_version_codename"
}

# Function to replace APT source for DEB-822 style
_replace_deb822_sources() {
    local distro_id="$1"
    local distro_version_codename="$2"
    local components=""
    local security_url_suffix=""

    case "$distro_id" in
        "debian")
        components="main contrib non-free"
        security_url_suffix="-security"
        ;;
        "ubuntu")
        components="main restricted universe multiverse"
        security_url_suffix=""
        ;;
        *)
        echo "Error: Unsupported distribution '$distro_id'."
        return 1
        ;;
    esac

    local new_source_file="${APT_SOURCES_LIST_D}/${distro_id}.sources"
    local mirror_url="${APT_MIRROR}/${distro_id}/"
    local security_mirror_url="${APT_MIRROR}/${distro_id}${security_url_suffix}/"

    echo "Detected DEB-822 format."
    echo "Backing up old sources file to ${new_source_file}.bak..."
    [[ -f "$new_source_file" ]] && mv "$new_source_file" "${new_source_file}.bak"

    echo "Creating new DEB-822 style sources file at: $new_source_file"
    
    cat << EOF > "$new_source_file"
Types: deb
URIs: ${mirror_url}
Suites: ${distro_version_codename} ${distro_version_codename}-updates ${distro_version_codename}-backports
Components: ${components}
Signed-By: /usr/share/keyrings/${distro_id}-archive-keyring.gpg

Types: deb
URIs: ${security_mirror_url}
Suites: ${distro_version_codename}-security
Components: ${components}
Signed-By: /usr/share/keyrings/${distro_id}-archive-keyring.gpg
EOF

    return 0
}

# Replace APT mirror main function.
# Both legacy sources.list and DEB-822 format are suppported.
replace_apt_mirror() {
    if [[ -z "$APT_MIRROR" ]]; then
        echo ">>> Skipping APT mirrors setting up."
        echo "To set up APT mirrors, set APT_MIRROR."
        return 1
    fi

    echo "--- Setting up APT mirrors ---"
    _check_apt_prerequisites || return 1

    . /etc/os-release
    local distro_id="$ID"
    local distro_version_codename="$VERSION_CODENAME"

    if [[ -z "$distro_version_codename" ]]; then
        echo "Error: Could not determine the distribution codename. This is required to proceed."
        return 1
    fi

    local source_style=$(_get_source_style)

    case "$source_style" in
        "deb-822")
        _replace_deb822_sources "$distro_id" "$distro_version_codename"
        ;;
        "legacy")
        _replace_legacy_sources "$distro_id" "$distro_version_codename"
        ;;
        *)
        echo "Error: Unknown source file style detected."
        return 1
        ;;
    esac

    if [[ $? -eq 0 ]]; then
        echo "APT source successfully updated to '$APT_MIRROR'."
    else
        echo "APT source replacement failed."
        return 1
    fi
}

# Run apt update to fetch newest repository index.
update_apt_list() {
    echo "-- Refreshing APT package lists ---"
    apt update
    if [[ $? -eq 0 ]]; then
        echo "APT update completed successfully."
    else
        echo "Warning: 'apt update' failed. Please check your network connection and the mirror URL."
        return 1
    fi
    return 0
}

# Install APT packages.
install_packages() {
    if [[ -z "$PACKAGES_TO_INSTALL" ]]; then
        echo ">>> Skipping APT packages installation."
        echo "To install APT packages, set PACKAGES_TO_INSTALL."
        return 0
    fi

    if ! update_apt_list; then
        echo "!!! Failed to install APT packages."
        return 1
    fi

    echo "--- Installing packages: $PACKAGES_TO_INSTALL ---"
    # Use 'apt-get install -y' for non-interactive installation
    apt-get install -y $PACKAGES_TO_INSTALL
    echo "--- Package installation complete. ---"
}

# Add user SSH public key.
add_ssh_key() {
    # Check if the SSH_PUBLIC_KEY variable is non-empty
    if [[ -n "$SSH_PUBLIC_KEY" ]]; then
        echo "--- Adding SSH_PUBLIC_KEY ---"

        local ssh_dir="${HOME}/.ssh"
        local authorized_keys_file="${ssh_dir}/authorized_keys"

        # Check if the .ssh directory exists, if not, create it with correct permissions
        if [[ ! -d "$ssh_dir" ]]; then
            echo "Creating directory: $ssh_dir"
            mkdir -p "$ssh_dir"
            chmod 700 "$ssh_dir"
        fi

        # Check if the authorized_keys file exists, if not, create it with correct permissions
        if [[ ! -f "$authorized_keys_file" ]]; then
            echo "Creating file: $authorized_keys_file"
            touch "$authorized_keys_file"
            chmod 600 "$authorized_keys_file"
        fi

        # Use grep to check if the public key already exists in the file.
        # The 'grep -q' option suppresses output and exits with a 0 status if a match is found.
        if ! grep -q -F "$SSH_PUBLIC_KEY" "$authorized_keys_file"; then
            echo "Appending key to $authorized_keys_file"
            echo "$SSH_PUBLIC_KEY" >> "$authorized_keys_file"
        else
            echo "Key already exists in $authorized_keys_file."
        fi
    fi
}

# Install uv, a powerful replacement of pip.
install_uv() {
    if [[ -z "$UV_DOWNLOAD_URL" || -z "$UV_INSTALL_PATH" ]]; then
        echo ">>> Skipping uv installation..."
        echo "To install uv, set UV_DOWNLOAD_URL and UV_INSTALL_PATH."
        return 1
    fi

    echo "--- Installing uv ---"
    local url="${UV_DOWNLOAD_URL}"
    local install_path="${UV_INSTALL_PATH}"
    local shell_type=""
    local shell_profile=""

    if ! mkdir -p "$install_path"; then
        echo "Error: failed to create install directory: $install_path" >&2
        return 1
    fi

    # curl -sSL "$url" | tar --strip-components=1 -C "$install_path" -xzf -
    if ! temp_file=$(mktemp); then
        echo "Error: Failed to create temporary file" >&2
        return 1
    fi

    if ! download_file "$url" "$temp_file"; then
        rm -f "$temp_file"
        return 1
    fi

    if ! tar --strip-components=1 -C "$install_path" -xzf "$temp_file"; then
        echo "Error: Failed to extract archive" >&2
        rm -f "$temp_file"
        return 1
    fi
    rm -f "$temp_file"

    echo "uv installed to $install_path"

    if [ -n "$BASH_VERSION" ]; then
        shell_type="bash"
        shell_profile="${HOME}/.bashrc"
    elif [ -n "$ZSH_VERSION" ]; then
        shell_type="zsh"
        shell_profile="${HOME}/.zshrc"
    else
        shell_type="bash"
        shell_profile="${HOME}/.profile"
    fi

    # Ensure install path is in PATH
    if ! echo "$PATH" | grep -q "$install_path"; then
        echo "Adding $install_path to PATH and Shell..."
        export PATH="$install_path:$PATH"
        # Persist for future shells
        if ! grep -q "$install_path" "$shell_profile" 2>/dev/null; then
            echo "export PATH=\"$install_path:\$PATH\"" >> "$shell_profile"
        fi
    fi

    if ! grep -qE "uv(x)?(\s)*(-)*generate-shell-completion" "$shell_profile"; then
        echo "Adding shell completion for uv to profile..."
        echo "eval \"\$(uv generate-shell-completion ${shell_type})\"" >> "$shell_profile"
        echo "eval \"\$(uvx --generate-shell-completion ${shell_type})\"" >> "$shell_profile"
    fi

    echo "Installation complete. Run 'uv --version' to verify."
}

# Remove APT caches.
clean_apt_cache() {
    echo "--- Removing APT caches ---"
    rm -rf /var/cache/apt/archives/* /var/lib/apt/lists/*
}

# Clean bash history.
clean_bash_history() {
    echo "--- Cleaning bash history ---"
    history -c
    # Overwrite the history file with nothing
    > ~/.bash_history
}

# Show hints after finishing.
hint_after_finish() {
    echo "--- Initialization finished. ---"
    echo "Restart the shell to let the changes take effect."
    echo "The simplest way is re-connect the server by SSH."
}




###############################################################
# Main process
#
# The main part of this script.
###############################################################

# Check for root privileges at the beginning.
check_root

# Main functions
set_locale
replace_apt_mirror
install_packages
add_ssh_key

# Install uv
turbo_autodl_network
install_uv

# Execute the functions in the hook list.
execute_hooks

# Clean the caches as the final step.
clean_apt_cache
clean_bash_history

# Hint
hint_after_finish

# Exit with success status.
exit 0