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

script_dir="${0:A:h}"
display_state_helper="$script_dir/virtual-display-state.swift"

canonical_display_name() {
  local requested_name="$1"

  # shellcheck disable=SC2076 # zsh treats a quoted =~ operand as a regex.
  if [[ "$requested_name" =~ '^ZenttyTests( \([0-9]+\)|\+%28[0-9]+%29|%20%28[0-9]+%29)?$' ]]; then
    print -r -- "ZenttyTests"
  else
    print -r -- "$requested_name"
  fi
}

requested_display_name="${ZENTTY_TEST_SCREEN_NAME:-ZenttyTests}"
display_name="$(canonical_display_name "$requested_display_name")"
provider="${ZENTTY_TEST_DISPLAY_PROVIDER:-auto}"
display_name_slug="${display_name//[^A-Za-z0-9_.-]/_}"
display_state_dir="${ZENTTY_TEST_DISPLAY_STATE_DIR:-/tmp}"
display_lock_file="${display_state_dir}/zentty-test-display-${display_name_slug}.lock"
display_lease_dir="${display_state_dir}/zentty-test-display-${display_name_slug}.leases"
display_lease_token="${$}.${RANDOM}"
active_display_name="$display_name"
active_display_tag_id=""
original_main_display_uuid=""
display_lock_fd=""
display_lease_registered=0
display_cleanup_ran=0
betterdisplay_transport=""
betterdisplay_command_path=""

betterdisplay_command() {
  if command -v betterdisplaycli >/dev/null 2>&1; then
    command -v betterdisplaycli
  elif [[ -x /Applications/BetterDisplay.app/Contents/MacOS/BetterDisplay ]]; then
    print -r -- "/Applications/BetterDisplay.app/Contents/MacOS/BetterDisplay"
  else
    return 1
  fi
}

betterdisplay_http_base() {
  print -r -- "${ZENTTY_BETTERDISPLAY_HTTP_BASE:-http://localhost:55777}"
}

betterdisplay_http_available() {
  [[ -z "${ZENTTY_BETTERDISPLAY_COMMAND:-}" ]] || return 1
  command -v curl >/dev/null 2>&1 || return 1
  curl -fsS --max-time 2 "$(betterdisplay_http_base)/help" >/dev/null 2>&1
}

betterdisplay_http_call() {
  local operation="$1"
  shift

  local curl_args arg parameter
  curl_args=(-fsS --max-time "${ZENTTY_BETTERDISPLAY_HTTP_TIMEOUT:-10}" --get "$(betterdisplay_http_base)/${operation}")
  for arg in "$@"; do
    parameter="${arg#--}"
    parameter="${parameter#-}"
    curl_args+=(--data-urlencode "$parameter")
  done

  curl "${curl_args[@]}"
}

betterdisplay_http_request() {
  betterdisplay_http_call "$@" >/dev/null
}

resolve_betterdisplay_transport() {
  if [[ -n "$betterdisplay_transport" ]]; then
    return 0
  fi

  if [[ -n "${ZENTTY_BETTERDISPLAY_COMMAND:-}" ]]; then
    betterdisplay_transport="command"
    betterdisplay_command_path="$ZENTTY_BETTERDISPLAY_COMMAND"
    return 0
  fi

  if betterdisplay_http_available; then
    betterdisplay_transport="http"
    return 0
  fi

  if betterdisplay_command_path="$(betterdisplay_command)"; then
    betterdisplay_transport="command"
    return 0
  fi

  return 1
}

betterdisplay_request() {
  resolve_betterdisplay_transport || return 1

  case "$betterdisplay_transport" in
    command) "$betterdisplay_command_path" "$@" ;;
    http) betterdisplay_http_request "$@" ;;
    *) return 1 ;;
  esac
}

betterdisplay_query() {
  resolve_betterdisplay_transport || return 1

  case "$betterdisplay_transport" in
    command) "$betterdisplay_command_path" "$@" ;;
    http) betterdisplay_http_call "$@" ;;
    *) return 1 ;;
  esac
}

