#!/bin/sh
set -eu

ROLE="${1:-worker}"
shift || true

REPO="${AGENTMUX_REPO:-kinboyw/agentmux}"
VERSION="${AGENTMUX_VERSION:-latest}"
HUB_HTTP="https://agentmux.kinboy.wang"
HUB_WS="wss://agentmux.kinboy.wang"
BIN_DIR="${AGENTMUX_BIN_DIR:-$HOME/.local/bin}"
BIN="$BIN_DIR/agentmux"
HUB_BIN="$BIN_DIR/agentmux-hub"
TUI_BIN="$BIN_DIR/agentmux-tui"

mkdir -p "$BIN_DIR"

verify_sha256() {
  archive="$1"
  checksum_file="$2"
  expected="$(cut -d ' ' -f 1 "$checksum_file" | head -n 1)"
  if [ -z "$expected" ]; then
    echo "empty checksum file: $checksum_file" >&2
    exit 1
  fi
  if command -v sha256sum >/dev/null 2>&1; then
    actual="$(sha256sum "$archive" | cut -d ' ' -f 1)"
  elif command -v shasum >/dev/null 2>&1; then
    actual="$(shasum -a 256 "$archive" | cut -d ' ' -f 1)"
  else
    echo "sha256 verification requires sha256sum or shasum" >&2
    exit 1
  fi
  if [ "$actual" != "$expected" ]; then
    echo "sha256 mismatch for $archive" >&2
    echo "expected: $expected" >&2
    echo "actual:   $actual" >&2
    exit 1
  fi
}

case "$ROLE" in
  worker)
    ASSET_ROLE="worker"
    INSTALL_BIN="$BIN"
    ;;
  control)
    ASSET_ROLE="tui"
    INSTALL_BIN="$TUI_BIN"
    ;;
  hub)
    ASSET_ROLE="hub"
    INSTALL_BIN="$HUB_BIN"
    ;;
  *)
    echo "usage: install.sh worker|control|hub [agentmux flags]" >&2
    exit 2
    ;;
esac

if [ "$ROLE" = "control" ] && command -v agentmux-tui >/dev/null 2>&1; then
  INSTALL_BIN="$(command -v agentmux-tui)"
elif [ "$ROLE" = "hub" ] && command -v agentmux-hub >/dev/null 2>&1; then
  INSTALL_BIN="$(command -v agentmux-hub)"
elif [ "$ROLE" != "control" ] && command -v agentmux >/dev/null 2>&1; then
  INSTALL_BIN="$(command -v agentmux)"
elif command -v go >/dev/null 2>&1 && [ "$ROLE" = "hub" ] && [ -f "./cmd/agentmux-hub/main.go" ]; then
  go build -o "$INSTALL_BIN" ./cmd/agentmux-hub
elif command -v go >/dev/null 2>&1 && [ -f "./cmd/agentmux/main.go" ]; then
  go build -o "$INSTALL_BIN" ./cmd/agentmux
elif command -v curl >/dev/null 2>&1 && command -v tar >/dev/null 2>&1; then
  os="$(uname -s | tr '[:upper:]' '[:lower:]')"
  arch="$(uname -m)"
  case "$arch" in
    x86_64|amd64) arch="amd64" ;;
    aarch64|arm64) arch="arm64" ;;
    *) echo "unsupported architecture: $arch" >&2; exit 1 ;;
  esac
  case "$os" in
    linux|darwin) ;;
    *) echo "unsupported OS: $os" >&2; exit 1 ;;
  esac
  asset_base="agentmux-${ASSET_ROLE}-${os}-${arch}"
  asset="${asset_base}.tar.gz"
  if [ "$VERSION" = "latest" ]; then
    url="https://github.com/${REPO}/releases/latest/download/${asset}"
  else
    url="https://github.com/${REPO}/releases/download/${VERSION}/${asset}"
  fi
  tmp="$(mktemp -d)"
  trap 'rm -rf "$tmp"' EXIT
  echo "Downloading $url" >&2
  if ! curl -fsSL "$url" -o "$tmp/$asset"; then
    if [ "$ROLE" = "hub" ]; then
      exit 1
    fi
    if [ "$ROLE" = "control" ]; then
      fallback_base="agentmux-control-${os}-${arch}"
      fallback_asset="${fallback_base}.tar.gz"
      if [ "$VERSION" = "latest" ]; then
        fallback_url="https://github.com/${REPO}/releases/latest/download/${fallback_asset}"
      else
        fallback_url="https://github.com/${REPO}/releases/download/${VERSION}/${fallback_asset}"
      fi
      echo "Falling back to $fallback_url" >&2
      if curl -fsSL "$fallback_url" -o "$tmp/$fallback_asset"; then
        asset_base="$fallback_base"
        asset="$fallback_asset"
        url="$fallback_url"
      fi
    fi
    if [ ! -f "$tmp/$asset" ]; then
      legacy_base="agentmux-${os}-${arch}"
      asset_base="$legacy_base"
      asset="${legacy_base}.tar.gz"
      if [ "$VERSION" = "latest" ]; then
        url="https://github.com/${REPO}/releases/latest/download/${asset}"
      else
        url="https://github.com/${REPO}/releases/download/${VERSION}/${asset}"
      fi
      echo "Falling back to $url" >&2
      curl -fsSL "$url" -o "$tmp/$asset"
    fi
  fi
  checksum_url="${url}.sha256"
  echo "Verifying $checksum_url" >&2
  curl -fsSL "$checksum_url" -o "$tmp/$asset.sha256"
  verify_sha256 "$tmp/$asset" "$tmp/$asset.sha256"
  tar -xzf "$tmp/$asset" -C "$tmp"
  install -m 0755 "$tmp/${asset_base}" "$INSTALL_BIN"
else
  echo "agentmux binary is not installed." >&2
  echo "Install curl+tar, put agentmux in PATH, or run this script from a source checkout with Go installed." >&2
  exit 1
fi

case "$ROLE" in
  worker)
    exec "$INSTALL_BIN" worker join --hub "$HUB_WS" "$@"
    ;;
  control)
    exec "$INSTALL_BIN" --hub "$HUB_HTTP" "$@"
    ;;
  hub)
    exec "$INSTALL_BIN" hub "$@"
    ;;
esac