#!/usr/bin/env bash

set -euo pipefail

PROGRAM="avaron"
OWNER="${AVARON_GITHUB_OWNER:-AvaronInc}"
REPO="${AVARON_GITHUB_REPO:-core-cli}"
DEFAULT_BACKEND_URL="https://aim.avaron.ai"
DEFAULT_INSTALL_DIR="/usr/local/bin"
DEFAULT_RELEASE_BASE="https://github.com/${OWNER}/${REPO}/releases/download"

BACKEND_URL="${AVARON_BACKEND_URL:-$DEFAULT_BACKEND_URL}"
TOKEN="${AVARON_REGISTRATION_TOKEN:-}"
VERSION="${AVARON_VERSION:-latest}"
INSTALL_DIR="${AVARON_INSTALL_DIR:-$DEFAULT_INSTALL_DIR}"
DOWNLOAD_BASE_URL="${AVARON_DOWNLOAD_BASE_URL:-$DEFAULT_RELEASE_BASE}"
ENABLE_SERVICE=0
SKIP_ENROLL=0
RUN_AS_USER="${SUDO_USER:-$(id -un)}"

log() {
    printf '[%s] %s\n' "$(date '+%Y-%m-%d %H:%M:%S')" "$*"
}

fail() {
    log "ERROR: $*"
    exit 1
}

usage() {
    cat <<EOF
Install the Avaron CLI and optionally enroll this node.

Usage:
  curl -fsSL https://install.avaron.ai | bash -s -- [options]

Options:
  --backend-url URL      Backend URL to enroll against (default: ${DEFAULT_BACKEND_URL})
  --token TOKEN          Registration token from the Avaron dashboard
  --version VERSION      Release tag to install, or 'latest' (default: latest)
  --install-dir DIR      Binary install directory (default: ${DEFAULT_INSTALL_DIR})
  --enable-service       Create and enable a systemd service after install
  --skip-enroll          Install only; do not run 'avaron enroll'
  --help                 Show this help text

Environment overrides:
  AVARON_GITHUB_OWNER
  AVARON_GITHUB_REPO
  AVARON_DOWNLOAD_BASE_URL
  AVARON_BACKEND_URL
  AVARON_REGISTRATION_TOKEN
  AVARON_INSTALL_DIR
  AVARON_VERSION
EOF
}

need_cmd() {
    command -v "$1" >/dev/null 2>&1 || fail "required command not found: $1"
}

detect_os() {
    local os
    os="$(uname -s | tr '[:upper:]' '[:lower:]')"
    case "$os" in
        linux|darwin)
            printf '%s\n' "$os"
            ;;
        *)
            fail "unsupported operating system: $os"
            ;;
    esac
}

detect_arch() {
    local arch
    arch="$(uname -m)"
    case "$arch" in
        x86_64|amd64)
            printf 'amd64\n'
            ;;
        arm64|aarch64)
            printf 'arm64\n'
            ;;
        *)
            fail "unsupported architecture: $arch"
            ;;
    esac
}

resolve_version() {
    if [ "$VERSION" != "latest" ]; then
        printf '%s\n' "$VERSION"
        return 0
    fi

    if [ "$DOWNLOAD_BASE_URL" != "$DEFAULT_RELEASE_BASE" ]; then
        fail "AVARON_DOWNLOAD_BASE_URL requires --version to be set explicitly"
    fi

    local latest_url
    latest_url="$(curl -fsSLo /dev/null -w '%{url_effective}' "https://github.com/${OWNER}/${REPO}/releases/latest" || true)"
    [ -n "$latest_url" ] || fail "unable to resolve latest release version"
    printf '%s\n' "${latest_url##*/}"
}

