#!/usr/bin/env bash
set -euo pipefail

if [[ "${DEBUG:-}" == "1" || "${UNREGISTRY_DEBUG:-}" == "1" ]]; then
    set -x
fi

# Script version
VERSION="0.3.0"

# Return metadata expected by the Docker CLI plugin framework: https://github.com/docker/cli/pull/1564
if [[ "${1:-}" = "docker-cli-plugin-metadata" ]]; then
    cat <<EOF
{
  "SchemaVersion": "0.1.0",
  "Vendor": "https://github.com/psviderski",
  "Version": "${VERSION}",
  "ShortDescription": "Upload image to remote Docker daemon via SSH without external registry"
}
EOF
    exit 0
fi

# Pin the unregistry image version. The image doesn't change too often compared to the script,
# so we want to avoid unnecessary pulls of the latest image version.
UNREGISTRY_VERSION=0.1.3
UNREGISTRY_IMAGE=${UNREGISTRY_IMAGE:-ghcr.io/psviderski/unregistry:${UNREGISTRY_VERSION}}

DEFAULT_CONTAINERD_SOCKET="/run/containerd/containerd.sock"
# Containerd socket path on remote host. It's populated by find_containerd_socket function.
# Can be overridden by setting CONTAINERD_SOCKET environment variable.
CONTAINERD_SOCKET=${CONTAINERD_SOCKET:-"${DEFAULT_CONTAINERD_SOCKET}"}
# SSH strict host key checking mode. Can be set to "yes", "no", "ask", or "accept-new".
# Default is to use the SSH client's default behavior.
SSH_STRICT_HOST_KEY_CHECKING=${SSH_STRICT_HOST_KEY_CHECKING:-}

# Ensure localhost connections bypass proxy.
export no_proxy="${no_proxy:-},localhost,127.0.0.1"

# Colors and symbols for output.
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
BOLD='\033[1m'
RST='\033[0m' # reset formatting

info() {
    echo -e " ${BLUE}•${RST} $1"
}

success() {
    echo -e " ${GREEN}✓${RST} $1"
}

warning() {
    echo -e " ${YELLOW}!${RST} $1"
}

error() {
    echo -e "${RED}ERROR:${RST} $1" >&2
    exit 1
}

usage() {
    echo "Usage: docker pussh [OPTIONS] IMAGE[:TAG] [USER@]HOST[:PORT]"
    echo ""
    echo "Upload a Docker image to a remote Docker daemon via SSH without an external registry."
    echo ""
    echo "Options:"
    echo "  -h, --help              Show this help message."
    echo "  -i, --ssh-key path      Path to SSH private key for remote login (if not already added to SSH agent)."
    echo "      --no-host-key-check Skip SSH host key checking (use with caution)."
    echo "      --platform string   Push a specific platform for a multi-platform image (e.g., linux/amd64, linux/arm64)."
    echo "                          Local Docker has to use containerd image store to support multi-platform images."
    echo ""
    echo "Examples:"
    echo "  docker pussh myimage:latest user@host"
    echo "  docker pussh --platform linux/amd64 myimage:latest host"
    echo "  docker pussh myimage:latest user@host:2222 -i ~/.ssh/id_ed25519"
}

# SSH command arguments to be used for all ssh commands after establishing a shared "master" connection
# using ssh_remote.
declare -a SSH_ARGS=()

# Establish SSH connection to the remote server that will be reused by subsequent ssh commands via the control socket.
# It populates the SSH_ARGS array with arguments for reuse.
ssh_remote() {
    local ssh_addr="$1"
    local target port
    # Split out the port component, if exists
    if [[ "${ssh_addr}" =~ ^([^:]+)(:([0-9]+))?$ ]]; then
        target="${BASH_REMATCH[1]}"
        port="${BASH_REMATCH[3]:-}"
    else
        error "Invalid SSH address format. Expected format: [USER@]HOST[:PORT]"
    fi

    local ssh_opts=(
        -o "ControlMaster=auto"
        # Unique control socket path for this invocation.
        -o "ControlPath=/tmp/docker-pussh-$$.sock"
        # The connection will be automatically terminated after 1 minute of inactivity.
        -o "ControlPersist=1m"
        -o "ConnectTimeout=15"
    )
    # Add port if specified
    if [[ -n "${port}" ]]; then
        ssh_opts+=(-p "${port}")
    fi
    # Add SSH key option if provided.
    if [[ -n "${SSH_KEY}" ]]; then
        ssh_opts+=(-i "${SSH_KEY}")
    fi
    # Configure SSH host key checking if requested.
    if [[ -n "${SSH_STRICT_HOST_KEY_CHECKING}" ]]; then
        ssh_opts+=(-o "StrictHostKeyChecking=${SSH_STRICT_HOST_KEY_CHECKING}")
        if [[ "${SSH_STRICT_HOST_KEY_CHECKING}" == "no" ]]; then
            ssh_opts+=(-o "UserKnownHostsFile=/dev/null")
        fi
    fi

    # Establish ControlMaster connection in the background.
    if ! ssh "${ssh_opts[@]}" -f -N "${target}"; then
        error "Failed to connect to remote host via SSH: ${ssh_addr}"
    fi

    # Populate SSH_ARGS array for reuse in all subsequent commands.
    SSH_ARGS=("${ssh_opts[@]}")
    SSH_ARGS+=("${target}")
}

