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

### --- OS CHECK --- ###
if [[ $OSTYPE != linux-gnu* ]]; then
  echo "[ERROR] This script supports Linux only. Aborting."
  exit 1
fi

# --- Auto-save if running from a pipe ---
if [ -p /dev/stdin ]; then
  TMP_SCRIPT="/tmp/install-edge.sh"
  echo "[INFO] Detected script running from a pipe. Saving to $TMP_SCRIPT..."
  cat >"$TMP_SCRIPT"
  chmod +x "$TMP_SCRIPT"

  echo "[INFO] Re-running saved script..."
  exec /usr/bin/env bash "$TMP_SCRIPT" "$@"
fi

### --- CONFIGURATION --- ###
IMAGE_NAME="ghcr.io/autonomy-logic/orchestrator-agent:latest"
NETMON_IMAGE_NAME="ghcr.io/autonomy-logic/autonomy-netmon:latest"
CONTAINER_NAME="orchestrator_agent"
NETMON_CONTAINER_NAME="autonomy_netmon"
SOURCE_DIR="/tmp/orchestrator-agent"
SHARED_VOLUME="orchestrator-shared"
SERVER_DNS="api.autonomylogic.com"
SERVER_URL="https://$SERVER_DNS"
GET_ID_URL="$SERVER_URL/orchestrators/id"
ENROLL_URL="$SERVER_URL/orchestrators/enroll"
MTLS_DIR="$HOME/.mtls"
KEY_PATH="$MTLS_DIR/client.key"
CRT_PATH="$MTLS_DIR/client.crt"
CSR_PATH="$MTLS_DIR/client.csr"
CSR_CONFIG_FILE="$MTLS_DIR/client.conf"
POLL_INTERVAL_SECONDS=3

### --- ARGUMENT PARSING --- ###
# Shared by the install and uninstall paths. Parsed before anything is created
# or removed so `--help` and a bad flag cost nothing.
MODE="install"
ASSUME_YES=false
KEEP_CERTS=false
KEEP_IMAGES=false
FORCE_LOCAL=false

usage() {
  cat <<USAGE
Usage: $(basename "$0") [OPTIONS]

Installs or upgrades the Autonomy Edge orchestrator agent. With --uninstall,
removes the agent and every resource it manages from this device.

Options:
  --uninstall      Remove the orchestrator agent, its vPLC runtime containers,
                   orchestrator networks, shared volume, mTLS identity and
                   pulled images
  -y, --yes        Skip the uninstall confirmation prompt (required when the
                   script is piped and no terminal is attached)
  --keep-certs     Uninstall without deleting the mTLS client certificate/key
  --keep-images    Uninstall without deleting the pulled Docker images
  --force-local    Uninstall with host-side cleanup only, skipping the agent's
                   own teardown (use when the agent is stopped or wedged)
  -h, --help       Show this help and exit

Examples:
  $(basename "$0")
  $(basename "$0") --uninstall
  curl -fsSL https://getedge.me | bash -s -- --uninstall --yes
USAGE
}

# Every flag is a boolean, so read them without consuming "$@": check_root
# below re-execs the script under sudo with "$@", and shifting here would drop
# the flags on that hand-off -- turning `--uninstall` into a fresh install.
for arg in "$@"; do
  case "$arg" in
  --uninstall) MODE="uninstall" ;;
  -y | --yes) ASSUME_YES=true ;;
  --keep-certs) KEEP_CERTS=true ;;
  --keep-images) KEEP_IMAGES=true ;;
  --force-local) FORCE_LOCAL=true ;;
  -h | --help)
    usage
    exit 0
    ;;
  *)
    echo "[ERROR] Unknown option: $arg"
    echo
    usage
    exit 1
    ;;
  esac
done

if [ "$MODE" = "install" ]; then
  if [ "$KEEP_CERTS" = true ] || [ "$KEEP_IMAGES" = true ] || [ "$FORCE_LOCAL" = true ]; then
    echo "[WARN] --keep-certs, --keep-images and --force-local only apply to --uninstall; ignoring."
  fi
fi

# mktemp the response files so a malicious local user can't pre-create them as
# symlinks before sudo elevates this script.
ENROLL_RESP_FILE=$(mktemp -t orchestrator-enroll-resp.XXXXXX)
PENDING_RESP_FILE=$(mktemp -t orchestrator-pending-resp.XXXXXX)

# When true, the EXIT trap will remove the freshly-generated mTLS material.
# Set to true after we create $MTLS_DIR; cleared once validation succeeds, so a
# script that aborts before validation does not leave a dead identity behind.
CLEANUP_CERTS_ON_FAILURE=false

# Cleanup function for trap - ensures temp files are removed on exit and
# (when applicable) drops the unvalidated mTLS material.
cleanup_temp_files() {
  rm -f "${CSR_PATH:-}" "${CSR_CONFIG_FILE:-}" \
    "${ENROLL_RESP_FILE:-}" "${PENDING_RESP_FILE:-}" 2>/dev/null || true

  if [ "$CLEANUP_CERTS_ON_FAILURE" = true ]; then
    echo "[INFO] Removing unvalidated mTLS material from $MTLS_DIR"
    rm -f "$KEY_PATH" "$CRT_PATH" 2>/dev/null || true
    # Non-recursive on purpose: if the user has unrelated files in ~/.mtls we
    # leave them alone. Do NOT change this to rm -rf.
    rmdir "$MTLS_DIR" 2>/dev/null || true
  fi
}
trap cleanup_temp_files EXIT

# Check for root privileges
check_root() {
  if [[ $EUID -ne 0 ]]; then
    echo "[INFO] Root privileges are required. Trying to elevate with sudo..."
    # Re-run the script with sudo, passing all original arguments
    exec sudo /usr/bin/env bash "$0" "$@"
    # exec replaces the current shell with the new command, so the rest of the script continues as root
  fi
}

# Make sure we are root before proceeding
check_root "$@"

