#!/usr/bin/env bash

set -euo pipefail

# If set to 'true', only install the packages and dependencies, without running, reloading, or
# restarting services or systemd.
INSTALL_ONLY=${INSTALL_ONLY:-false}

INSTALL_BIN_DIR=${INSTALL_BIN_DIR:-/usr/local/bin}
INSTALL_SYSTEMD_DIR=${INSTALL_SYSTEMD_DIR:-/etc/systemd/system}

UNCLOUD_GITHUB_URL="https://github.com/psviderski/uncloud"
UNCLOUD_VERSION=${UNCLOUD_VERSION:-latest}
# Remove the 'v' prefix from the version if it exists.
UNCLOUD_VERSION=${UNCLOUD_VERSION#v}
UNCLOUD_USER="uncloud"
# Add the specified Linux user to group $UNCLOUD_USER to allow the user to run uncloud commands without sudo.
UNCLOUD_GROUP_ADD_USER=${UNCLOUD_GROUP_ADD_USER:-}
UNCLOUD_DATA_DIR=${UNCLOUD_DATA_DIR:-/var/lib/uncloud}

DOCKER_ALREADY_INSTALLED=false
CONTAINERD_IMAGE_STORE_ENABLED=false
DOCKER_DAEMON_CONFIG_FILE=${DOCKER_DAEMON_CONFIG_FILE:-/etc/docker/daemon.json}
# Docker daemon configuration optimised for Uncloud.
DOCKER_DAEMON_CONFIG='{
  "features": {
    "containerd-snapshotter": true
  },
  "live-restore": true
}'

log() {
    echo -e "\033[1;32m$1\033[0m"
}

warning() {
    echo -e "\033[1;33m$1\033[0m"
}

error() {
    echo -e "\033[1;31mERROR: $1\033[0m" >&2
    exit 1
}

command_exists() {
    command -v "$1" >/dev/null 2>&1
}

verify_system() {
  if [[ "$(uname -s)" != "Linux" ]]; then
      error "Uncloud machine must be a Linux system. Your system ($(uname -s)) is not supported."
  fi

  local arch
  arch=$(uname -m)
  if [[ "$arch" != "x86_64" && "$arch" != "aarch64" ]]; then
      error "Uncloud machine must have amd64 (x86_64) or arm64 (aarch64) architecture. \
Your system architecture ($arch) is not supported."
  fi

  if [[ ! -d /run/systemd/system && "${INSTALL_ONLY}" != "true" ]]; then
      error "Cannot find systemd to use as a service manager for the Uncloud machine daemon. \
Uncloud supports only systemd-based Linux systems for now."
  fi
}

install_docker() {
    if command_exists dockerd; then
        log "✓ Docker is already installed."
        DOCKER_ALREADY_INSTALLED=true

        if [[ "${INSTALL_ONLY}" == "true" ]]; then
            return
        fi

        docker version

        # Check if the installed Docker configured to use the containerd image store.
        local driver_status
        driver_status=$(docker info -f '{{ .DriverStatus }}' 2>/dev/null)
        if [[ "$driver_status" == *"io.containerd.snapshotter"* ]]; then
            CONTAINERD_IMAGE_STORE_ENABLED="true"
        fi

        return
    fi

    log "⏳ Installing Docker..."
    curl -fsSL https://get.docker.com | sh

    # Configure Docker daemon for new installation.
    # Create Docker daemon config directory if it doesn't exist.
    local docker_config_dir
    docker_config_dir=$(dirname "${DOCKER_DAEMON_CONFIG_FILE}")
    if [ ! -d "${docker_config_dir}" ]; then
        mkdir -p "${docker_config_dir}"
    fi

    log "⏳ Configuring Docker daemon (${DOCKER_DAEMON_CONFIG_FILE}) to optimise it for Uncloud..."
    echo "${DOCKER_DAEMON_CONFIG}" > "${DOCKER_DAEMON_CONFIG_FILE}"

    if [[ "${INSTALL_ONLY}" != "true" ]]; then
        systemctl restart docker
    fi

    log "✓ Docker installed and configured successfully."
}

create_uncloud_user_and_group() {
    if id "${UNCLOUD_USER}" &> /dev/null; then
        log "✓ Linux user '${UNCLOUD_USER}' already exists."
    else
        # In addition to creating the user, create a group with the same name as the user.
        if ! useradd --system --home-dir /nonexistent --shell /usr/sbin/nologin --user-group "${UNCLOUD_USER}"; then
            error "Failed to create Linux user '${UNCLOUD_USER}'."
        fi
        log "✓ Linux user and group '${UNCLOUD_USER}' created."
    fi

    if [ -n "${UNCLOUD_GROUP_ADD_USER}" ]; then
        if ! gpasswd --add "${UNCLOUD_GROUP_ADD_USER}" "${UNCLOUD_USER}" > /dev/null; then
            error "Failed to add user '${UNCLOUD_GROUP_ADD_USER}' to group '${UNCLOUD_USER}'."
        fi
        log "✓ Linux user '${UNCLOUD_GROUP_ADD_USER}' added to group '${UNCLOUD_USER}'."
    fi
}