# sudo prefix for remote docker commands. It's set to "sudo -n" if the remote user is not root and requires sudo
# to run docker commands.
REMOTE_SUDO=""

# Check if the remote host has Docker installed and if we can run docker commands.
# If sudo is required, it sets the REMOTE_SUDO variable to "sudo -n".
check_remote_docker() {
    # Check if docker command is available.
    if ! ssh "${SSH_ARGS[@]}" "command -v docker" >/dev/null 2>&1; then
        error "'docker' command not found on remote host. Please ensure Docker is installed."
    fi
    # Check if we need sudo to run docker commands.
    if ! ssh "${SSH_ARGS[@]}" "docker version" >/dev/null 2>&1; then
        # Check if we're not root and if sudo docker works.
        if ssh "${SSH_ARGS[@]}" "[ \$(id -u) -ne 0 ] && sudo -n docker version" >/dev/null; then
            REMOTE_SUDO="sudo -n"
        else
            error "Failed to run docker commands on remote host. Please ensure:
  - Docker is installed and running on the remote host
  - SSH user has permissions to run docker commands (user is root or non-root user is in 'docker' group)
  - If sudo is required, ensure the user can run 'sudo docker' without a password prompt"
        fi
    fi
}

# Generate a random port in range 55000-65535.
random_port() {
    echo $((55000 + RANDOM % 10536))
}

# Container name for the unregistry instance on remote host. It's populated by run_unregistry function.
UNREGISTRY_CONTAINER=""
# Unregistry port on the remote host that is bound to localhost. It's populated by run_unregistry function.
UNREGISTRY_PORT=""

# Find the containerd socket path on the remote host
# If no socket is found, keeps the default value to avoid regression
find_containerd_socket() {
    # Skip detection if CONTAINERD_SOCKET was explicitly set by user
    if [[ "${CONTAINERD_SOCKET}" != "${DEFAULT_CONTAINERD_SOCKET}" ]]; then
        return 0
    fi

    local socket_paths=(
        "${DEFAULT_CONTAINERD_SOCKET}"
        "/var/run/docker/containerd/containerd.sock"
        "/var/run/containerd/containerd.sock"
        "/run/docker/containerd/containerd.sock"
        "/run/snap.docker/containerd/containerd.sock"
    )

    for socket_path in "${socket_paths[@]}"; do
        # Try without sudo first, then with sudo if REMOTE_SUDO is set
        # shellcheck disable=SC2029
        if ssh "${SSH_ARGS[@]}" "test -S '${socket_path}'" 2>/dev/null ||
           ssh "${SSH_ARGS[@]}" "sudo -n test -S '${socket_path}'" 2>/dev/null; then
            CONTAINERD_SOCKET="${socket_path}"
            return 0
        fi
    done

    # If no socket found, keep the default and let the container startup handle the error
    # This ensures we don't introduce a regression for users who had working setups
}