### --- COLOR DETECTION --- ###
if [ -t 1 ] && command -v tput >/dev/null && [ "$(tput colors 2>/dev/null)" -ge 8 ]; then
  GREEN="$(tput setaf 2)"
  CYAN="$(tput setaf 6)"
  YELLOW="$(tput setaf 3)"
  RED="$(tput setaf 1)"
  GRAY="$(tput setaf 8)"
  BOLD="$(tput bold)"
  RESET="$(tput sgr0)"
else
  GREEN=""
  CYAN=""
  YELLOW=""
  RED=""
  GRAY=""
  BOLD=""
  RESET=""
fi

### --- UNINSTALL --- ###
# The uninstall path mirrors the cloud-side `delete_orchestrator` teardown
# (src/use_cases/docker_manager/selfdestruct.py) and runs in two stages:
#
#   1. Agent-driven (preferred): exec the agent's own uninstall entry point so
#      the real use case runs inside the container -- the persisted client
#      registry as the source of truth, netmon-mediated dedicated-NIC and
#      Proxy ARP cleanup, and the same ordering the cloud command uses.
#   2. Host-side sweep (fallback): plain docker/ip/iptables commands, used when
#      the agent is absent, stopped, wedged, or predates the entry point --
#      i.e. a partially installed or broken orchestrator.
#
# The sweep always runs, even after a successful agent teardown: it is
# idempotent, and some resources cannot be removed by the agent from inside
# itself (the shared volume it mounts, its own image, the host's mTLS material,
# leftover iptables rules).

UNINSTALL_ENTRYPOINTS="src/uninstall.pyc src/uninstall.py"
UPGRADER_CONTAINER_NAME="orchestrator_upgrader"
AGENT_IMAGE_PREFIX="ghcr.io/autonomy-logic/orchestrator-agent"
NETMON_IMAGE_PREFIX="ghcr.io/autonomy-logic/autonomy-netmon"
RUNTIME_IMAGE_PREFIX="ghcr.io/autonomy-logic/openplc-runtime"

# Managed-resource name patterns, kept in sync with selfdestruct.py.
UUID_PATTERN='[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}'
INTERNAL_NET_PATTERN="^${UUID_PATTERN}_internal\$"
MACVLAN_NET_PATTERN='^macvlan_[a-zA-Z0-9]+_[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+_[0-9]+$'

# How long to wait for the agent to finish removing itself before falling back.
# Overridable so tests (and operators on very large fleets) can retune it.
AGENT_TEARDOWN_TIMEOUT="${AGENT_TEARDOWN_TIMEOUT:-180}"

# Asks netmon to drop every Proxy ARP veth and neighbour entry, exactly as
# selfdestruct.py's _cleanup_proxy_arp_veths() does over the same socket.
# netmon greets every new client with network_discovery/device_discovery pushes,
# so the first line on the socket is not our answer: skip pushed events (they
# carry "type", never "success") until the command response arrives. Exiting
# non-zero without a confirmed response drops us to the host-side cleanup.
NETMON_PROXY_ARP_PY='
import json, socket, sys

try:
    s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    s.settimeout(15)
    s.connect("/var/orchestrator/netmon.sock")
    s.sendall(json.dumps({"command": "cleanup_all_proxy_arp"}).encode() + b"\n")

    buffer = b""
    response = None
    while response is None:
        chunk = s.recv(4096)
        if not chunk:
            break
        buffer += chunk
        while b"\n" in buffer:
            line, buffer = buffer.split(b"\n", 1)
            line = line.strip()
            if not line:
                continue
            try:
                message = json.loads(line.decode())
            except ValueError:
                continue
            if isinstance(message, dict) and "success" in message:
                response = message
                break
    s.close()

    if response is None:
        print("netmon did not answer cleanup_all_proxy_arp")
        sys.exit(1)
    print(json.dumps(response))
    sys.exit(0 if response.get("success") else 1)
except Exception as exc:
    print("netmon cleanup_all_proxy_arp failed: %s" % exc)
    sys.exit(1)
'

u_step() {
  echo
  echo -e "${BOLD}==> $1${RESET}"
}
u_info() { echo "    $1"; }
u_ok() { echo -e "    ${GREEN}[OK]${RESET} $1"; }
u_warn() { echo -e "    ${YELLOW}[WARN]${RESET} $1"; }

container_exists() {
  docker inspect --type container "$1" >/dev/null 2>&1
}

container_running() {
  [ "$(docker inspect --type container --format '{{.State.Running}}' "$1" 2>/dev/null || true)" = "true" ]
}

# Where the mTLS identity actually lives. Read from the agent's mount table
# rather than $HOME, because sudo changes $HOME and the certs may predate it.
# Must be called before the agent container is removed.
detect_mtls_dir() {
  local detected=""
  if container_exists "$CONTAINER_NAME"; then
    detected=$(docker inspect "$CONTAINER_NAME" \
      --format '{{range .Mounts}}{{if eq .Destination "/root/.mtls"}}{{.Source}}{{end}}{{end}}' 2>/dev/null || true)
  fi
  if [ -n "$detected" ]; then
    printf '%s' "$detected"
  else
    printf '%s' "$MTLS_DIR"
  fi
}

# Managed vPLC runtime containers, discovered from three independent angles so a
# broken install with partial state is still cleaned up:
#   1. the persisted client registry on the shared volume
#   2. containers attached to a <uuid>_internal control-plane network
#   3. containers running an openplc-runtime image
find_runtime_containers() {
  local candidates=""
  local volume_mount=""
  local net=""
  local name=""

  volume_mount=$(docker volume inspect "$SHARED_VOLUME" --format '{{.Mountpoint}}' 2>/dev/null || true)
  if [ -n "$volume_mount" ] && [ -f "$volume_mount/data/clients.json" ]; then
    candidates="$candidates
$(grep -oEi "$UUID_PATTERN" "$volume_mount/data/clients.json" 2>/dev/null || true)"
  fi

  for net in $(docker network ls --format '{{.Name}}' 2>/dev/null | grep -Ei "$INTERNAL_NET_PATTERN" || true); do
    candidates="$candidates
$(docker network inspect "$net" --format '{{range .Containers}}{{println .Name}}{{end}}' 2>/dev/null || true)"
  done

  candidates="$candidates
$(docker ps -a --format '{{.Names}} {{.Image}}' 2>/dev/null |
    awk -v prefix="$RUNTIME_IMAGE_PREFIX" 'index($2, prefix) == 1 { print $1 }' || true)"

  # Drop blanks, our own containers, and anything that no longer exists.
  for name in $(echo "$candidates" | sort -u); do
    case "$name" in
    "" | "$CONTAINER_NAME" | "$NETMON_CONTAINER_NAME" | "$UPGRADER_CONTAINER_NAME") continue ;;
    esac
    if container_exists "$name"; then
      echo "$name"
    fi
  done
}