retry_betterdisplay_request_with_command_transport() {
  [[ "$betterdisplay_transport" == "http" ]] || return 1

  if ! betterdisplay_command_path="$(betterdisplay_command)"; then
    return 1
  fi

  betterdisplay_transport="command"
  "$betterdisplay_command_path" "$@"
}

required_betterdisplay_request() {
  betterdisplay_request "$@" || retry_betterdisplay_request_with_command_transport "$@"
}

print_betterdisplay_integration_help() {
  print -u2 "error: BetterDisplay integration not available"
  print -u2 "Enable BetterDisplay HTTP integration, install CLI with: brew install waydabber/betterdisplay/betterdisplaycli"
  print -u2 "Fallback supported if BetterDisplay.app is installed at /Applications/BetterDisplay.app."
  print -u2 "Tests may also provide ZENTTY_BETTERDISPLAY_COMMAND to override the BetterDisplay command."
}

capture_main_display_identity() {
  local identity display_id
  if [[ -n "${ZENTTY_TEST_MAIN_DISPLAY_COMMAND:-}" ]]; then
    display_id="$("$ZENTTY_TEST_MAIN_DISPLAY_COMMAND" "$display_name")"
    original_main_display_uuid="${ZENTTY_TEST_MAIN_DISPLAY_UUID:-TEST-MAIN-UUID}"
  else
    identity="$(/usr/bin/swift "$display_state_helper" main-display-identity "$display_name")" || return 1
    original_main_display_uuid="${identity%%$'\t'*}"
    display_id="${identity##*$'\t'}"
  fi

  # shellcheck disable=SC2076 # zsh treats a quoted =~ operand as a regex.
  if [[ ! "$display_id" =~ '^[1-9][0-9]*$' ]]; then
    print -u2 "error: could not capture the current main display before connecting '${display_name}'"
    return 1
  fi

  [[ -n "$original_main_display_uuid" ]] || return 1
}

resolve_original_main_display_id() {
  if [[ -n "${ZENTTY_TEST_RESOLVE_DISPLAY_ID_COMMAND:-}" ]]; then
    "$ZENTTY_TEST_RESOLVE_DISPLAY_ID_COMMAND" "$original_main_display_uuid"
  else
    /usr/bin/swift "$display_state_helper" display-id-for-uuid "$original_main_display_uuid"
  fi
}

raw_main_display_id() {
  if [[ -n "${ZENTTY_TEST_RAW_MAIN_DISPLAY_COMMAND:-}" ]]; then
    "$ZENTTY_TEST_RAW_MAIN_DISPLAY_COMMAND"
  else
    /usr/bin/swift "$display_state_helper" raw-main-display-id
  fi
}

restore_original_main_display() {
  [[ -n "$original_main_display_uuid" ]] || return 0

  local attempt current_main expected_main
  expected_main="$(resolve_original_main_display_id)" || return 1
  current_main="$(raw_main_display_id)" || return 1
  if [[ "$current_main" == "$expected_main" ]]; then
    return 0
  fi
  if [[ "$resolved_provider" != "betterdisplay" ]]; then
    print -u2 "error: provider '${resolved_provider}' cannot restore display ${original_main_display_uuid} as main"
    return 1
  fi
  if ! required_betterdisplay_request set -UUID="$original_main_display_uuid" -main=on >/dev/null 2>&1; then
    print -u2 "error: could not restore display ${original_main_display_uuid} as the main display during cleanup"
    return 1
  fi
  for attempt in {1..20}; do
    expected_main="$(resolve_original_main_display_id)" || return 1
    current_main="$(raw_main_display_id)" || return 1
    [[ "$current_main" == "$expected_main" ]] && return 0
    sleep 0.1
  done
  print -u2 "error: display ${original_main_display_uuid} did not remain the main display during cleanup"
  return 1
}

display_topology_status() {
  local test_display_name="$1"
  local expected_main_display_id="$2"

  if [[ -n "${ZENTTY_TEST_TOPOLOGY_STATUS_COMMAND:-}" ]]; then
    "$ZENTTY_TEST_TOPOLOGY_STATUS_COMMAND" "$test_display_name" "$expected_main_display_id"
    return
  fi

  /usr/bin/swift "$display_state_helper" topology-status "$test_display_name" "$expected_main_display_id"
}