# Run unregistry container on remote host with retry logic for port binding conflicts.
# Sets UNREGISTRY_PORT and UNREGISTRY_CONTAINER global variables.
run_unregistry() {
    local output

    # Find containerd socket first
    find_containerd_socket

    # Pull unregistry image if it doesn't exist on the remote host. This is done separately to not capture the output
    # and print the pull progress to the terminal.
    # shellcheck disable=SC2029
    if ! ssh "${SSH_ARGS[@]}" "${REMOTE_SUDO} docker image inspect ${UNREGISTRY_IMAGE}" >/dev/null 2>&1; then
        ssh "${SSH_ARGS[@]}" "${REMOTE_SUDO} docker pull ${UNREGISTRY_IMAGE}"
    fi

    for _ in {1..10}; do
        UNREGISTRY_PORT=$(random_port)
        UNREGISTRY_CONTAINER="unregistry-pussh-$$-${UNREGISTRY_PORT}"

        # shellcheck disable=SC2029
        if output=$(ssh "${SSH_ARGS[@]}" "${REMOTE_SUDO} docker run -d \
            --name ${UNREGISTRY_CONTAINER} \
            -p 127.0.0.1:${UNREGISTRY_PORT}:5000 \
            -v ${CONTAINERD_SOCKET}:/run/containerd/containerd.sock \
            --userns=host \
            --user root:root \
            ${UNREGISTRY_IMAGE}" 2>&1);
        then
            # Wait a moment for the container to start
            sleep 1

            # Verify the container is actually running and healthy
            if ssh "${SSH_ARGS[@]}" "${REMOTE_SUDO} docker ps --filter name=${UNREGISTRY_CONTAINER} --filter status=running --quiet" | grep -q .; then
                return 0
            fi
        fi

        # Remove the container that failed to start if it was created.
        # shellcheck disable=SC2029
        ssh "${SSH_ARGS[@]}" "${REMOTE_SUDO} docker rm -f ${UNREGISTRY_CONTAINER}" >/dev/null 2>&1 || true
        # Check if the error is due to port binding.
        if ! echo "${output}" | grep -q --ignore-case "bind.*${UNREGISTRY_PORT}"; then
            error "Failed to start unregistry container:\n${output}"
        fi
    done

    error "Failed to start unregistry container:\n${output}"
}

# Forward a local port to a remote port over the established SSH connection.
# Returns the local port that was successfully bound.
forward_port() {
    local remote_port="$1"
    local local_port
    local output

    for _ in {1..10}; do
        local_port=$(random_port)

        # Check if port is already in use locally.
        # TODO: handle the case when nc is not available.
        if command -v nc >/dev/null && nc -z 127.0.0.1 "${local_port}" 2>/dev/null; then
            continue
        fi

        if output=$(ssh "${SSH_ARGS[@]}" -O forward -L "${local_port}:127.0.0.1:${remote_port}" 2>&1); then
            echo "${local_port}"
            return 0
        fi

        error "Failed to forward local port ${local_port} to remote unregistry port 127.0.0.1:${remote_port}: ${output}"
    done

    error "Failed to find an available local port to forward to remote unregistry port. Please try again."
}

# Check if the local Docker server needs additional tunneling.
is_additional_tunneling_needed() {
    # Read all output to a variable to avoid issues with pipefail when 'grep -q' exits early.
    local output
    output=$(docker version 2>/dev/null)
    echo "${output}" | grep -E -q "Docker Desktop|colima" && return 0
    return 1
}

# Container name for the Docker Desktop tunnel. It's populated by run_docker_desktop_tunnel function.
DOCKER_DESKTOP_TUNNEL_CONTAINER=""
# Port on localhost that docker in Docker Desktop should push to. It's populated by run_docker_desktop_tunnel function.
DOCKER_DESKTOP_TUNNEL_PORT=""

# Run a socat tunnel container for pushing images from Docker Desktop VM to the forwarded port on the host.
run_docker_desktop_tunnel() {
    local host_port="$1"
    local output

    DOCKER_DESKTOP_TUNNEL_CONTAINER="docker-pussh-tunnel-$$"
    for _ in {1..10}; do
        DOCKER_DESKTOP_TUNNEL_PORT=$(random_port)

        if output=$(docker run -d --rm \
            --name "${DOCKER_DESKTOP_TUNNEL_CONTAINER}" \
            -p "127.0.0.1:${DOCKER_DESKTOP_TUNNEL_PORT}:5000" \
            alpine/socat \
            TCP-LISTEN:5000,fork,reuseaddr \
            "TCP-CONNECT:host.docker.internal:${host_port}" 2>&1);
        then
            return 0
        fi

        # Remove the container that failed to start if it was created.
        docker rm -f "${DOCKER_DESKTOP_TUNNEL_CONTAINER}" >/dev/null 2>&1 || true
        # Check if error is due to port binding.
        if ! echo "${output}" | grep -q --ignore-case "bind.*${DOCKER_DESKTOP_TUNNEL_PORT}"; then
            error "Failed to create a tunnel from Docker Desktop VM to localhost:${host_port}:\n${output}"
        fi
    done

    error "Failed to create a tunnel from Docker Desktop VM to localhost:${host_port}:\n${output}"
}