remove_container() {
  local name="$1"
  if ! container_exists "$name"; then
    u_info "$name not present"
    return 0
  fi
  docker stop -t 10 "$name" >/dev/null 2>&1 || true
  if docker rm -f "$name" >/dev/null 2>&1; then
    u_ok "removed container $name"
  else
    u_warn "could not remove container $name"
  fi
}

# Remove orchestrator-created networks. Networks holding a container that is not
# ours are skipped, mirroring selfdestruct.py's refusal to disrupt unrelated
# applications; ours are force-disconnected first, like remove_internal_network().
remove_orchestrator_networks() {
  local nets=""
  local net=""
  local attached=""
  local name=""
  local foreign=""
  local removed=0
  local skipped=0

  nets=$(docker network ls --format '{{.Name}}' 2>/dev/null |
    grep -Ei "$INTERNAL_NET_PATTERN|$MACVLAN_NET_PATTERN" || true)

  if [ -z "$nets" ]; then
    u_info "no orchestrator networks found"
    return 0
  fi

  for net in $nets; do
    attached=$(docker network inspect "$net" --format '{{range .Containers}}{{println .Name}}{{end}}' 2>/dev/null || true)
    foreign=false
    for name in $attached; do
      case "$name" in
      "$CONTAINER_NAME" | "$NETMON_CONTAINER_NAME" | "$UPGRADER_CONTAINER_NAME") continue ;;
      esac
      # An unknown container on the network means something else is using it.
      if container_exists "$name"; then
        foreign=true
      fi
    done

    if [ "$foreign" = true ]; then
      u_warn "network $net still has non-orchestrator containers attached, skipping"
      skipped=$((skipped + 1))
      continue
    fi

    for name in $attached; do
      docker network disconnect -f "$net" "$name" >/dev/null 2>&1 || true
    done

    if docker network rm "$net" >/dev/null 2>&1; then
      u_ok "removed network $net"
      removed=$((removed + 1))
    else
      u_warn "could not remove network $net"
      skipped=$((skipped + 1))
    fi
  done

  u_info "networks removed: $removed, skipped: $skipped"
}

