#!/bin/bash

#################################################
# @Author: Einic <einicyeo AT gmail.com>        #
# @Description:                                 #
# @File: docker-manage.sh                       #
# @Version: 1.0.0                               #
# @Date: 2025/3/13 14:03                        #
# @BLOG:  https://www.infvie.com                #
# @Project home page:                           #
#     @https://github.com/Einic/EnvoyinStack    #
################################################# 

export PATH=/sbin:/bin:/usr/sbin:/usr/bin:/usr/local/sbin:/usr/local/bin
set -euo pipefail


readonly DOCKER_VERSION="28.0.0"
readonly ARCH="$(uname -m)"
readonly REGISTRY_MIRROR="https://d8b3zdiw.mirror.aliyuncs.com"
readonly INSECURE_REGISTRY="127.0.0.1:5000"
readonly DATA_ROOT="/data/docker"
readonly DOCKER_BIN_ROOT="/data/docker-server"
readonly SYSTEMD_SERVICE="/usr/lib/systemd/system/docker.service"
readonly SYSTEMD_REGISTRY_SERVICE="/usr/lib/systemd/system/registry.service"
readonly DOCKER_COMPOSE_BIN="/usr/local/bin/docker-compose"
readonly DOCKER_REGISTRY_BIN="/usr/local/bin/registry"
readonly DATA_REGISTRY="/data/docker-registry"
readonly JQ_BIN="/usr/local/bin/jq"

readonly RED='\033[0;31m'
readonly GREEN='\033[0;32m'
readonly YELLOW='\033[0;33m'
readonly NC='\033[0m' # No Color


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


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


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


error_exit() {
    echo -e "${RED}[✗] $(date '+%Y-%m-%d %H:%M:%S') $1${NC}" >&2
    exit 1
}


check_exists() {
    if [[ ! -e "$1" ]]; then
        error_exit "File or directory not found: $1"
    fi
}

# Check if user is root
[ $(id -u) != "0" ] && { warn "You must be root to run this script"; exit 1; }


validate_architecture() {
    case "${ARCH}" in
        x86_64|aarch64) 
            log "Validated architecture: ${ARCH}"
            ;;
        *) 
            error_exit "Unsupported architecture: ${ARCH}"
            ;;
    esac
}


clean_directory() {
    local dir_path="$1"
    if [[ -d "${dir_path}" ]]; then
        log "Removing directory: ${dir_path}"
        rm -rf "${dir_path}" || warn "Failed to remove ${dir_path}"
    else
        log "Directory not found: ${dir_path}"
    fi
}


create_directories() {
    local dirs=(
        "${DATA_ROOT}" 
		"${DATA_REGISTRY}" 
        "/etc/docker/registry" 
        "${DOCKER_BIN_ROOT}"
    )
    
    log "Creating system directories..."
    for dir in "${dirs[@]}"; do
        mkdir -p "${dir}" || error_exit "Failed to create directory: ${dir}"
        log "Created directory: ${dir}"
    done
}