verify_display_topology() {
  local expected_main_display_id topology_state

  expected_main_display_id="$(resolve_original_main_display_id)" || return 1
  topology_state="$(display_topology_status "$active_display_name" "$expected_main_display_id")"
  case "$topology_state" in
    ok) return 0 ;;
    mirrored)
      print -u2 "error: virtual display '${display_name}' is mirrored; refusing to run tests that would change another display's resolution"
      ;;
    main-changed)
      if [[ "$resolved_provider" != "betterdisplay" ]]; then
        print -u2 "error: virtual display '${display_name}' became the main display and provider '${resolved_provider}' cannot restore it automatically"
        return 1
      fi
      if ! required_betterdisplay_request set -UUID="$original_main_display_uuid" -main=on >/dev/null 2>&1; then
        print -u2 "error: could not restore display ${original_main_display_uuid} as the main display"
        return 1
      fi

      local attempt
      for attempt in {1..20}; do
        expected_main_display_id="$(resolve_original_main_display_id)" || return 1
        topology_state="$(display_topology_status "$active_display_name" "$expected_main_display_id")"
        if [[ "$topology_state" == "ok" ]]; then
          return 0
        fi
        if [[ "$topology_state" == "mirrored" ]]; then
          print -u2 "error: virtual display '${display_name}' became mirrored while restoring the main display"
          return 1
        fi
        sleep 0.1
      done

      print -u2 "error: display ${expected_main_display_id} did not return as the main display"
      return 1
      ;;
    missing)
      print -u2 "error: virtual display '${display_name}' is missing during topology verification"
      ;;
    *)
      print -u2 "error: could not verify the display topology for '${display_name}'"
      ;;
  esac
  return 1
}

registered_betterdisplay_plan() {
  local registered_identifiers

  if [[ -n "${ZENTTY_TEST_REGISTERED_SCREEN_COMMAND:-}" ]]; then
    "$ZENTTY_TEST_REGISTERED_SCREEN_COMMAND" "$display_name"
    return
  fi

  if ! registered_identifiers="$(betterdisplay_query get -type=VirtualScreen -identifiers 2>/dev/null)"; then
    return 2
  fi

  ZENTTY_REGISTERED_SCREEN_IDENTIFIERS="$registered_identifiers" \
    /usr/bin/swift "$display_state_helper" registered-screen-plan "$display_name"
}

matched_screen_name() {
  if [[ -n "${ZENTTY_TEST_SCREEN_EXISTS_COMMAND:-}" ]]; then
    local matched_name
    matched_name="$("$ZENTTY_TEST_SCREEN_EXISTS_COMMAND" "$display_name" 2>/dev/null)" || return 1
    print -r -- "${matched_name:-$display_name}"
    return
  fi

  /usr/bin/swift "$display_state_helper" matched-screen-name "$display_name" 2>/dev/null
}

screen_exists() {
  matched_screen_name >/dev/null 2>&1
}

acquire_display_lock() {
  zmodload zsh/system
  touch "$display_lock_file"

  if ! zsystem flock -t 60 -f display_lock_fd "$display_lock_file"; then
    print -u2 "error: timed out waiting for virtual display lock '${display_lock_file}'"
    return 1
  fi
}

release_display_lock() {
  if [[ -n "$display_lock_fd" ]]; then
    zsystem flock -u "$display_lock_fd" >/dev/null 2>&1 || true
    display_lock_fd=""
  fi
}