# Proxy ARP veths live in the host namespace. Ask netmon first (it created them
# and owns iproute2), then fall back to the host's own ip command.
cleanup_proxy_arp() {
  local output=""
  local iface=""
  local entry=""
  local addr=""
  local dev=""

  if container_running "$NETMON_CONTAINER_NAME"; then
    if output=$(docker exec "$NETMON_CONTAINER_NAME" python -c "$NETMON_PROXY_ARP_PY" 2>&1); then
      u_ok "netmon Proxy ARP cleanup: $output"
      return 0
    fi
    u_warn "netmon Proxy ARP cleanup failed, falling back to host cleanup"
  else
    u_info "netmon not running, cleaning Proxy ARP interfaces from the host"
  fi

  if ! command -v ip >/dev/null 2>&1; then
    u_warn "no 'ip' command on this host, skipping Proxy ARP cleanup"
    return 0
  fi

  # Neighbour entries first: "<ip> dev <iface> proxy". Only ours (veth-*).
  for entry in $(ip neighbor show proxy 2>/dev/null | awk '$3 ~ /^veth-/ { print $1 "|" $3 }' || true); do
    addr=${entry%%|*}
    dev=${entry##*|}
    ip neighbor del proxy "$addr" dev "$dev" >/dev/null 2>&1 || true
    u_ok "removed proxy ARP entry $addr dev $dev"
  done

  for iface in $(ip -o link show 2>/dev/null | sed -n 's/^[0-9]*: \([^:@]*\).*/\1/p' | grep '^veth-' || true); do
    if ip link del "$iface" >/dev/null 2>&1; then
      u_ok "removed Proxy ARP veth $iface"
    else
      u_warn "could not remove veth $iface"
    fi
  done
}

# netmon adds FORWARD ACCEPT rules per Proxy ARP veth and never removes them in
# bulk. With the veths and netmon gone the rules are dangling, so drop them.
cleanup_iptables_rules() {
  local rule=""
  local removed=0

  if ! command -v iptables >/dev/null 2>&1; then
    u_info "no 'iptables' command on this host, nothing to clean"
    return 0
  fi

  while IFS= read -r rule; do
    [ -n "$rule" ] || continue
    # shellcheck disable=SC2086 # $rule must word-split into iptables arguments
    if iptables $rule >/dev/null 2>&1; then
      removed=$((removed + 1))
    fi
  done <<EOF
$(iptables -S FORWARD 2>/dev/null | grep -E '(-i|-o) veth-' | sed 's/^-A /-D /' || true)
EOF

  u_info "iptables FORWARD rules removed: $removed"
}

remove_shared_volume() {
  if ! docker volume inspect "$SHARED_VOLUME" >/dev/null 2>&1; then
    u_info "volume $SHARED_VOLUME not present"
    return 0
  fi
  if docker volume rm -f "$SHARED_VOLUME" >/dev/null 2>&1; then
    u_ok "removed volume $SHARED_VOLUME"
  else
    u_warn "could not remove volume $SHARED_VOLUME (still in use?); try 'docker volume prune'"
  fi
}

remove_mtls_material() {
  local dir="$1"

  if [ "$KEEP_CERTS" = true ]; then
    u_info "keeping mTLS material in $dir (--keep-certs)"
    return 0
  fi

  if [ ! -d "$dir" ]; then
    u_info "no mTLS directory at $dir"
    return 0
  fi

  # Non-recursive on purpose: unrelated files the user keeps in ~/.mtls stay put.
  # Do NOT change this to rm -rf.
  rm -f "$dir/client.key" "$dir/client.crt" "$dir/client.csr" "$dir/client.conf" 2>/dev/null || true
  rmdir "$dir" 2>/dev/null || true
  u_ok "removed mTLS client certificate and key from $dir"
  u_info "this device's identity is gone; reinstalling enrolls a new orchestrator ID"
}

remove_images() {
  local prefix=""
  local image=""
  local removed=0

  if [ "$KEEP_IMAGES" = true ]; then
    u_info "keeping pulled images (--keep-images)"
    return 0
  fi

  for prefix in "$AGENT_IMAGE_PREFIX" "$NETMON_IMAGE_PREFIX" "$RUNTIME_IMAGE_PREFIX"; do
    for image in $(docker images --format '{{.Repository}}:{{.Tag}}' 2>/dev/null | grep "^${prefix}:" || true); do
      if docker rmi -f "$image" >/dev/null 2>&1; then
        u_ok "removed image $image"
        removed=$((removed + 1))
      else
        u_warn "could not remove image $image"
      fi
    done
  done

  u_info "images removed: $removed (run 'docker image prune' to drop untagged leftovers)"
}

remove_installer_leftovers() {
  if [ -d "$SOURCE_DIR/.git" ]; then
    rm -rf "$SOURCE_DIR" 2>/dev/null || true
    u_ok "removed local source clone $SOURCE_DIR"
  fi
  rm -f /tmp/install-edge.sh /tmp/install-edge-staging.sh 2>/dev/null || true
}

# Stage 1: let the agent tear itself down with the same code the cloud calls.
# Returns 0 only when the agent removed its own container.
run_agent_teardown() {
  local candidate=""
  local entrypoint=""
  local waited=0

  if ! container_running "$CONTAINER_NAME"; then
    if container_exists "$CONTAINER_NAME"; then
      u_info "agent container exists but is not running; host cleanup will handle it"
    else
      u_info "no agent container on this device; host cleanup will handle leftovers"
    fi
    return 1
  fi

  for candidate in $UNINSTALL_ENTRYPOINTS; do
    if docker exec "$CONTAINER_NAME" sh -c "test -f '$candidate'" >/dev/null 2>&1; then
      entrypoint="$candidate"
      break
    fi
  done

  if [ -z "$entrypoint" ]; then
    u_warn "this agent image predates the uninstall entry point; using host cleanup"
    return 1
  fi

  u_info "running the agent's own teardown: docker exec $CONTAINER_NAME python $entrypoint"
  echo
  # The agent removes its own container last, which kills this exec, so a
  # non-zero status here is expected and not an error signal.
  docker exec "$CONTAINER_NAME" python "$entrypoint" 2>&1 | sed 's/^/    /' || true
  echo

  while [ "$waited" -lt "$AGENT_TEARDOWN_TIMEOUT" ]; do
    if ! container_exists "$CONTAINER_NAME"; then
      u_ok "agent removed itself and the resources it manages"
      return 0
    fi
    sleep 1
    waited=$((waited + 1))
  done

  u_warn "agent did not remove itself within ${AGENT_TEARDOWN_TIMEOUT}s; using host cleanup"
  return 1
}

confirm_uninstall() {
  local answer=""

  if [ "$ASSUME_YES" = true ]; then
    return 0
  fi

  # A piped script leaves stdin at EOF, so prefer the terminal when there is one.
  if [ -r /dev/tty ]; then
    printf "%s" "Remove the orchestrator agent and all vPLCs from this device? [y/N] " >/dev/tty
    read -r answer </dev/tty || answer=""
  elif [ -t 0 ]; then
    printf "%s" "Remove the orchestrator agent and all vPLCs from this device? [y/N] "
    read -r answer || answer=""
  else
    echo -e "${RED}[ERROR] No terminal available to confirm. Re-run with --yes to uninstall.${RESET}"
    exit 1
  fi

  case "$answer" in
  y | Y | yes | YES | Yes) return 0 ;;
  *)
    echo "Aborted. Nothing was removed."
    exit 1
    ;;
  esac
}

