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

relay_url="http://127.0.0.1:4173"
pairing_code=""
interval_seconds=60
once=false
install_mode="manual"
state_dir="$HOME/Library/Application Support/AllReach"
state_path="$state_dir/agent-state.json"
keychain_service="AllReach Agent"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --relay-url)
      relay_url="${2:-}"
      shift 2
      ;;
    --pairing-code)
      pairing_code="${2:-}"
      shift 2
      ;;
    --interval-seconds)
      interval_seconds="${2:-60}"
      shift 2
      ;;
    --install-mode)
      install_mode="${2:-manual}"
      shift 2
      ;;
    --once)
      once=true
      shift
      ;;
    *)
      printf "Unknown option: %s\n" "$1" >&2
      exit 1
      ;;
  esac
done

screen_sharing_enabled() {
  launchctl print system/com.apple.screensharing >/dev/null 2>&1
}

vnc_listening() {
  nc -z 127.0.0.1 5900 >/dev/null 2>&1
}

json_bool() {
  if "$@"; then
    printf "true"
  else
    printf "false"
  fi
}

json_escape() {
  printf "%s" "$1" | sed 's/\\/\\\\/g; s/"/\\"/g'
}

ipv4_addresses() {
  ifconfig | awk '/inet / && $2 != "127.0.0.1" { print $2 }' | sort -u
}

addresses_json_from_system() {
  local first=true address
  printf "["
  while IFS= read -r address; do
    if [[ -z "$address" ]]; then
      continue
    fi
    if [[ "$first" == true ]]; then
      first=false
    else
      printf ","
    fi
    printf '"%s"' "$(json_escape "$address")"
  done < <(ipv4_addresses)
  printf "]"
}

agent_status_json() {
  local computer_name macos_version enabled listening generated_at addresses_json
  computer_name="$(scutil --get ComputerName 2>/dev/null || hostname)"
  macos_version="$(sw_vers -productVersion)"
  addresses_json="$(addresses_json_from_system)"
  enabled="$(json_bool screen_sharing_enabled)"
  listening="$(json_bool vnc_listening)"
  generated_at="$(date -u +"%Y-%m-%dT%H:%M:%SZ")"

  printf '{'
  printf '"deviceName":"%s",' "$(json_escape "$computer_name")"
  printf '"platform":"mac",'
  printf '"edition":"macOS %s",' "$(json_escape "$macos_version")"
  printf '"installMode":"%s",' "$(json_escape "$install_mode")"
  printf '"online":true,'
  printf '"addresses":%s,' "$addresses_json"
  printf '"protocols":[{"name":"vnc","enabled":%s,"port":5900,"listening":%s}],' "$enabled" "$listening"
  printf '"routes":{"lan":true,"mesh":false,"relay":true},'
  printf '"generatedAt":"%s"' "$(json_escape "$generated_at")"
  printf '}'
}