prune_stale_display_leases() {
  local lease pid

  mkdir -p "$display_lease_dir"
  for lease in "$display_lease_dir"/*(.N); do
    pid="${lease:t}"
    pid="${pid%%.*}"
    if [[ -z "$pid" ]] || ! kill -0 "$pid" >/dev/null 2>&1; then
      command rm -f "$lease"
    fi
  done
}

register_display_lease() {
  mkdir -p "$display_lease_dir"
  print -r -- "$$" > "$display_lease_dir/$display_lease_token"
  display_lease_registered=1
}

release_display_lease() {
  if [[ "$display_lease_registered" != "1" ]]; then
    return
  fi

  command rm -f "$display_lease_dir/$display_lease_token"
  display_lease_registered=0
  prune_stale_display_leases
}

has_active_display_leases() {
  local leases

  leases=("$display_lease_dir"/*(.N))
  (( ${#leases[@]} > 0 ))
}

wait_for_screen() {
  local attempt
  # shellcheck disable=SC2034 # zsh needs a loop variable even though its value is unused.
  for attempt in {1..30}; do
    if screen_exists; then
      return 0
    fi
    sleep 0.5
  done

  print -u2 "error: virtual display '${display_name}' did not appear in NSScreen.screens"
  return 1
}

resolve_provider() {
  case "$provider" in
    auto)
      if resolve_betterdisplay_transport; then
        print -r -- "betterdisplay"
      elif command -v simpledisplayctl >/dev/null 2>&1; then
        print -r -- "simpledisplay"
      else
        print -u2 "error: no virtual display provider integration found"
        print_betterdisplay_integration_help
        print -u2 "Or install SimpleDisplay CLI and set ZENTTY_TEST_DISPLAY_PROVIDER=simpledisplay."
        return 1
      fi
      ;;
    betterdisplay|simpledisplay)
      print -r -- "$provider"
      ;;
    *)
      print -u2 "error: unsupported ZENTTY_TEST_DISPLAY_PROVIDER='${provider}'"
      print -u2 "Supported values: auto, betterdisplay, simpledisplay"
      return 1
      ;;
  esac
}

connect_existing_betterdisplay() {
  local tag_id="$1"
  betterdisplay_request set -tagID="$tag_id" -connected=on >/dev/null 2>&1
}

# macOS remembers the last arrangement for a display pair. If the test screen
# was ever mirrored onto the built-in display, every reconnect restores that
# hardware mirror and the laptop screen shows the test display. Break the
# mirror right after connecting; the unmirrored layout is what macOS then
# remembers for the next run.
unmirror_betterdisplay() {
  local tag_id="$1" builtin_uuid target_uuid
  if [[ -n "${ZENTTY_TEST_MAIN_DISPLAY_COMMAND:-}" ]]; then
    return 0
  fi
  builtin_uuid="$(/usr/bin/swift "$display_state_helper" builtin-display-uuid 2>/dev/null || true)"
  for target_uuid in "$builtin_uuid" "$original_main_display_uuid"; do
    [[ -n "$target_uuid" ]] || continue
    betterdisplay_request set -tagID="$tag_id" -mirror=off -targetUUID="$target_uuid" >/dev/null 2>&1 || true
  done
}

resolve_registered_betterdisplay() {
  local registration_plan registration_status line kind record_fields tag_id display_id
  local canonical_tag_id=""
  local -a legacy_records

  if registration_plan="$(registered_betterdisplay_plan)"; then
    # shellcheck disable=SC2066,SC2296 # ${(@f)...} is zsh newline splitting.
    for line in "${(@f)registration_plan}"; do
      kind="${line%%:*}"
      record_fields="${line#*:}"
      tag_id="${record_fields%%:*}"
      display_id="${line##*:}"
      case "$kind" in
        canonical) canonical_tag_id="$tag_id" ;;
        legacy) legacy_records+=("$tag_id:$display_id") ;;
        *)
          print -u2 "error: invalid BetterDisplay registration record '${line}'"
          return 2
          ;;
      esac
    done
  else
    registration_status=$?
    return "$registration_status"
  fi

  if [[ -n "$canonical_tag_id" ]]; then
    local legacy_record legacy_tag_id legacy_display_id
    for legacy_record in "${legacy_records[@]}"; do
      legacy_tag_id="${legacy_record%%:*}"
      legacy_display_id="${legacy_record##*:}"
      if [[ "$legacy_display_id" != "0" ]]; then
        print -u2 "error: legacy Zentty test display tagID ${legacy_tag_id} is still connected; refusing to alter an active duplicate"
        return 2
      fi
      if ! required_betterdisplay_request discard -tagID="$legacy_tag_id" >/dev/null 2>&1; then
        print -u2 "error: could not discard disconnected legacy Zentty test display tagID ${legacy_tag_id}"
        return 2
      fi
    done
    active_display_tag_id="$canonical_tag_id"
    return 0
  fi

  if [[ "${#legacy_records[@]}" == "1" ]]; then
    tag_id="${legacy_records[1]%%:*}"
    display_id="${legacy_records[1]##*:}"
    if [[ "$display_id" != "0" ]]; then
      print -u2 "error: legacy Zentty test display tagID ${tag_id} is still connected; disconnect it before retrying"
      return 2
    fi
    if ! required_betterdisplay_request set -tagID="$tag_id" -virtualScreenName="$display_name" >/dev/null 2>&1; then
      print -u2 "error: could not rename legacy Zentty test display tagID ${tag_id} to '${display_name}'"
      return 2
    fi
    active_display_tag_id="$tag_id"
    return 0
  fi

  print -u2 "error: multiple legacy Zentty test displays are registered; refusing to guess which one to reuse"
  return 2
}

ensure_betterdisplay() {
  if ! resolve_betterdisplay_transport; then
    print_betterdisplay_integration_help
    return 1
  fi

  local registration_status
  if resolve_registered_betterdisplay; then
    if screen_exists; then
      active_display_name="$(matched_screen_name || print -r -- "$display_name")"
      print -r -- "Using existing display '${active_display_name}'. It will be disconnected by the last test run using it."
      return 0
    fi

    if ! connect_existing_betterdisplay "$active_display_tag_id"; then
      print -u2 "error: registered virtual display tagID ${active_display_tag_id} could not be connected; refusing to create a duplicate"
      return 1
    fi
    wait_for_screen
    unmirror_betterdisplay "$active_display_tag_id"
    active_display_name="$(matched_screen_name || print -r -- "$display_name")"
    print -r -- "Connected existing BetterDisplay virtual display '${active_display_name}'. It will be disconnected by the last test run using it."
    return 0
  else
    registration_status=$?
  fi

  if [[ "$registration_status" != "1" ]]; then
    print -u2 "error: could not safely resolve the registered BetterDisplay '${display_name}' screen; refusing to create a duplicate"
    return 1
  fi

  required_betterdisplay_request create -type=VirtualScreen -virtualScreenName="$display_name" -aspectWidth=16 -aspectHeight=10 -virtualScreenHiDPI=on
  if ! resolve_registered_betterdisplay || [[ -z "$active_display_tag_id" ]]; then
    print -u2 "error: BetterDisplay created '${display_name}' but did not return its stable tagID"
    return 1
  fi
  required_betterdisplay_request set -tagID="$active_display_tag_id" -connected=on
  wait_for_screen
  unmirror_betterdisplay "$active_display_tag_id"
  active_display_name="$(matched_screen_name || print -r -- "$display_name")"
}

cleanup_betterdisplay() {
  resolve_betterdisplay_transport || return 0
  # Disconnect rather than discard. Discarding destroys the virtual screen, so
  # macOS mints a fresh display UUID on the next create. Menu-bar managers like
  # Thaw/Ice key per-display settings by UUID and never prune them, so a
  # discard-per-run leaks a new ghost display into their config every run; it
  # also lets concurrent agents stack duplicate "Name (1)" displays. Keeping the
  # screen registered preserves its UUID, so the next run reconnects the same
  # screen (by stable tagID) and Thaw keeps one entry. BetterDisplay can briefly
  # reconnect a screen during layout restoration, so verify that disconnect held.
  [[ -n "$active_display_tag_id" ]] || return 0

  local attempt registration_plan line record_fields tag_id display_id
  # shellcheck disable=SC2034 # zsh needs a loop variable even though its value is unused.
  for attempt in {1..5}; do
    betterdisplay_request set -tagID="$active_display_tag_id" -connected=off >/dev/null 2>&1 || true
    sleep 0.2

    display_id=""
    # shellcheck disable=SC2034 # Consumed by zsh's ${(@f)...} expansion below.
    if registration_plan="$(registered_betterdisplay_plan 2>/dev/null)"; then
      # shellcheck disable=SC2066,SC2296 # ${(@f)...} is zsh newline splitting.
      for line in "${(@f)registration_plan}"; do
        record_fields="${line#*:}"
        tag_id="${record_fields%%:*}"
        if [[ "$tag_id" == "$active_display_tag_id" ]]; then
          display_id="${line##*:}"
          break
        fi
      done
    fi
    if [[ "$display_id" == "0" ]]; then
      restore_original_main_display
      return
    fi
  done

  if ! required_betterdisplay_request discard -tagID="$active_display_tag_id" >/dev/null 2>&1; then
    restore_original_main_display || true
    print -u2 "error: BetterDisplay kept reconnecting virtual display tagID ${active_display_tag_id}, and emergency discard failed"
    return 1
  fi
  restore_original_main_display || true
  print -u2 "error: BetterDisplay kept reconnecting virtual display tagID ${active_display_tag_id}; it was discarded for display safety"
  return 1
}

ensure_simpledisplay() {
  if ! command -v simpledisplayctl >/dev/null 2>&1; then
    print -u2 "error: simpledisplayctl not found"
    return 1
  fi

  if screen_exists; then
    active_display_name="$(matched_screen_name || print -r -- "$display_name")"
    print -r -- "Using existing display '${active_display_name}'. It will be removed by the last test run using it."
    return 0
  fi

  simpledisplayctl create --width 2732 --height 2048 --name "$display_name" --hidpi
  wait_for_screen
  active_display_name="$(matched_screen_name || print -r -- "$display_name")"
}

cleanup_simpledisplay() {
  command -v simpledisplayctl >/dev/null 2>&1 || return 0
  simpledisplayctl remove --name "$active_display_name" >/dev/null 2>&1 || \
    simpledisplayctl remove --name "$display_name" >/dev/null 2>&1 || true
}

cleanup_display_if_last_lease() {
  local already_locked=0
  local cleanup_status=0

  if [[ "$display_cleanup_ran" == "1" ]]; then
    return
  fi
  display_cleanup_ran=1

  if [[ -n "$display_lock_fd" ]]; then
    already_locked=1
  else
    acquire_display_lock
  fi

  release_display_lease

  if ! has_active_display_leases; then
    case "$resolved_provider" in
      betterdisplay) cleanup_betterdisplay || cleanup_status=$? ;;
      simpledisplay)
        cleanup_simpledisplay || cleanup_status=$?
        restore_original_main_display || cleanup_status=$?
        ;;
    esac
    command rmdir "$display_lease_dir" >/dev/null 2>&1 || true
  fi

  if [[ "$already_locked" == "0" ]]; then
    release_display_lock
  fi

  return "$cleanup_status"
}

resolved_provider="$(resolve_provider)"
trap cleanup_display_if_last_lease EXIT INT TERM

register_display_lease
acquire_display_lock
prune_stale_display_leases
capture_main_display_identity
case "$resolved_provider" in
  betterdisplay) ensure_betterdisplay ;;
  simpledisplay) ensure_simpledisplay ;;
esac
if ! verify_display_topology; then
  cleanup_display_if_last_lease
  exit 1
fi
release_display_lock

export ZENTTY_TEST_SCREEN_NAME="$active_display_name"

if TEST_RUNNER_SWIFT_BACKTRACE=enable=no \
    TEST_RUNNER_ZENTTY_TEST_SCREEN_NAME="$active_display_name" \
    /usr/bin/caffeinate -di xcodebuild test \
    -scheme Zentty \
    -destination 'platform=macOS' \
    "$@"; then
  test_status=0
else
  test_status=$?
fi

cleanup_status=0
cleanup_display_if_last_lease || cleanup_status=$?
trap - EXIT INT TERM

if [[ "$cleanup_status" != "0" ]]; then
  exit "$cleanup_status"
fi
exit "$test_status"