run_uninstall() {
  local mtls_dir_detected=""
  local runtimes=""
  local name=""
  local agent_teardown_ok=false

  echo
  echo -e "${BOLD}${RED}=====================================================${RESET}"
  echo -e "${BOLD}${RED}  UNINSTALL ORCHESTRATOR AGENT                       ${RESET}"
  echo -e "${BOLD}${RED}=====================================================${RESET}"
  echo
  echo "This removes, from this device:"
  echo "  - every vPLC runtime container and its stored programs"
  echo "  - the orchestrator agent and network monitor containers"
  echo "  - orchestrator MACVLAN and internal networks"
  echo "  - the $SHARED_VOLUME volume (client registry, vNIC configs, logs)"
  if [ "$KEEP_CERTS" != true ]; then
    echo "  - the mTLS client certificate and key (this device's identity)"
  fi
  if [ "$KEEP_IMAGES" != true ]; then
    echo "  - the pulled orchestrator, netmon and runtime images"
  fi
  echo
  echo -e "${YELLOW}Running PLC logic will stop. This cannot be undone.${RESET}"
  echo

  confirm_uninstall

  if ! command -v docker >/dev/null 2>&1; then
    u_warn "docker is not installed on this host; only host files can be cleaned up"
    u_step "Removing mTLS material"
    remove_mtls_material "$MTLS_DIR"
    u_step "Removing installer leftovers"
    remove_installer_leftovers
    echo
    echo -e "${BOLD}${GREEN}UNINSTALL COMPLETE${RESET} (no Docker resources to remove)"
    return 0
  fi

  if ! docker info >/dev/null 2>&1; then
    echo -e "${RED}[ERROR] Cannot talk to the Docker daemon. Start Docker and re-run.${RESET}"
    exit 1
  fi

  # Read the cert location while the agent container is still around.
  mtls_dir_detected=$(detect_mtls_dir)

  u_step "Stage 1: agent-driven teardown (delete_orchestrator equivalent)"
  if [ "$FORCE_LOCAL" = true ]; then
    u_info "skipped (--force-local): going straight to host cleanup"
  elif run_agent_teardown; then
    agent_teardown_ok=true
  fi

  u_step "Stage 2: host cleanup"
  if [ "$agent_teardown_ok" = true ]; then
    u_info "verifying nothing the agent could not reach was left behind"
  else
    u_info "removing orchestrator resources directly"
  fi

  u_step "Removing managed runtime containers"
  runtimes=$(find_runtime_containers)
  if [ -z "$runtimes" ]; then
    u_info "no runtime containers found"
  else
    for name in $runtimes; do
      remove_container "$name"
    done
  fi

  u_step "Cleaning up Proxy ARP interfaces"
  cleanup_proxy_arp

  u_step "Removing network monitor and upgrader containers"
  remove_container "$NETMON_CONTAINER_NAME"
  remove_container "$UPGRADER_CONTAINER_NAME"

  u_step "Removing the orchestrator agent container"
  remove_container "$CONTAINER_NAME"

  # After every container is gone: networks have no endpoints left and the
  # shared volume is finally unreferenced.
  u_step "Removing orchestrator networks"
  remove_orchestrator_networks

  u_step "Removing dangling iptables rules"
  cleanup_iptables_rules

  u_step "Removing the shared volume"
  remove_shared_volume

  u_step "Removing mTLS material"
  remove_mtls_material "$mtls_dir_detected"

  u_step "Removing Docker images"
  remove_images

  u_step "Removing installer leftovers"
  remove_installer_leftovers

  echo
  echo -e "${BOLD}${GREEN}UNINSTALL COMPLETE${RESET}"
  echo -e "${GRAY}=====================================================${RESET}"
  echo
  echo "The orchestrator agent has been removed from this device."
  echo
  echo -e "${YELLOW}One step is left, and it is not local:${RESET} the device still"
  echo "appears in Autonomy Edge. Delete it there so the cloud stops"
  echo "expecting a heartbeat from this orchestrator."
  echo
  if [ "$KEEP_CERTS" = true ]; then
    echo -e "${GRAY}The mTLS identity was kept, so a reinstall can reuse it.${RESET}"
  fi
  echo -e "${GRAY}Dedicated NICs return to the host when their container is${RESET}"
  echo -e "${GRAY}destroyed; reboot if an interface kept an orchestrator name.${RESET}"
  echo -e "${GRAY}=====================================================${RESET}"
}

if [ "$MODE" = "uninstall" ]; then
  run_uninstall
  exit 0
fi

### --- DEPENDENCIES --- ###
echo "Checking and installing required dependencies..."
PKG_MANAGER=""

# Detect package manager
if command -v apt-get &>/dev/null; then
  PKG_MANAGER="apt-get"
elif command -v dnf &>/dev/null; then
  PKG_MANAGER="dnf"
elif command -v yum &>/dev/null; then
  PKG_MANAGER="yum"
else
  echo "[ERROR] No supported package manager found (apt, dnf, or yum). Install dependencies manually."
  echo "Required packages: curl, jq, openssl, docker"
  echo "Attempting to continue without automatic dependency installation..."
  PKG_MANAGER="none"
fi

# Define package names per package manager
declare -A PKG_MAP
if [[ "$PKG_MANAGER" == "apt-get" ]]; then
  PKG_MAP=(
    [curl]="curl"
    [jq]="jq"
    [openssl]="openssl"
    [docker]="docker.io"
  )
elif [[ "$PKG_MANAGER" == "dnf" ]]; then
  PKG_MAP=(
    [curl]="curl"
    [jq]="jq"
    [openssl]="openssl"
    [docker]="docker"
  )
elif [[ "$PKG_MANAGER" == "yum" ]]; then
  PKG_MAP=(
    [curl]="curl"
    [jq]="jq"
    [openssl]="openssl"
    [docker]="docker"
  )
fi

# Collect missing packages
MISSING_PKGS=()
for cmd in curl jq openssl docker; do
  if ! command -v "$cmd" &>/dev/null; then
    echo "Missing dependency: $cmd"
    if [[ -n "${PKG_MAP[$cmd]}" ]]; then
      MISSING_PKGS+=("${PKG_MAP[$cmd]}")
    fi
  else
    echo "[SUCCESS] $cmd is already installed."
  fi
done

# Install missing packages
if [ ${#MISSING_PKGS[@]} -ne 0 ]; then
  echo "Updating package lists and installing missing dependencies: ${MISSING_PKGS[*]}"
  case "$PKG_MANAGER" in
  apt-get)
    sudo apt-get update -y
    sudo apt-get install -y "${MISSING_PKGS[@]}"
    ;;
  dnf)
    sudo dnf install -y "${MISSING_PKGS[@]}"
    ;;
  yum)
    sudo yum install -y "${MISSING_PKGS[@]}"
    ;;
  none)
    echo "[ERROR] Cannot install dependencies automatically. Please install: ${MISSING_PKGS[*]}"
    exit 1
    ;;
  esac
fi

echo "Creating shared volume for container communication..."
if docker volume inspect "$SHARED_VOLUME" &>/dev/null; then
  echo "[SUCCESS] Shared volume $SHARED_VOLUME already exists"
else
  docker volume create "$SHARED_VOLUME"
  echo "[SUCCESS] Created shared volume $SHARED_VOLUME"
fi

### --- UPGRADE DETECTION --- ###
# If an orchestrator is already running with valid mTLS certs, perform an
# in-place upgrade instead of a full install. This preserves the existing
# identity, certificates, vPLC containers, and network configuration.
EXISTING_ORCHESTRATOR=false
if docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
  EXISTING_MTLS_DIR=$(docker inspect "$CONTAINER_NAME" --format '{{range .Mounts}}{{if eq .Destination "/root/.mtls"}}{{.Source}}{{end}}{{end}}' 2>/dev/null || true)
  if [ -n "$EXISTING_MTLS_DIR" ] && [ -f "$EXISTING_MTLS_DIR/client.crt" ] && [ -f "$EXISTING_MTLS_DIR/client.key" ]; then
    EXISTING_ORCHESTRATOR=true
  fi