install_dependencies() {
    log "Installing required dependencies..."

    # dnf > yum > apt-get > pacman > zypper > apk > emerge > nix > pkg > snap > flatpak
    if command -v dnf >/dev/null 2>&1; then
        log "Detected dnf package manager"
        dnf install -y tar unzip git wget curl vim || error_exit "Failed to install dependencies with dnf"
    elif command -v yum >/dev/null 2>&1; then
        log "Detected yum package manager"
        yum install -y tar unzip git wget curl vim || error_exit "Failed to install dependencies with yum"
    elif command -v apt-get >/dev/null 2>&1; then
        log "Detected apt package manager"
        apt-get update || error_exit "Failed to update apt repositories"
        apt-get install -y tar unzip git wget curl vim || error_exit "Failed to install dependencies with apt"
    elif command -v pacman >/dev/null 2>&1; then
        log "Detected pacman package manager"
        pacman -Syu --noconfirm tar unzip git wget curl vim || error_exit "Failed to install dependencies with pacman"
    elif command -v zypper >/dev/null 2>&1; then
        log "Detected zypper package manager"
        zypper refresh || error_exit "Failed to refresh zypper repositories"
        zypper install -y tar unzip git wget curl vim || error_exit "Failed to install dependencies with zypper"
    elif command -v apk >/dev/null 2>&1; then
        log "Detected apk package manager"
        apk update || error_exit "Failed to update apk repositories"
        apk add tar unzip git wget curl vim || error_exit "Failed to install dependencies with apk"
    elif command -v emerge >/dev/null 2>&1; then
        log "Detected emerge package manager"
        emerge --sync || error_exit "Failed to sync emerge repositories"
        emerge tar unzip git wget curl vim || error_exit "Failed to install dependencies with emerge"
    elif command -v nix-env >/dev/null 2>&1; then
        log "Detected nix package manager"
        nix-env -iA nixpkgs.tar nixpkgs.unzip nixpkgs.git nixpkgs.wget nixpkgs.curl nixpkgs.vim || error_exit "Failed to install dependencies with nix"
    elif command -v pkg >/dev/null 2>&1; then
        log "Detected pkg package manager"
        pkg install -y tar unzip git wget curl vim || error_exit "Failed to install dependencies with pkg"
    elif command -v snap >/dev/null 2>&1; then
        log "Detected snap package manager"
        snap install tar unzip git wget curl vim || error_exit "Failed to install dependencies with snap"
    elif command -v flatpak >/dev/null 2>&1; then
        log "Detected flatpak package manager"
        flatpak install -y flathub tar unzip git wget curl vim || error_exit "Failed to install dependencies with flatpak"
    else
        error_exit "Unsupported package manager. Supported: dnf/yum/apt/pacman/zypper/apk/emerge/nix/pkg/snap/flatpak"
    fi
}

update_mysql_password() {
    local db_password
    db_password=$(openssl rand -base64 12)
    local escaped_password
    escaped_password=$(printf '%s\n' "$db_password" | sed -e 's/[\/&]/\\&/g')
    local compose_file="docker-compose.yml"

    sed -i.bak -E "s|(MYSQL_ROOT_PASSWORD: ).*|\1\"${escaped_password}\"|" "$compose_file"
    sed -i.bak -E "s|(DB_PASSWORD=).*|\1\"${escaped_password}\"|" "$compose_file"

    success "Generated DB_PASSWORD: $db_password"
}


update_project_name() {
   local compose_file="docker-compose.yml"
   local app_dir="./app"

   # Check if app directory exists
   if [[ ! -d "$app_dir" ]]; then
       error_exit "Error: Directory '$app_dir' does not exist."
       return 1
   fi

   # Get the name of the first subdirectory under the app directory
   local first_subdir=$(find "$app_dir" -mindepth 1 -maxdepth 1 -type d | head -n 1)

   # Check if subdirectory is found
   if [[ -z "$first_subdir" ]]; then
      error_exit "Error: Subdirectory not found in directory '$app_dir'."
      return 1
   fi

   # Extract subdirectory name
   local project_name=$(basename "$first_subdir")

   # Update PROJECT_NAME value in docker-compose.yml using sed
   sed -i.bak -E "s/(PROJECT_NAME=)[^\"]*/\1$project_name/" "$compose_file"

   if [[ $? -eq 0 ]]; then
      success "Updated 'PROJECT_NAME' to '$project_name'. "
      return 0
   else
      error_exit " Error updating 'docker-compose.yml' file. "
      return 1
   fi
}

