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

# ------------------------------
# Constants and style
# ------------------------------
readonly DEFAULT_MANIFEST_URL="https://get.laviq.io/releases/stable/manifest.json"
readonly DEFAULT_BIN_PATH="/opt/laviq/laviq"
readonly DEFAULT_BPF_PATH="/opt/laviq/laviq.bpf.o"
readonly DEFAULT_PACKAGE_TYPE="tar.gz"
readonly DEFAULT_AGENT_ID_PATH="/var/lib/laviq/agent-id"

readonly COLOR_RED='\033[0;31m'
readonly COLOR_GREEN='\033[0;32m'
readonly COLOR_YELLOW='\033[1;33m'
readonly COLOR_BLUE='\033[0;34m'
readonly COLOR_RESET='\033[0m'

STEP_COUNTER=0

# ------------------------------
# Logging and step helpers
# ------------------------------
log_info() {
  echo -e "${COLOR_BLUE}$*${COLOR_RESET}"
}

log_warn() {
  echo -e "${COLOR_YELLOW}$*${COLOR_RESET}"
}

log_ok() {
  echo -e "${COLOR_GREEN}✓ $*${COLOR_RESET}"
}

log_fail() {
  echo -e "${COLOR_RED}✗ $*${COLOR_RESET}" >&2
}

run_step() {
  local title="$1"
  shift
  STEP_COUNTER=$((STEP_COUNTER + 1))
  echo
  echo "Step ${STEP_COUNTER}: ${title}"
  if "$@"; then
    log_ok "${title}"
  else
    log_fail "${title}"
    exit 1
  fi
}

# ------------------------------
# Utility functions
# ------------------------------
require_cmd() {
  command -v "$1" >/dev/null 2>&1 || {
    log_fail "missing required command: $1"
    return 1
  }
}

preflight_checks() {
  require_cmd tar
  require_cmd curl
  require_cmd sha256sum
  require_cmd python3
}

detect_platform_key() {
  local arch
  arch="$(uname -m | tr '[:upper:]' '[:lower:]')"
  if [[ "${arch}" == "aarch64" || "${arch}" == "arm64" ]]; then
    echo "linux-arm64"
  else
    echo "linux-amd64"
  fi
}

detect_distro() {
  if [[ -r /etc/os-release ]]; then
    . /etc/os-release
    echo "${PRETTY_NAME:-Linux}"
  else
    uname -a
  fi
}

write_package_env() {
  local bin_path="$1"
  local bpf_path="$2"
  if [[ -n "${PACKAGE_ENV_FILE:-}" ]]; then
    printf 'export LAVIQ_BIN_PATH=%s\nexport LAVIQ_BPF_OBJECT_PATH=%s\n' \
      "${bin_path}" "${bpf_path}" > "${PACKAGE_ENV_FILE}"
    log_info "Wrote package env: ${PACKAGE_ENV_FILE}"
  fi
}

extract_version_from_manifest() {
  local manifest_file="$1"
  python3 - "${manifest_file}" <<'PY'
import json, sys
with open(sys.argv[1], 'r', encoding='utf-8') as f:
    print((json.load(f).get('version') or 'unknown'))
PY
}

resolve_manifest_artifact() {
  local manifest_file="$1"
  local platform_key="$2"
  python3 - "${manifest_file}" "${platform_key}" <<'PY'
import json, sys
with open(sys.argv[1], 'r', encoding='utf-8') as f:
    data = json.load(f)
art = (data.get('artifacts') or {}).get(sys.argv[2])
if not art:
    print('ERROR:missing_artifact')
    sys.exit(2)
print(art.get('url', ''))
print(art.get('sha256', ''))
PY
}

download_to_file() {
  local url="$1"
  local out_file="$2"
  curl -fsS "${url}" -o "${out_file}"
  log_info "Downloaded OK: ${url}"
}

verify_sha256() {
  local file_path="$1"
  local expected_sha="$2"
  local actual_sha
  actual_sha="$(sha256sum "${file_path}" | awk '{print $1}')"
  [[ "${actual_sha}" == "${expected_sha}" ]]
  log_info "Checksum OK: ${actual_sha}"
}

install_from_tar() {
  local tar_pkg="$1"
  local bin_path="$2"
  local bpf_path="$3"

  [[ -f "${tar_pkg}" ]] || {
    log_fail "artifact tar not found: ${tar_pkg}"
    return 1
  }

  tar -xzf "${tar_pkg}" -C /
  [[ -x "${bin_path}" ]]
  [[ -f "${bpf_path}" ]]
  log_info "Installed OK: ${bin_path} and ${bpf_path}"

  if [[ "${LOCAL_PACKAGE_GUARD:-1}" == "1" ]]; then
    local expected_bin_sha expected_bpf_sha actual_bin_sha actual_bpf_sha
    expected_bin_sha="$(tar -xOzf "${tar_pkg}" ./opt/laviq/laviq | sha256sum | awk '{print $1}')"
    expected_bpf_sha="$(tar -xOzf "${tar_pkg}" ./opt/laviq/laviq.bpf.o | sha256sum | awk '{print $1}')"
    actual_bin_sha="$(sha256sum "${bin_path}" | awk '{print $1}')"
    actual_bpf_sha="$(sha256sum "${bpf_path}" | awk '{print $1}')"
    [[ "${expected_bin_sha}" == "${actual_bin_sha}" ]]
    [[ "${expected_bpf_sha}" == "${actual_bpf_sha}" ]]
    log_info "Package guard OK: installed binaries match tar payload"
  fi
}