install_uncloud_binaries() {
    local arch
    local file_arch

    arch=$(uname -m)
    case $arch in
        x86_64)
            file_arch="amd64"
            ;;
        aarch64)
            file_arch="arm64"
            ;;
        *)
            error "Unsupported architecture: ${arch}"
            ;;
    esac

    local uncloudd_install_path="${INSTALL_BIN_DIR}/uncloudd"
    local installed_version=""
    if [ -f "${uncloudd_install_path}" ]; then
        installed_version=$("${uncloudd_install_path}" version -o '{{.Version}}' 2>/dev/null || true)
    fi

    # Decide whether to (re)install based on UNCLOUD_VERSION and the already installed binary, and pick the
    # download URL for the same channel/tag.
    local uncloudd_url uninstall_ref
    local uncloudd_archive_name="uncloudd_linux_${file_arch}.tar.gz"

    case "${UNCLOUD_VERSION}" in
        nightly)
            log "⏳ Installing uncloudd from the nightly channel (installed: ${installed_version:-none})..."
            uncloudd_url="${UNCLOUD_GITHUB_URL}/releases/download/nightly/${uncloudd_archive_name}"
            # The 'nightly' tag is deleted and recreated on every nightly build, so pin uninstall.sh to main
            # to avoid racing.
            uninstall_ref="refs/heads/main"
            ;;
        latest)
            # Resolve the redirect of releases/latest to discover the concrete latest version.
            local latest_url latest_version
            latest_url=$(curl -sLI -o /dev/null -w '%{url_effective}' \
                "${UNCLOUD_GITHUB_URL}/releases/latest" 2>/dev/null || true)
            latest_version="${latest_url##*/}"
            latest_version="${latest_version#v}"

            uninstall_ref="refs/tags/v${latest_version}"

            if [ -z "${latest_version}" ] || [ "${latest_version}" = "latest" ]; then
                warning "Could not resolve the version of the latest release on GitHub."
                # Fall back to the redirecting URL so curl follows whatever GitHub considers latest.
                uncloudd_url="${UNCLOUD_GITHUB_URL}/releases/latest/download/${uncloudd_archive_name}"
                uninstall_ref="refs/heads/main"
            elif [ -z "${installed_version}" ]; then
                log "⏳ Installing uncloudd ${latest_version} (latest stable)..."
                uncloudd_url="${UNCLOUD_GITHUB_URL}/releases/download/v${latest_version}/${uncloudd_archive_name}"
            elif [ "${installed_version}" = "${latest_version}" ]; then
                log "✓ uncloudd ${installed_version} is already the latest stable version."
                return 0
            else
                # Substitute '-' with '~' so sort -V treats pre-release versions per SemVer (Debian-style rules):
                # 0.20.0~nightly-abc < 0.20.0 < 0.21.0~nightly-def.
                # latest_version is always a clean stable tag from releases/latest, so the substitution is one-sided.
                local newest
                newest=$(printf '%s\n%s\n' "${installed_version//-/~}" "${latest_version}" | sort -V | tail -n1)
                if [ "${newest}" = "${latest_version}" ]; then
                    log "⏳ Upgrading uncloudd ${installed_version} → ${latest_version}..."
                    uncloudd_url="${UNCLOUD_GITHUB_URL}/releases/download/v${latest_version}/${uncloudd_archive_name}"
                else
                    log "✓ uncloudd ${installed_version} is newer than the latest stable ${latest_version}, keeping it."
                    return 0
                fi
            fi
            ;;
        *)
            # Explicit version. Install if it differs from the installed one (covers upgrade and downgrade).
            if [ "${installed_version}" = "${UNCLOUD_VERSION}" ]; then
                log "✓ uncloudd ${installed_version} matches the requested version, keeping it."
                return 0
            fi
            log "⏳ Installing uncloudd ${UNCLOUD_VERSION} (replacing ${installed_version:-none})..."
            uncloudd_url="${UNCLOUD_GITHUB_URL}/releases/download/v${UNCLOUD_VERSION}/${uncloudd_archive_name}"
            uninstall_ref="refs/tags/v${UNCLOUD_VERSION}"
            ;;
    esac
    local uninstall_url="https://raw.githubusercontent.com/psviderski/uncloud/${uninstall_ref}/scripts/uninstall.sh"

    # Create a temporary directory for downloads.
    local tmp_dir
    tmp_dir=$(mktemp -d)
    # Ensure the temporary directory is deleted on script exit.
    # shellcheck disable=SC2064
    trap "rm -rf '$tmp_dir'" EXIT

    local uncloudd_download_path="${tmp_dir}/uncloudd.tar.gz"
    local uninstall_download_path="${tmp_dir}/uninstall.sh"

    log "⏳ Downloading uncloudd binary: ${uncloudd_url}"
    if ! curl -fsSL -o "${uncloudd_download_path}" "${uncloudd_url}"; then
        error "Failed to download uncloudd binary."
    fi
    tar -xf "${uncloudd_download_path}" --directory "${tmp_dir}"
    if ! install "${tmp_dir}/uncloudd" "${uncloudd_install_path}"; then
        error "Failed to install uncloud binary to ${uncloudd_install_path}"
    fi
    log "✓ uncloudd binary installed: ${uncloudd_install_path}"

    log "⏳ Downloading uninstall script: ${uninstall_url}"
    if ! curl -fsSL -o "${uninstall_download_path}" "${uninstall_url}"; then
        error "Failed to download uninstall script."
    fi
    local uninstall_install_path="${INSTALL_BIN_DIR}/uncloud-uninstall"
    if ! install "${uninstall_download_path}" "${uninstall_install_path}"; then
        error "Failed to install uninstall.sh script to ${uninstall_install_path}"
    fi
    log "✓ uncloud-uninstall script installed: ${uninstall_install_path}"

    # TODO: install uncloud CLI binary and create a uc alias.
}