register_agent() {
  local status body response agent_id agent_secret device_id device_name
  status="$(agent_status_json)"
  body="$(printf '{"pairingCode":"%s",%s' "$(json_escape "$pairing_code")" "${status#\{}")"
  response="$(curl -fsS -X POST "$relay_url/api/agents/register" -H "Content-Type: application/json" -d "$body")"
  agent_id="$(printf "%s" "$response" | sed -n 's/.*"id":"\([^"]*\)".*/\1/p')"
  agent_secret="$(printf "%s" "$response" | sed -n 's/.*"agentSecret":"\([^"]*\)".*/\1/p')"
  device_id="$(printf "%s" "$response" | sed -n 's/.*"deviceId":"\([^"]*\)".*/\1/p')"
  device_name="$(printf "%s" "$response" | sed -n 's/.*"deviceName":"\([^"]*\)".*/\1/p')"

  if [[ -z "$agent_id" ]]; then
    printf "AllReach pairing failed: %s\n" "$response" >&2
    exit 1
  fi

  save_agent_state "$agent_id" "$agent_secret" "$device_id" "$device_name" "$(date -u +"%Y-%m-%dT%H:%M:%SZ")"
  printf "AllReach agent paired as %s for %s.\n" "$agent_id" "$device_name"
}

read_agent_id() {
  if [[ ! -f "$state_path" ]]; then
    printf "No paired AllReach agent was found. Run with --pairing-code NN-NN-NN first.\n" >&2
    exit 1
  fi

  sed -n 's/.*"agentId":"\([^"]*\)".*/\1/p' "$state_path"
}

read_saved_relay_url() {
  if [[ ! -f "$state_path" ]]; then
    return
  fi

  sed -n 's/.*"relayUrl":"\([^"]*\)".*/\1/p' "$state_path"
}

read_agent_secret() {
  if [[ ! -f "$state_path" ]]; then
    return
  fi

  local agent_id
  agent_id="$(read_agent_id)"
  if [[ -n "$agent_id" ]]; then
    security find-generic-password -a "$agent_id" -s "$keychain_service" -w 2>/dev/null && return
  fi

  sed -n 's/.*"agentSecret":"\([^"]*\)".*/\1/p' "$state_path"
}

save_agent_secret() {
  local agent_id="$1" agent_secret="$2"
  if [[ -z "$agent_id" || -z "$agent_secret" ]]; then
    return
  fi

  security add-generic-password -a "$agent_id" -s "$keychain_service" -w "$agent_secret" -U >/dev/null 2>&1 || true
}

read_state_field() {
  local field="$1"
  if [[ ! -f "$state_path" ]]; then
    return
  fi

  sed -n "s/.*\"$field\":\"\\([^\"]*\\)\".*/\\1/p" "$state_path"
}

save_agent_state() {
  local agent_id="$1" agent_secret="$2" device_id="$3" device_name="$4" paired_at="$5" last_heartbeat_at="${6:-}"
  mkdir -p "$state_dir"
  save_agent_secret "$agent_id" "$agent_secret"
  printf '{"agentId":"%s","agentSecretStorage":"keychain","deviceId":"%s","deviceName":"%s","relayUrl":"%s","pairedAt":"%s","lastHeartbeatAt":"%s"}\n' \
    "$(json_escape "$agent_id")" \
    "$(json_escape "$device_id")" \
    "$(json_escape "$device_name")" \
    "$(json_escape "$relay_url")" \
    "$(json_escape "$paired_at")" \
    "$(json_escape "$last_heartbeat_at")" > "$state_path"
}

update_state_from_response() {
  local response="$1" agent_id agent_secret device_id device_name paired_at last_seen
  agent_id="$(printf "%s" "$response" | sed -n 's/.*"id":"\([^"]*\)".*/\1/p')"
  agent_secret="$(printf "%s" "$response" | sed -n 's/.*"agentSecret":"\([^"]*\)".*/\1/p')"
  device_id="$(printf "%s" "$response" | sed -n 's/.*"deviceId":"\([^"]*\)".*/\1/p')"
  device_name="$(printf "%s" "$response" | sed -n 's/.*"deviceName":"\([^"]*\)".*/\1/p')"
  last_seen="$(printf "%s" "$response" | sed -n 's/.*"lastSeen":"\([^"]*\)".*/\1/p')"

  agent_id="${agent_id:-$(read_state_field agentId)}"
  agent_secret="${agent_secret:-$(read_state_field agentSecret)}"
  device_id="${device_id:-$(read_state_field deviceId)}"
  device_name="${device_name:-$(read_state_field deviceName)}"
  paired_at="$(read_state_field pairedAt)"
  paired_at="${paired_at:-$(date -u +"%Y-%m-%dT%H:%M:%SZ")}"
  save_agent_state "$agent_id" "$agent_secret" "$device_id" "$device_name" "$paired_at" "${last_seen:-$(date -u +"%Y-%m-%dT%H:%M:%SZ")}"
}

send_heartbeat() {
  local agent_id agent_secret status response last_seen device_name
  agent_id="$(read_agent_id)"
  agent_secret="$(read_agent_secret || true)"
  status="$(agent_status_json)"
  if [[ -n "$agent_secret" ]]; then
    response="$(curl -fsS -X POST "$relay_url/api/agents/$agent_id/heartbeat" -H "Content-Type: application/json" -H "X-AllReach-Agent-Token: $agent_secret" -d "$status")"
  else
    response="$(curl -fsS -X POST "$relay_url/api/agents/$agent_id/heartbeat" -H "Content-Type: application/json" -d "$status")"
  fi
  update_state_from_response "$response"
  last_seen="$(printf "%s" "$response" | sed -n 's/.*"lastSeen":"\([^"]*\)".*/\1/p')"
  device_name="$(printf "%s" "$response" | sed -n 's/.*"deviceName":"\([^"]*\)".*/\1/p')"
  printf "Heartbeat sent for %s at %s.\n" "${device_name:-unknown}" "${last_seen:-unknown}"
}

claim_pending_sessions() {
  local agent_id agent_secret response compact session_id agent_token protocol connect_port claim_body claimed_status
  agent_id="$(read_agent_id)"
  agent_secret="$(read_agent_secret || true)"
  if [[ -n "$agent_secret" ]]; then
    response="$(curl -fsS -H "X-AllReach-Agent-Token: $agent_secret" "$relay_url/api/agents/$agent_id/sessions")"
  else
    response="$(curl -fsS "$relay_url/api/agents/$agent_id/sessions")"
  fi
  compact="$(printf "%s" "$response" | tr -d '\n')"
  session_id="$(printf "%s" "$compact" | sed -n 's/.*"sessionId":"\([^"]*\)".*"status":"pending".*/\1/p')"

  if [[ -z "$session_id" ]]; then
    return
  fi

  agent_token="$(printf "%s" "$compact" | sed -n 's/.*"agentToken":"\([^"]*\)".*/\1/p')"
  protocol="$(printf "%s" "$compact" | sed -n 's/.*"protocol":"\([^"]*\)".*/\1/p')"
  connect_port="$(printf "%s" "$compact" | sed -n 's/.*"connectPort":\([0-9]*\).*/\1/p')"
  claim_body="$(printf '{"agentId":"%s","agentToken":"%s"}' "$(json_escape "$agent_id")" "$(json_escape "$agent_token")")"
  claimed_status="$(curl -fsS -X POST "$relay_url/api/sessions/$session_id/claim" -H "Content-Type: application/json" -d "$claim_body" | sed -n 's/.*"status":"\([^"]*\)".*/\1/p')"
  printf "Claimed %s session %s on local port %s with status %s.\n" "${protocol:-unknown}" "$session_id" "${connect_port:-unknown}" "${claimed_status:-unknown}"
}

if [[ -n "$pairing_code" ]]; then
  register_agent
elif [[ "$relay_url" == "http://127.0.0.1:4173" ]]; then
  saved_relay_url="$(read_saved_relay_url || true)"
  if [[ -n "${saved_relay_url:-}" ]]; then
    relay_url="$saved_relay_url"
  fi
fi

while true; do
  send_heartbeat
  claim_pending_sessions
  if [[ "$once" == true ]]; then
    break
  fi
  sleep "$interval_seconds"
done