fi

if [ "$EXISTING_ORCHESTRATOR" = true ]; then
  echo ""
  echo "======================================================"
  echo "  EXISTING ORCHESTRATOR DETECTED — UPGRADING IN-PLACE"
  echo "======================================================"
  echo ""
  echo "mTLS certificates: $EXISTING_MTLS_DIR"
  echo "Shared volume: $SHARED_VOLUME"
  echo ""

  # Pull latest images
  echo "Pulling latest orchestrator image..."
  docker pull "$IMAGE_NAME"
  echo "Pulling latest netmon image..."
  docker pull "$NETMON_IMAGE_NAME"

  # Upgrade netmon sidecar
  echo "Upgrading network monitor sidecar..."
  if docker ps -a --format '{{.Names}}' | grep -q "^${NETMON_CONTAINER_NAME}$"; then
    docker rm -f "$NETMON_CONTAINER_NAME"
  fi
  docker run -d \
    --name "$NETMON_CONTAINER_NAME" \
    --network=host \
    --pid=host \
    --privileged \
    --restart unless-stopped \
    -v "$SHARED_VOLUME:/var/orchestrator" \
    -v /dev:/dev \
    -v /run/udev:/run/udev:ro \
    "$NETMON_IMAGE_NAME"
  echo "[SUCCESS] Network monitor sidecar upgraded"

  # Clean up any leftover upgrader from a previous attempt
  docker rm -f orchestrator_upgrader 2>/dev/null || true

  # Record the old container ID so we can detect when it's been replaced
  OLD_CONTAINER_ID=$(docker inspect "$CONTAINER_NAME" --format '{{.Id}}' 2>/dev/null || true)
  echo "Old orchestrator container: ${OLD_CONTAINER_ID:0:12}"

  # Spawn the upgrader container — same process as the v1.0.0+ remote upgrade.
  echo "Spawning upgrader container..."
  docker run -d \
    --name orchestrator_upgrader \
    -v /var/run/docker.sock:/var/run/docker.sock \
    -e UPGRADE_MODE=true \
    -e TARGET_CONTAINER="$CONTAINER_NAME" \
    -e NEW_IMAGE="$IMAGE_NAME" \
    -e MTLS_HOST_PATH="$EXISTING_MTLS_DIR" \
    -e SHARED_VOLUME="$SHARED_VOLUME" \
    "$IMAGE_NAME" \
    python src/tools/upgrade_self.py

  # Step 1: Wait for old orchestrator to be stopped and removed
  echo "Waiting for old orchestrator to be removed..."
  for i in $(seq 1 30); do
    CURRENT_ID=$(docker inspect "$CONTAINER_NAME" --format '{{.Id}}' 2>/dev/null || true)
    if [ -z "$CURRENT_ID" ] || [ "$CURRENT_ID" != "$OLD_CONTAINER_ID" ]; then
      echo "[OK] Old orchestrator removed"
      break
    fi
    sleep 1
  done

  # Step 2: Wait for new orchestrator container to appear and be running
  echo "Waiting for new orchestrator to start..."
  UPGRADE_OK=false
  for i in $(seq 1 30); do
    NEW_ID=$(docker inspect "$CONTAINER_NAME" --format '{{.Id}}' 2>/dev/null || true)
    if [ -n "$NEW_ID" ] && [ "$NEW_ID" != "$OLD_CONTAINER_ID" ]; then
      NEW_STATUS=$(docker inspect "$CONTAINER_NAME" --format '{{.State.Status}}' 2>/dev/null || true)
      if [ "$NEW_STATUS" = "running" ]; then
        echo "[OK] New orchestrator running: ${NEW_ID:0:12}"
        UPGRADE_OK=true
        break
      fi
    fi
    sleep 1
  done

  if [ "$UPGRADE_OK" != "true" ]; then
    echo "[ERROR] Upgrade failed. Check logs with: docker logs orchestrator_upgrader"
    exit 1
  fi

  echo ""
  echo "======================================================"
  echo "  UPGRADE COMPLETE"
  echo "======================================================"
  echo ""
  echo "The orchestrator has been upgraded to the latest version."
  echo "All vPLC containers and configurations have been preserved."
  echo ""
  exit 0
fi

### --- STEP 1: REQUEST CUSTOM ID --- ###
echo "Requesting ID from $GET_ID_URL..."
response=$(curl -fsSL "$GET_ID_URL")

# Validate JSON format
if ! echo "$response" | jq empty 2>/dev/null; then
  echo "[ERROR] Invalid server response: not JSON."
  echo "$response"
  exit 1
fi

CUSTOM_ID=$(echo "$response" | jq -r '.data.id')
EXPIRES_AT=$(echo "$response" | jq -r '.data.expiresAt')
EXPIRES_IN=$(echo "$response" | jq -r '.data.expiresIn')

if [[ -z "$CUSTOM_ID" || "$CUSTOM_ID" == "null" ]]; then
  echo "[ERROR] Failed to retrieve ID from server."
  exit 1
fi

# Sanitize EXPIRES_IN — fall back to 300s if the server omits it.
if ! [[ "$EXPIRES_IN" =~ ^[0-9]+$ ]] || [ "$EXPIRES_IN" -lt 30 ]; then
  EXPIRES_IN=300
fi

PENDING_URL="$SERVER_URL/orchestrators/$CUSTOM_ID/pending"

### --- STEP 2: GENERATE CLIENT KEY AND CSR --- ###
echo "Generating client key and certificate signing request..."
mkdir -p "$MTLS_DIR"
chmod 700 "$MTLS_DIR"
# From this point on, any failure before validation should leave the host clean.
CLEANUP_CERTS_ON_FAILURE=true

echo ""
echo "Generating client certificate for orchestrator ID: $CUSTOM_ID"
echo ""

# STEP 2.1: Generate client private key
echo "--- 1. Generating client private key ---"
if ! openssl genrsa -out "$KEY_PATH" 4096 2>/dev/null; then
  echo "[ERROR] Failed to generate client private key. Aborting."
  exit 1