install_uncloud_systemd() {
    local uncloud_service_path="${INSTALL_SYSTEMD_DIR}/uncloud.service"
    cat > "${uncloud_service_path}" << EOF
[Unit]
Description=Uncloud machine daemon
After=network-online.target docker.service
Wants=network-online.target

[Service]
Type=notify
ExecStart=${INSTALL_BIN_DIR}/uncloudd
TimeoutStartSec=15
Restart=always
RestartSec=2

# Hardening options.
NoNewPrivileges=true
ProtectSystem=full
ProtectControlGroups=true
ProtectHome=read-only
ProtectKernelTunables=true
PrivateTmp=true
RestrictAddressFamilies=AF_INET AF_INET6 AF_UNIX AF_NETLINK
RestrictNamespaces=true

[Install]
WantedBy=multi-user.target
EOF
    log "✓ Systemd unit file created: ${uncloud_service_path}"


    if [[ "${INSTALL_ONLY}" != "true" ]]; then
        # Reload systemd to recognize the new or updated unit file.
        systemctl daemon-reload
    fi
    systemctl enable uncloud.service
}

start_uncloud() {
    if [[ "${INSTALL_ONLY}" == "true" ]]; then
        return
    fi

    log "⏳ Starting Uncloud machine daemon (uncloud.service)..."
    systemctl restart uncloud.service
    log "✓ Uncloud machine daemon started."
}

log "⏳ Running Uncloud install script..."

if [ "$EUID" -ne 0 ]; then
    error "Please run the install script with sudo or as root."
fi

verify_system
install_docker
create_uncloud_user_and_group
install_uncloud_binaries
install_uncloud_systemd
start_uncloud

# Show warning if Docker was already installed without containerd image store enabled.
if [ "$DOCKER_ALREADY_INSTALLED" = "true" ] && [ "$CONTAINERD_IMAGE_STORE_ENABLED" = "false" ]; then
    echo ""
    warning "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    warning "⚠️  IMPORTANT: Containerd image store configuration"
    warning "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
    warning "Docker was already installed on the machine but it doesn't use the containerd"
    warning "image store. Uncloud works best with the containerd image store enabled in Docker."
    warning "It allows Docker to directly use the images stored in containerd (pushed with"
    warning "'uc deploy' or 'uc image push') without duplicating them in Docker. This saves"
    warning "disk space and makes image management more efficient."
    echo ""
    warning "See https://docs.docker.com/engine/storage/containerd/ for more details."
    echo ""
    warning "To enable it, run the following commands on the machine:"
    echo ""
    echo "sudo bash -c 'cat > ${DOCKER_DAEMON_CONFIG_FILE} << EOF"
    echo "${DOCKER_DAEMON_CONFIG}"
    echo "EOF'"
    echo "sudo systemctl restart docker"
    echo ""
    warning "WARNING: Switching to containerd image store causes you to temporarily lose images"
    warning "and containers created using the classic storage driver. Those resources still"
    warning "exist on your filesystem, and you can retrieve them by turning off the containerd"
    warning "image store feature."
    warning "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
fi

log "✓ Uncloud installed on the machine successfully! 🎉"