install_docker_registry() {
    local registry_files=("./${ARCH}/registry"*)

    if [[ ${#registry_files[@]} -eq 0 ]]; then
        error_exit "No docker-registry files found in ./${ARCH}/"
    fi

    local src_file="${registry_files[0]}"
    check_exists "${src_file}"
    
    if ! command -v "${DOCKER_REGISTRY_BIN}" >/dev/null 2>&1; then
       log "Installing docker-registry from ${src_file}..."
       cp -v "${src_file}" "${DOCKER_REGISTRY_BIN}" || error_exit "Copy failed"
       chmod -v +x "${DOCKER_REGISTRY_BIN}" || error_exit "Permission modification failed"
       success "docker-registry installed successfully"
   fi
}

load_docker_registry_images() {
    local image_dir="./${ARCH}/images"
    check_exists "${image_dir}"
	
    for image in "${image_dir}"/*; do
        log "Loading Docker image from ${image}..."
        docker load -i "${image}" || error_exit "Failed to load image from ${image}"

        local image_name=$(docker load -i "${image}" | awk 'END{print $NF}')
        log "Tagging image ${image_name} for local registry..."
        docker tag "${image_name}" "127.0.0.1:5000/${image_name}" || error_exit "Failed to tag image ${image_name}"

        log "Pushing image ${image_name} to local registry..."
        docker push "127.0.0.1:5000/${image_name}" || error_exit "Failed to push image ${image_name} to local registry"
    done

    success "All Docker images have been loaded and pushed to the local registry successfully"
}

install_jq() {
    local jq_files=("./${ARCH}/jq"*)

    if [[ ${#jq_files[@]} -eq 0 ]]; then
        error_exit "No jq files found in ./${ARCH}/"
    fi

    local src_file="${jq_files[0]}"
    check_exists "${src_file}"

    log "Installing jq from ${src_file}..."
    cp -v "${src_file}" "${JQ_BIN}" || error_exit "Copy failed"
    chmod -v +x "${JQ_BIN}" || error_exit "Permission modification failed"
    success "jq installed successfully"
}

install_docker_compose() {
    local compose_files=("./${ARCH}/docker-compose"*)

    if [[ ${#compose_files[@]} -eq 0 ]]; then
        error_exit "No docker-compose files found in ./${ARCH}/"
    fi

    local src_file="${compose_files[0]}"
    check_exists "${src_file}"

    if ! command -v "${DOCKER_COMPOSE_BIN}" >/dev/null 2>&1; then
       log "Installing docker-compose from ${src_file}..."
       cp -v "${src_file}" "${DOCKER_COMPOSE_BIN}" || error_exit "Copy failed"
       chmod -v +x "${DOCKER_COMPOSE_BIN}" || error_exit "Permission modification failed"
       success "docker-compose installed successfully"
    fi
}

deploy_docker_binaries() {
    local tar_file="./${ARCH}/docker-${DOCKER_VERSION}.tgz"
    check_exists "${tar_file}"
    
    log "Extracting docker binaries..."
    mkdir -p "${DOCKER_BIN_ROOT}" || error_exit "Cannot create docker bin directory"
 
    if ! command -v docker >/dev/null 2>&1; then
       tar -xvf "${tar_file}" -C "${DOCKER_BIN_ROOT}" --strip-components=1 || error_exit "Failed to extract ${tar_file}"

       log "Deploying docker binaries..."
       cp -v "${DOCKER_BIN_ROOT}"/* /usr/local/bin/ || error_exit "Binary deployment failed"
       success "docker binaries installed successfully"
    fi
}

configure_registry() {
    local registry_file="/etc/docker/registry/config.yml"
    log "Generating Docker registry configuration..."
    cat > "${registry_file}" << EOF
version: 0.1
log:
  fields:
    service: registry
storage:
  filesystem:
    rootdirectory: /data/docker-registry
http:
  addr: :5000
EOF
    check_exists "${registry_file}"
}

configure_daemon() {
    local daemon_file="/etc/docker/daemon.json"
    log "Generating Docker daemon configuration..."
    cat > "${daemon_file}" << EOF
{
    "exec-opts": ["native.cgroupdriver=systemd"],
    "max-concurrent-downloads": 10,
    "max-concurrent-uploads": 5,
    "log-driver": "json-file",
    "storage-driver": "overlay2",
    "data-root": "${DATA_ROOT}",
    "registry-mirrors": ["${REGISTRY_MIRROR}"],
    "insecure-registries": ["${INSECURE_REGISTRY}"],
    "log-opts": {
        "max-size": "300m",
        "max-file": "2"
    },
    "live-restore": true
}
EOF
    check_exists "${daemon_file}"
}

configure_registry_systemd() {
    log "Creating registry systemd service unit..."
    cat > "${SYSTEMD_REGISTRY_SERVICE}" << EOF
[Unit]
Description=Docker Registry Service
After=network.target

[Service]
ExecStart=/usr/local/bin/registry serve /etc/docker/registry/config.yml
Restart=always
User=root
Group=root

[Install]
WantedBy=multi-user.target
EOF
    check_exists "${SYSTEMD_REGISTRY_SERVICE}"
}

configure_systemd() {
    log "Creating systemd service unit..."
    cat > "${SYSTEMD_SERVICE}" << EOF
[Unit]
Description=Docker Application Container Engine
Documentation=https://docs.docker.com
After=network-online.target firewalld.service
Wants=network-online.target

[Service]
Type=notify
ExecStart=/usr/local/bin/dockerd
ExecReload=/bin/kill -s HUP \$MAINPID
LimitNOFILE=infinity
LimitNPROC=infinity
LimitCORE=infinity
TimeoutStartSec=0
Delegate=yes
KillMode=process
Restart=on-failure
StartLimitBurst=3
StartLimitInterval=60s

[Install]
WantedBy=multi-user.target
EOF
    check_exists "${SYSTEMD_SERVICE}"
}

start_registry_services() {
    log "Starting docker registry services..."
    systemctl daemon-reload || error_exit "Systemd reload failed"
    systemctl enable registry --now || error_exit "Service management failed"
}

start_services() {
    log "Starting docker services..."
    systemctl daemon-reload || error_exit "Systemd reload failed"
    systemctl enable docker --now || error_exit "Service management failed"
}


uninstall_binaries() {

    if [[ -f "${DOCKER_COMPOSE_BIN}" ]]; then
        log "Removing docker-compose..."
        rm -vf "${DOCKER_COMPOSE_BIN}" || warn "Failed to remove docker-compose"
    fi


    if [[ -f "${DOCKER_REGISTRY_BIN}" ]]; then
        log "Removing docker-registry..."
        rm -vf "${DOCKER_REGISTRY_BIN}" || warn "Failed to remove docker-registry"
    fi
	

    if [[ -d "${DOCKER_BIN_ROOT}" ]]; then
        log "Removing docker binaries..."
        find "${DOCKER_BIN_ROOT}" -type f -exec rm -vf /usr/local/bin/{} \; || warn "Binary removal failed"
        clean_directory "${DOCKER_BIN_ROOT}"
    fi
}

uninstall_service() {

    if systemctl is-active docker &>/dev/null; then
        log "Stopping docker service..."
        systemctl stop docker || warn "Failed to stop docker service"
    fi

    if systemctl is-enabled docker &>/dev/null; then
        log "Disabling docker service..."
        systemctl disable docker || warn "Failed to disable docker service"
    fi


    if [[ -f "${SYSTEMD_SERVICE}" ]]; then
        log "Removing systemd service..."
        rm -vf "${SYSTEMD_SERVICE}" || warn "Failed to remove service file"
        systemctl daemon-reload
    fi
}

uninstall_registry_service() {

    if systemctl is-active registry &>/dev/null; then
        log "Stopping docker registry service..."
        systemctl stop registry || warn "Failed to stop docker registry service"
    fi

    if systemctl is-enabled registry &>/dev/null; then
        log "Disabling docker registry service..."
        systemctl disable registry || warn "Failed to disable docker registry service"
    fi


    if [[ -f "${SYSTEMD_REGISTRY_SERVICE}" ]]; then
        log "Removing systemd registry service..."
        rm -vf "${SYSTEMD_REGISTRY_SERVICE}" || warn "Failed to remove registry service file"
        systemctl daemon-reload
    fi
}

purge_data() {
    log "Purging docker data..."
    clean_directory "${DATA_ROOT}"
	#clean_directory "${DATA_REGISTRY}"
    clean_directory "/etc/docker"
}


validate_installation() {
    local result=0

    check_binary() {
        local bin_path="$1"
        if [[ -x "${bin_path}" ]]; then
            success "Found executable: ${bin_path}"
            "${bin_path}" --version || {
                warn "Version check failed: ${bin_path}"
                result=1
            }
        else
            warn "Missing executable: ${bin_path}"
            result=1
        fi
    }

    check_binary "/usr/local/bin/docker"
    check_binary "/usr/local/bin/docker-compose"


    if systemctl is-active docker &>/dev/null; then
        success "Docker service is running"
    else
        warn "Docker service is not running"
        result=1
    fi


    check_directory "${DATA_ROOT}"
    check_directory "/etc/docker"


    check_file "/etc/docker/daemon.json"

    pull_private_images

    log "Validating container operation..."
    if docker run --rm alpine:3.19&>/dev/null; then
        success "Container validation passed"
    else
        warn "Container validation failed"
        result=1
    fi

    return ${result}
}

pull_private_images() {
    local registry="127.0.0.1:5000"
    local result=0

    log "Fetching list of repositories from private registry..."
    local repositories=$(curl -s http://${registry}/v2/_catalog | jq -r '.repositories[]')

    if [[ -z "${repositories}" ]]; then
        warn "No repositories found in private registry."
        return 1
    fi

    for repo in ${repositories}; do
        log "Fetching tags for repository: ${repo}..."
        local tags=$(curl -s http://${registry}/v2/${repo}/tags/list | jq -r '.tags[]')

        if [[ -z "${tags}" ]]; then
            warn "No tags found for repository: ${repo}."
            result=1
            continue
        fi

        for tag in ${tags}; do
            log "Pulling image: ${registry}/${repo}:${tag}..."
            if docker pull ${registry}/${repo}:${tag}; then
                success "Successfully pulled image: ${registry}/${repo}:${tag}"
            else
                warn "Failed to pull image: ${registry}/${repo}:${tag}"
                result=1
            fi
        done
    done

    return ${result}
}

check_file() {
    local file_path="$1"
    if [[ -f "${file_path}" ]]; then
        success "Found configuration file: ${file_path}"
    else
        warn "Missing configuration file: ${file_path}"
    fi
}

check_directory() {
    local dir_path="$1"
    if [[ -d "${dir_path}" ]]; then
        success "Found directory: ${dir_path}"
    else
        warn "Missing directory: ${dir_path}"
    fi
}

backup_db() {
  local backup_dir="./docker/mysql/data"

  if [ ! -d "$backup_dir" ] || [ -z "$(ls -A "$backup_dir")" ]; then
    warn "No database backup files found in ${backup_dir}/"
  else
    mv "$backup_dir" "${backup_dir}_$(date +%Y%m%d%H%M%S)"
  fi
}


shutdown_docker_compose() {
   if command -v docker-compose >/dev/null 2>&1; then
      log "docker-compose command detected, execute docker-compose down -v..."
      docker-compose down -v
   else
      log "docker-compose command not detected, skip operation."
fi
}

install() {
    log "===== Starting Docker Installation ====="
    validate_architecture
	create_directories
    install_dependencies
    install_jq
	install_docker_registry
    configure_registry
    configure_registry_systemd
    start_registry_services
    install_docker_compose
    deploy_docker_binaries
    configure_daemon
    configure_systemd
    start_services
	load_docker_registry_images
    success "===== Docker Installation Completed ====="
}

uninstall() {
    log "===== Starting Docker Uninstallation ====="
    shutdown_docker_compose
    backup_db
    uninstall_registry_service
    uninstall_service
    uninstall_binaries
    purge_data
    success "===== Docker Uninstallation Completed ====="
    warn "Note: All docker data in ${DATA_ROOT} has been removed!"
}

validate() {
    log "===== Starting System Validation ====="
    if validate_installation; then
        success "===== All Checks Passed ====="
    else
        error_exit "===== Validation Failed ====="
    fi
}

usage() {
clear
printf "
#################################################
# @Author: Einic <einicyeo AT gmail.com>        #
# @Description:                                 #
# @File: docker-manage.sh                       #
# @Version: 1.0.0                               #
# @Date: 2025/3/11 14:03                        #
# @BLOG:  https://www.infvie.com                #
# @Project home page:                           #
#     @https://github.com/Einic/EnvoyinStack    #
################################################# 
"
    cat << EOF
Usage: $0 <command>

Commands:
  init        Initialize the environment
  install     Install the application
  uninstall   Uninstall the application
  validate    Validate the installation
  app-build   Build the application
  app-run     Run the application
  app-ps      Display application process status

EOF
    exit 1
}


app-build() {
    docker-compose --build
}

app-run() {
    docker-compose up -d
}

app-ps() {
    docker-compose ps
}

main() {
    case "${1:-}" in
	    init)
            update_project_name
			update_mysql_password
            ;;
        install)
            install
            ;;
        uninstall)
            uninstall
            ;;
        validate)
            validate
            ;;
		app-build)
            app-build
            ;;
		app-run)
            app-run
            ;;	
		app-ps)
            app-ps
            ;;		
        *)
            usage
            ;;
    esac
}

main "$@"