DOCKER_PLATFORM=""
SSH_KEY=""
IMAGE=""
SSH_ADDRESS=""

# Skip 'pussh' if called as Docker CLI plugin.
if [[ "${1:-}" = "pussh" ]]; then
    shift
fi

# Parse options and arguments.
help_command="Run 'docker pussh --help' for usage information."
while [[ $# -gt 0 ]]; do
    case "$1" in
        -i|--ssh-key)
            if [[ -z "${2:-}" ]]; then
                error "-i/--ssh-key option requires an argument.\n${help_command}"
            fi
            SSH_KEY="$2"
            shift 2
            ;;
        --no-host-key-check)
            SSH_STRICT_HOST_KEY_CHECKING="no"
            shift
            ;;
        --platform)
            if [[ -z "${2:-}" ]]; then
                error "--platform option requires an argument.\n${help_command}"
            fi
            DOCKER_PLATFORM="$2"
            shift 2
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        -v|--version)
            echo "docker-pussh, version ${VERSION}"
            echo "unregistry image: ${UNREGISTRY_IMAGE}"
            exit 0
            ;;
        -*)
            error "Unknown option: $1\n${help_command}"
            ;;
        *)
            # First non-option argument is the image.
            if [[ -z "${IMAGE}" ]]; then
                IMAGE="$1"
            # Second non-option argument is the SSH address.
            elif [[ -z "${SSH_ADDRESS}" ]]; then
                SSH_ADDRESS="$1"
            else
                error "Too many arguments.\n${help_command}"
            fi
            shift
            ;;
    esac
done

# Validate required arguments.
if [[ -z "${IMAGE}" ]] || [[ -z "${SSH_ADDRESS}" ]]; then
    error "IMAGE and HOST are required.\n${help_command}"
fi
# Validate SSH key file exists if provided.
if [[ -n "${SSH_KEY}" ]] && [[ ! -f "${SSH_KEY}" ]]; then
    error "SSH key file not found: ${SSH_KEY}"
fi

# Replace colon in the registry part of the image name with dash to make it a valid Docker image name component.
# For example, "localhost:5000/myimage" -> "localhost-5000/myimage"
normalise_registry_port() {
    local image=$1

    # Check if there's a port in registry address (colon followed by digits before the first slash).
    if [[ "${image}" =~ ^([^/]+):([0-9]+)(/.*)$ ]]; then
        # Replace the colon with a dash.
        echo "${BASH_REMATCH[1]}-${BASH_REMATCH[2]}${BASH_REMATCH[3]}"
    else
        # No port found, return as-is.
        echo "${image}"
    fi
}