fi
chmod 600 "$KEY_PATH"
echo "Client private key generated: $KEY_PATH"

# STEP 2.2: Create CSR configuration file
echo "--- 2. Creating CSR configuration ---"
cat <<EOF > "$CSR_CONFIG_FILE"
[req]
distinguished_name = req_distinguished_name
req_extensions = v3_req
prompt = no

[req_distinguished_name]
C=BR
ST=SP
L=SaoPaulo
O=AutonomyLogic
OU=Production
CN=${CUSTOM_ID}

[v3_req]
keyUsage = keyEncipherment, dataEncipherment
extendedKeyUsage = clientAuth
EOF

echo "CSR configuration created"

# STEP 2.3: Generate CSR (Certificate Signing Request)
echo "--- 3. Generating Certificate Signing Request (CSR) ---"
openssl req -new \
  -key "$KEY_PATH" \
  -out "$CSR_PATH" \
  -config "$CSR_CONFIG_FILE" 2>/dev/null

echo "CSR generated: $CSR_PATH"

### --- STEP 3: ENROLL WITH SERVER (CSR SIGNING) --- ###
echo "--- 4. Submitting CSR to server for signing ---"
echo "Enrolling with $ENROLL_URL..."
http_code=$(curl -sS -w "%{http_code}" -o "$ENROLL_RESP_FILE" \
  -X POST "$ENROLL_URL" \
  -F "csr=@$CSR_PATH") || {
  echo "[ERROR] Enrollment request failed. Please check network connectivity."
  exit 1
}

if [[ -z "$http_code" || ! "$http_code" =~ ^[0-9]{3}$ ]]; then
  echo "[ERROR] Invalid HTTP response code received: $http_code"
  exit 1
fi

if [[ "$http_code" -ne 200 ]]; then
  echo "[ERROR] Enrollment failed. HTTP code: $http_code"
  echo "Server response:"
  cat "$ENROLL_RESP_FILE"
  echo
  exit 1
fi

# Extract certificate from response JSON
certificate=$(jq -r '.data.certificate' "$ENROLL_RESP_FILE")
message=$(jq -r '.data.message' "$ENROLL_RESP_FILE")
id_resp=$(jq -r '.data.id' "$ENROLL_RESP_FILE")
status=$(jq -r '.statusCode' "$ENROLL_RESP_FILE")

if [[ "$status" != "200" ]]; then
  echo "[WARNING] Unexpected server status: $status"
  cat "$ENROLL_RESP_FILE"
  echo
  exit 1
fi

if [[ -z "$certificate" || "$certificate" == "null" ]]; then
  echo "[ERROR] No certificate received from server"
  cat "$ENROLL_RESP_FILE"
  echo
  exit 1
fi

# Save the signed certificate
echo "$certificate" > "$CRT_PATH"
chmod 644 "$CRT_PATH"
echo "Client certificate signed and saved: $CRT_PATH"

# Verify the certificate has correct extensions for mTLS client auth
echo "--- 5. Verifying certificate extensions ---"
if ! openssl x509 -in "$CRT_PATH" -noout -purpose 2>/dev/null | grep -q "SSL client : Yes"; then
  echo "[ERROR] Server returned a certificate that is not valid for SSL client authentication."
  echo "This may indicate a server misconfiguration. Please contact support."
  exit 1
fi
echo "Certificate verified: valid for SSL client authentication"

echo ""
echo "CERTIFICATE ENROLLMENT COMPLETE"
echo "=================================================="
echo "Files created:"
echo "   - Client key:         $KEY_PATH"
echo "   - Client certificate: $CRT_PATH"
echo ""
echo "   Valid until:"
openssl x509 -in "$CRT_PATH" -noout -enddate
echo ""
echo "   Fingerprint SHA256:"
openssl x509 -in "$CRT_PATH" -noout -fingerprint -sha256
echo ""
echo "=================================================="
echo ""
echo "[SUCCESS] Enrollment completed: $message (ID: $id_resp)"

### --- STEP 4: WAIT FOR USER VALIDATION ON AUTONOMY EDGE --- ###
echo
echo -e "${BOLD}${GREEN}=====================================================${RESET}"
echo -e "${BOLD}${GREEN}  ACTION REQUIRED: VALIDATE YOUR ORCHESTRATOR ID     ${RESET}"
echo -e "${BOLD}${GREEN}=====================================================${RESET}"
echo
echo -e "Orchestrator ID: ${BOLD}${CYAN}${CUSTOM_ID}${RESET}"
echo
echo "Open the Autonomy Edge application and link a new orchestrator"
echo "using the ID above. The installation will continue automatically"
echo "as soon as the validation is recorded."
echo
echo -e "${GRAY}If the ID is not validated within ${EXPIRES_IN}s, the installation${RESET}"
echo -e "${GRAY}will be aborted and the unvalidated certificate removed.${RESET}"
echo

START_TS=$(date +%s)
USE_CR=false
[ -t 1 ] && USE_CR=true

VALIDATED=false
FAIL_REASON=""
TRANSIENT_WARN_PRINTED=false

while :; do
  ELAPSED=$(( $(date +%s) - START_TS ))
  REMAINING=$(( EXPIRES_IN - ELAPSED ))
  if [ "$REMAINING" -le 0 ]; then
    FAIL_REASON="timeout"
    break
  fi

  poll_code=$(curl -sS -o "$PENDING_RESP_FILE" -w "%{http_code}" "$PENDING_URL" 2>/dev/null || echo "000")

  if [ "$poll_code" = "200" ]; then
    poll_status=$(jq -r '.data.status // .status // empty' "$PENDING_RESP_FILE" 2>/dev/null || echo "")
    if [ "$poll_status" = "validated" ]; then
      VALIDATED=true
      break
    fi
  elif [ "$poll_code" = "404" ]; then
    FAIL_REASON="server-rejected"
    break
  else
    # 000 timeout, 5xx, etc. — keep polling, but surface a one-shot warning so
    # a permanently-unhealthy server doesn't look like a silent countdown.
    if [ "$TRANSIENT_WARN_PRINTED" = false ]; then
      if [ "$USE_CR" = true ]; then printf "\r\033[K"; fi
      echo -e "${YELLOW}[WARN] Server returned ${poll_code} while polling for validation; will keep retrying.${RESET}"
      TRANSIENT_WARN_PRINTED=true
    fi
  fi

  if [ "$USE_CR" = true ]; then
    printf "\r${YELLOW}Waiting for validation… %3ds remaining${RESET}\033[K" "$REMAINING"
  elif (( REMAINING % 30 == 0 )) || [ "$REMAINING" -le 10 ]; then
    echo "Waiting for validation… ${REMAINING}s remaining"
  fi

  sleep "$POLL_INTERVAL_SECONDS"