ensure_agent_identity_file() {
  local agent_id_path="${1}"
  local agent_id_dir
  agent_id_dir="$(dirname "${agent_id_path}")"
  mkdir -p "${agent_id_dir}"

  if [[ -f "${agent_id_path}" ]]; then
    python3 - "${agent_id_path}" <<'PY'
import pathlib, sys, uuid
value = pathlib.Path(sys.argv[1]).read_text(encoding='utf-8').strip()
if not value:
    raise SystemExit(2)
uuid.UUID(value)
PY
    log_info "Agent identity file already present: ${agent_id_path}"
    chmod 600 "${agent_id_path}"
    return 0
  fi

  python3 - "${agent_id_path}" <<'PY'
import pathlib, sys, uuid
path = pathlib.Path(sys.argv[1])
path.write_text(f"{uuid.uuid4()}\n", encoding='utf-8')
PY
  chmod 600 "${agent_id_path}"
  log_info "Initialized agent identity file: ${agent_id_path}"
}

enroll_installed_agent() {
  local bin_path="$1"
  local api_key="${API_KEY:-}"

  if [[ -z "${api_key}" && -r /dev/tty ]]; then
    read -r -p "Enter your API key: " api_key </dev/tty
  fi

  api_key="$(echo "${api_key}" | tr -d '\r' | xargs)"
  [[ -n "${api_key}" ]] || {
    log_fail "API_KEY is required"
    return 1
  }

  "${bin_path}" enroll --help >/dev/null 2>&1 || {
    log_fail "installed laviq binary does not support the enroll subcommand"
    return 1
  }

  if ! "${bin_path}" enroll --api-key "${api_key}"; then
    log_fail "Enrollment failed"
    return 1
  fi
  log_info "Enrollment OK"
}

# ------------------------------
# Main flow (linear, readable)
# ------------------------------
echo "Laviq Agent Installer (laviq.io)"
echo "─────────────────────"

OS_NAME="$(uname -s 2>/dev/null || echo unknown)"
DISTRO="$(detect_distro)"
PLATFORM_KEY="$(detect_platform_key)"
PACKAGE_TYPE="${PACKAGE_TYPE:-${DEFAULT_PACKAGE_TYPE}}"
LAVIQ_BIN_PATH="${LAVIQ_BIN_PATH:-${DEFAULT_BIN_PATH}}"
LAVIQ_BPF_OBJECT_PATH="${LAVIQ_BPF_OBJECT_PATH:-${DEFAULT_BPF_PATH}}"
LAVIQ_AGENT_ID_PATH="${LAVIQ_AGENT_ID_PATH:-${DEFAULT_AGENT_ID_PATH}}"

if [[ "${OS_NAME}" != "Linux" ]]; then
  log_fail "installer currently supports Linux only"
  exit 1
fi

if [[ "${EUID}" -ne 0 ]]; then
  log_fail "Root access is required to load the eBPF kernel module"
  echo "Run: curl -sSL https://get.laviq.io/install.sh | sudo bash"
  exit 1
fi

echo "Distro: ${DISTRO}"
echo "Platform: ${PLATFORM_KEY}"
echo "Package type: ${PACKAGE_TYPE}"

tmp_dir="$(mktemp -d)"
trap 'rm -rf "${tmp_dir}"' EXIT
manifest_file="${tmp_dir}/manifest.json"
artifact_file="${tmp_dir}/artifact.tar.gz"
artifact_url=""
artifact_sha=""
release_version="unknown"

# 1) Preflight checks
run_step "Preflight checks" preflight_checks

# 2) Resolve artifact source
if [[ -n "${INSTALL_ARTIFACT_TAR:-}" ]]; then
  artifact_file="${INSTALL_ARTIFACT_TAR}"
  release_version="local-artifact"
  log_info "Using local artifact: ${artifact_file}"
else
  manifest_url="${MANIFEST_URL:-${DEFAULT_MANIFEST_URL}}"
  run_step "Fetch manifest" download_to_file "${manifest_url}" "${manifest_file}"

  release_version="$(extract_version_from_manifest "${manifest_file}")"
  py_out="$(resolve_manifest_artifact "${manifest_file}" "${PLATFORM_KEY}")"
  artifact_url="$(echo "${py_out}" | sed -n '1p')"
  artifact_sha="$(echo "${py_out}" | sed -n '2p')"
  [[ "${artifact_url}" == ERROR:* ]] && {
    log_fail "manifest does not contain artifact for ${PLATFORM_KEY}"
    exit 1
  }

  # 3) Download artifact from manifest
  run_step "Download artifact" download_to_file "${artifact_url}" "${artifact_file}"

  # 4) Verify checksum from manifest
  run_step "Verify checksum" verify_sha256 "${artifact_file}" "${artifact_sha}"
fi

echo "Version: ${release_version}"

# 5) Optional validation-only mode
if [[ "${INSTALL_VALIDATE_ONLY:-0}" == "1" ]]; then
  log_ok "Validation-only mode complete"
  exit 0
fi

# 6) Install files from tar payload
run_step "Install package payload" install_from_tar "${artifact_file}" "${LAVIQ_BIN_PATH}" "${LAVIQ_BPF_OBJECT_PATH}"
write_package_env "${LAVIQ_BIN_PATH}" "${LAVIQ_BPF_OBJECT_PATH}"

# 7) Ensure persisted agent identity exists
run_step "Initialize agent identity" ensure_agent_identity_file "${LAVIQ_AGENT_ID_PATH}"

# 8) Enroll agent (required)
run_step "Enroll agent" enroll_installed_agent "${LAVIQ_BIN_PATH}"

echo
log_ok "Installation completed"
echo "Binary: ${LAVIQ_BIN_PATH}"
echo "BPF object: ${LAVIQ_BPF_OBJECT_PATH}"
echo "Version: ${release_version}"