# Clean up resources on exit, including on errors. Uses global variables to determine what needs to be cleaned up.
cleanup() {
    local exit_code=$?

    if [[ "${exit_code}" -ne 0 ]]; then
        warning "Cleaning up after error..."
    fi

    # Remove Docker Desktop tunnel container if exists.
    if [[ -n "${DOCKER_DESKTOP_TUNNEL_CONTAINER}" ]]; then
        docker rm -f "${DOCKER_DESKTOP_TUNNEL_CONTAINER}" >/dev/null 2>&1 || true
    fi

    # Clean up the temporary registry image tag.
    if [[ -n "${REGISTRY_IMAGE:-}" ]]; then
        docker rmi "${REGISTRY_IMAGE}" >/dev/null 2>&1 || true
    fi

    # Stop and remove unregistry container on remote host.
    if [[ -n "${UNREGISTRY_CONTAINER}" ]]; then
        # shellcheck disable=SC2029
        ssh "${SSH_ARGS[@]}" "${REMOTE_SUDO} docker rm -f ${UNREGISTRY_CONTAINER}" >/dev/null 2>&1 || true
    fi

    # Terminate the shared SSH connection if it was established.
    if [[ ${#SSH_ARGS[@]} -ne 0 ]]; then
        ssh "${SSH_ARGS[@]}" -O exit 2>/dev/null || true
    fi
}
trap cleanup EXIT

info "Connecting to ${SSH_ADDRESS}..."
ssh_remote "${SSH_ADDRESS}"
check_remote_docker

info "Starting unregistry container on remote host..."
run_unregistry
success "Unregistry is listening localhost:${UNREGISTRY_PORT} on remote host."

# Forward random local port to remote unregistry port through established SSH connection.
LOCAL_PORT=$(forward_port "${UNREGISTRY_PORT}")
success "Forwarded localhost:${LOCAL_PORT} to unregistry over SSH connection."

PUSH_PORT=${LOCAL_PORT}
# Handle virtualized Docker on macOS (e.g., Docker Desktop, Colima, etc.)
# shellcheck disable=SC2310
if is_additional_tunneling_needed; then
    info "Detected virtualised Docker locally, creating additional tunnel to localhost:${LOCAL_PORT}..."
    run_docker_desktop_tunnel "${LOCAL_PORT}"
    PUSH_PORT=${DOCKER_DESKTOP_TUNNEL_PORT}
    success "Additional tunnel created: localhost:${PUSH_PORT} → localhost:${LOCAL_PORT}"
fi

REMOTE_IMAGE=$(normalise_registry_port "${IMAGE}")
# Tag and push the image to unregistry through the forwarded port.
REGISTRY_IMAGE="localhost:${PUSH_PORT}/${REMOTE_IMAGE}"
docker tag "${IMAGE}" "${REGISTRY_IMAGE}"
info "Pushing ${REGISTRY_IMAGE} to unregistry..."

DOCKER_PUSH_OPTS=()
if [[ -n "${DOCKER_PLATFORM}" ]]; then
    DOCKER_PUSH_OPTS+=("--platform" "${DOCKER_PLATFORM}")
fi

# Try push with retry logic for connection issues
PUSH_RETRY_COUNT=3
PUSH_SUCCESS=false
PUSH_SLEEP_INTERVAL=3

for attempt in $(seq 1 "${PUSH_RETRY_COUNT}"); do
    # That DOCKER_PUSH_OPTS expansion is needed to avoid issues with empty array expansion in older bash versions.
    if docker push ${DOCKER_PUSH_OPTS[@]+"${DOCKER_PUSH_OPTS[@]}"} "${REGISTRY_IMAGE}"; then
        PUSH_SUCCESS=true
        break
    else
        if [[ "${attempt}" -lt "${PUSH_RETRY_COUNT}" ]]; then
            warning "Push attempt ${attempt} failed, retrying in ${PUSH_SLEEP_INTERVAL} seconds..."
            sleep "${PUSH_SLEEP_INTERVAL}"
        fi
    fi
done

if [[ "${PUSH_SUCCESS}" = false ]]; then
    error "Failed to push image after ${PUSH_RETRY_COUNT} attempts."
fi

REMOTE_RETAG_IMAGE=""
if [[ "${REMOTE_IMAGE}" != "${IMAGE}" ]]; then
    REMOTE_RETAG_IMAGE="${REMOTE_IMAGE}"
fi

# Pull image from unregistry if remote Docker doesn't use containerd image store.
# shellcheck disable=SC2029
if ! ssh "${SSH_ARGS[@]}" "${REMOTE_SUDO} docker info -f '{{ .DriverStatus }}' | grep -q 'containerd.snapshotter'"; then
    info "Remote Docker doesn't use containerd image store. Pulling image from unregistry..."
    REMOTE_REGISTRY_IMAGE="localhost:${UNREGISTRY_PORT}/${REMOTE_IMAGE}"
    if ! ssh "${SSH_ARGS[@]}" "${REMOTE_SUDO} docker pull ${REMOTE_REGISTRY_IMAGE}"; then
        error "Failed to pull image from unregistry on remote host."
    fi
    REMOTE_RETAG_IMAGE="${REMOTE_REGISTRY_IMAGE}"
fi

# Retag the image to the original name if needed and remove the temporary tag.
if [[ -n "${REMOTE_RETAG_IMAGE}" ]]; then
    # shellcheck disable=SC2029
    if ! ssh "${SSH_ARGS[@]}" "${REMOTE_SUDO} docker tag ${REMOTE_RETAG_IMAGE} ${IMAGE}"; then
        error "Failed to retag image on remote host ${REMOTE_RETAG_IMAGE} → ${IMAGE}"
    fi
    # shellcheck disable=SC2029
    ssh "${SSH_ARGS[@]}" "${REMOTE_SUDO} docker rmi ${REMOTE_RETAG_IMAGE}" >/dev/null || true
    success "Retagged image on remote host ${REMOTE_RETAG_IMAGE} → ${IMAGE}"
fi

info "Removing unregistry container on remote host..."
# shellcheck disable=SC2029
ssh "${SSH_ARGS[@]}" "${REMOTE_SUDO} docker rm -f ${UNREGISTRY_CONTAINER}" >/dev/null || true

success "Successfully pushed ${BOLD}${IMAGE}${RST} to ${BOLD}${SSH_ADDRESS}${RST}"