asset_names() {
    local version="$1"
    local os="$2"
    local arch="$3"
    cat <<EOF
${PROGRAM}_${version#v}_${os}_${arch}.tar.gz
${PROGRAM}_${version#v}_${os}-${arch}.zip
${PROGRAM}_${version#v}_${os}_${arch}.zip
EOF
}

asset_url() {
    local version="$1"
    local asset="$2"

    if [ "$DOWNLOAD_BASE_URL" = "$DEFAULT_RELEASE_BASE" ]; then
        printf '%s/%s/%s\n' "$DOWNLOAD_BASE_URL" "$version" "$asset"
        return 0
    fi

    printf '%s/%s/%s\n' "${DOWNLOAD_BASE_URL%/}" "$version" "$asset"
}

as_root() {
    if [ "$(id -u)" -eq 0 ]; then
        "$@"
    elif command -v sudo >/dev/null 2>&1; then
        sudo "$@"
    else
        fail "this step requires root privileges and sudo is not installed"
    fi
}

run_for_install_user() {
    if [ "$(id -u)" -eq 0 ] && [ "$RUN_AS_USER" != "root" ]; then
        su - "$RUN_AS_USER" -c "$*"
    else
        eval "$@"
    fi
}

install_binary() {
    local version="$1"
    local os="$2"
    local arch="$3"
    local url tmp_dir archive_path asset downloaded_asset
    tmp_dir="$(mktemp -d)"
    trap 'rm -rf "$tmp_dir"' EXIT

    log "Downloading ${PROGRAM} ${version} for ${os}/${arch}"
    downloaded_asset=""
    while IFS= read -r asset; do
        [ -n "$asset" ] || continue
        url="$(asset_url "$version" "$asset")"
        archive_path="${tmp_dir}/${asset}"
        if curl -fsSL "$url" -o "$archive_path"; then
            downloaded_asset="$asset"
            break
        fi
    done <<EOF
$(asset_names "$version" "$os" "$arch")
EOF

    [ -n "$downloaded_asset" ] || fail "download failed for ${os}/${arch} from ${DOWNLOAD_BASE_URL}"

    case "$downloaded_asset" in
        *.tar.gz)
            tar -xzf "$archive_path" -C "$tmp_dir"
            ;;
        *.zip)
            unzip -j "$archive_path" -d "$tmp_dir" >/dev/null
            ;;
        *)
            fail "unsupported archive format: $downloaded_asset"
            ;;
    esac

    [ -f "${tmp_dir}/${PROGRAM}" ] || fail "archive did not contain ${PROGRAM}"

    chmod +x "${tmp_dir}/${PROGRAM}"
    as_root mkdir -p "$INSTALL_DIR"
    as_root install "${tmp_dir}/${PROGRAM}" "${INSTALL_DIR}/${PROGRAM}"

    log "Installed ${PROGRAM} to ${INSTALL_DIR}/${PROGRAM}"
    "${INSTALL_DIR}/${PROGRAM}" version || fail "installed binary failed version check"
}

write_service_file() {
    local service_path
    service_path="/etc/systemd/system/avaron.service"

    [ "$(detect_os)" = "linux" ] || fail "--enable-service is only supported on Linux"
    command -v systemctl >/dev/null 2>&1 || fail "systemctl not found"

    local tmp_dir service_file
    tmp_dir="$(mktemp -d)"
    service_file="${tmp_dir}/avaron.service"
    trap 'rm -rf "$tmp_dir"' EXIT

    cat >"$service_file" <<EOF
[Unit]
Description=Avaron Execution Daemon
After=network-online.target
Wants=network-online.target

[Service]
ExecStart=${INSTALL_DIR}/${PROGRAM} daemon start
Restart=on-failure
RestartSec=5s
User=${RUN_AS_USER}

[Install]
WantedBy=multi-user.target
EOF

    as_root install -m 0644 "$service_file" "$service_path"
    as_root systemctl daemon-reload
    as_root systemctl enable avaron
    as_root systemctl restart avaron

    log "Systemd service enabled: ${service_path}"
}

run_enroll() {
    [ -n "$TOKEN" ] || fail "--token is required unless --skip-enroll is used"
    run_for_install_user \
        "\"${INSTALL_DIR}/${PROGRAM}\" enroll register --backend-url '${BACKEND_URL}' --token '${TOKEN}'"
}

while [ $# -gt 0 ]; do
    case "$1" in
        --backend-url)
            BACKEND_URL="$2"
            shift 2
            ;;
        --token)
            TOKEN="$2"
            shift 2
            ;;
        --version)
            VERSION="$2"
            shift 2
            ;;
        --install-dir)
            INSTALL_DIR="$2"
            shift 2
            ;;
        --enable-service)
            ENABLE_SERVICE=1
            shift
            ;;
        --skip-enroll)
            SKIP_ENROLL=1
            shift
            ;;
        --help|-h)
            usage
            exit 0
            ;;
        *)
            fail "unknown argument: $1"
            ;;
    esac
done

need_cmd curl
need_cmd tar
need_cmd unzip
need_cmd uname

OS="$(detect_os)"
ARCH="$(detect_arch)"
RESOLVED_VERSION="$(resolve_version)"

install_binary "$RESOLVED_VERSION" "$OS" "$ARCH"

if [ "$SKIP_ENROLL" -eq 0 ]; then
    log "Enrolling node with ${BACKEND_URL}"
    run_enroll
fi

if [ "$ENABLE_SERVICE" -eq 1 ]; then
    write_service_file
fi

log "Done."