done

if [ "$USE_CR" = true ]; then printf "\r\033[K"; fi

if [ "$VALIDATED" != true ]; then
  echo
  if [ "$FAIL_REASON" = "server-rejected" ]; then
    echo -e "${BOLD}${RED}Server reports the orchestrator ID is unknown or expired.${RESET}"
    echo -e "${YELLOW}Reverting: removing the unvalidated certificate. No containers were created.${RESET}"
    echo "Re-run the installer to request a fresh ID and try again."
  else
    echo -e "${BOLD}${YELLOW}Installation was not validated on Autonomy Edge.${RESET}"
    echo -e "${YELLOW}Reverting: removing the unvalidated certificate. No containers were created.${RESET}"
    echo "Re-run the installer and complete the validation in the app to finish setup."
  fi
  exit 1
fi

# Validation succeeded — keep the cert files past the trap.
CLEANUP_CERTS_ON_FAILURE=false
echo
echo -e "${BOLD}${GREEN}Validation confirmed. Continuing installation…${RESET}"
echo

### --- STEP 5: DEPLOY NETWORK MONITOR SIDECAR --- ###
echo "Deploying network monitor sidecar container..."

if docker pull "$NETMON_IMAGE_NAME" 2>/dev/null; then
  echo "[SUCCESS] Pulled network monitor image: $NETMON_IMAGE_NAME"
else
  echo "[WARNING] No prebuilt netmon image found. Falling back to local build..."

  if [ -d "$SOURCE_DIR/.git" ]; then
    echo "Updating existing source clone..."
    if ! git -C "$SOURCE_DIR" pull --rebase; then
      echo "Pull failed, stashing local changes and retrying..."
      git -C "$SOURCE_DIR" stash push --include-untracked -m "installer-auto-stash $(date +%s)" || true
      if ! git -C "$SOURCE_DIR" pull --rebase; then
        echo "[ERROR] git pull still failing after stash. Please inspect $SOURCE_DIR."
        exit 1
      fi
    fi
  else
    echo "Cloning source to $SOURCE_DIR..."
    git clone https://github.com/autonomy-logic/orchestrator-agent.git "$SOURCE_DIR"
  fi

  echo "Building network monitor image locally..."
  docker build -t "$NETMON_IMAGE_NAME" -f "$SOURCE_DIR/install/Dockerfile.netmon" "$SOURCE_DIR/install"

  echo "[SUCCESS] Local netmon build completed: $NETMON_IMAGE_NAME"
fi

if docker ps -a --format '{{.Names}}' | grep -q "^${NETMON_CONTAINER_NAME}$"; then
  echo "Removing existing network monitor container..."
  docker rm -f "$NETMON_CONTAINER_NAME"
fi


docker run -d \
  --name "$NETMON_CONTAINER_NAME" \
  --network=host \
  --pid=host \
  --privileged \
  --restart unless-stopped \
  -v "$SHARED_VOLUME:/var/orchestrator" \
  -v /dev:/dev \
  -v /run/udev:/run/udev:ro \
  "$NETMON_IMAGE_NAME"

echo "[SUCCESS] Network monitor sidecar started"

### --- STEP 6: PULL ORCHESTRATOR AGENT IMAGE AND CREATE CONTAINER --- ###
echo "Pulling Docker image: $IMAGE_NAME"
if docker pull "$IMAGE_NAME"; then
  echo "Pulled image: $IMAGE_NAME"
else
  echo "[WARNING] No prebuilt image found for this host architecture. Falling back to local build..."
  # Clone or update the source tree
  if [ -d "$SOURCE_DIR/.git" ]; then
    echo "Updating existing source clone..."
    if ! git -C "$SOURCE_DIR" pull --rebase; then
      echo "Pull failed, stashing local changes and retrying..."
      git -C "$SOURCE_DIR" stash push --include-untracked -m "installer-auto-stash $(date +%s)" || true
      if ! git -C "$SOURCE_DIR" pull --rebase; then
        echo "[ERROR] git pull still failing after stash. Please inspect $SOURCE_DIR."
        exit 1
      fi
    fi
  else
    echo "Cloning source to $SOURCE_DIR..."
    git clone https://github.com/autonomy-logic/orchestrator-agent.git "$SOURCE_DIR"
  fi

  # Build locally for the host architecture
  # Use 'docker build' which builds for the local machine arch (simplest and most reliable for fallback)
  echo "Building Docker image locally for this host architecture..."
  docker build -t "$IMAGE_NAME" "$SOURCE_DIR"

  echo "Local build completed: $IMAGE_NAME"
fi

if docker ps -a --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
  echo "Existing container detected. Removing and recreating..."
  docker rm -f "$CONTAINER_NAME"
fi

echo "Creating new container: $CONTAINER_NAME"
docker run -d \
  --name "$CONTAINER_NAME" \
  --restart unless-stopped \
  -v "$MTLS_DIR:/root/.mtls:ro" \
  -v /var/run/docker.sock:/var/run/docker.sock \
  -v "$SHARED_VOLUME:/var/orchestrator" \
  "$IMAGE_NAME"

echo
echo
echo -e "${BOLD}${GREEN}INSTALLATION COMPLETE${RESET}"
echo -e "${GRAY}=====================================================${RESET}"
echo
echo -e "Orchestrator ID: ${BOLD}${CYAN}${CUSTOM_ID}${RESET}"
echo "The orchestrator container is running and connected to Autonomy Edge."
echo -e "${GRAY}=====================================================${RESET}"