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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT="${ROOT:-$(cd "${SCRIPT_DIR}/.." && pwd)}"
PD_IMAGE="${PD_IMAGE:-hub-new.pingcap.net/keyspace/pd:v8.5.0}"
CLIENT_PORT="${CLIENT_PORT:-2379}"
PEER_PORT="${PEER_PORT:-2380}"
TSO_PORT="${TSO_PORT:-2479}"
RUNNER_LABEL="${RUNNER_LABEL:-local}"
PD_RS_INTEGRATION_BASENAME="${PD_RS_INTEGRATION_BASENAME:-pd-rs-integration-${RUNNER_LABEL}-${GITHUB_RUN_ID:-manual}-${GITHUB_RUN_ATTEMPT:-1}}"
PD_SOURCE_KIND="${PD_SOURCE_KIND:-}"
PD_SOURCE_DETAIL="${PD_SOURCE_DETAIL:-}"

require_command() {
  if ! command -v "$1" >/dev/null 2>&1; then
    echo "missing required command: $1" >&2
    exit 1
  fi
}

pick_tmp_dir() {
  local dir
  for dir in \
    "/tmp/${PD_RS_INTEGRATION_BASENAME}" \
    "${TMPDIR:-/tmp}/${PD_RS_INTEGRATION_BASENAME}" \
    "/var/tmp/${PD_RS_INTEGRATION_BASENAME}" \
    "${ROOT}/.tmp/${PD_RS_INTEGRATION_BASENAME}"
  do
    if mkdir -p "$dir" 2>/dev/null && [ -w "$dir" ]; then
      printf '%s\n' "$dir"
      return 0
    fi
  done
  return 1
}

check_port_free() {
  local port="$1"
  shift
  python3 - "$port" "$@" <<'PY'
import socket
import sys

port = int(sys.argv[1])
reserved = {int(value) for value in sys.argv[2:] if value}
if port in reserved:
    sys.exit(1)

s = socket.socket()
s.settimeout(0.2)
try:
    s.connect(("127.0.0.1", port))
except Exception:
    sys.exit(0)
else:
    sys.exit(1)
finally:
    s.close()
PY
}

pick_port() {
  local preferred="$1"
  shift
  local port
  for port in $(seq "$preferred" "$((preferred + 100))"); do
    if check_port_free "$port" "$@"; then
      printf '%s\n' "$port"
      return 0
    fi
  done
  echo "unable to find a free port near ${preferred}" >&2
  exit 1
}

wait_for_port() {
  local port="$1"
  for _ in $(seq 1 120); do
    if python3 - "$port" <<'PY'
import socket
import sys

port = int(sys.argv[1])
s = socket.socket()
s.settimeout(0.2)
try:
    s.connect(("127.0.0.1", port))
except Exception:
    sys.exit(1)
else:
    sys.exit(0)
finally:
    s.close()
PY
    then
      return 0
    fi
    sleep 1
  done
  echo "timed out waiting for 127.0.0.1:${port}" >&2
  return 1
}

pd_dir_from_dotenv() {
  local dotenv_file="${ROOT}/.env"
  local value

  if [[ ! -f "${dotenv_file}" ]]; then
    return 1
  fi

  value="$(awk -F= '/^PD_DIR=/{sub(/^PD_DIR=/, ""); print; exit}' "${dotenv_file}")"
  if [[ -z "${value}" ]]; then
    return 1
  fi

  value="${value%$'\r'}"
  if [[ "${value}" =~ ^\".*\"$ || "${value}" =~ ^\'.*\'$ ]]; then
    value="${value:1:${#value}-2}"
  fi

  if [[ -z "${value}" ]]; then
    return 1
  fi

  if [[ "${value}" != /* ]]; then
    value="${ROOT}/${value}"
  fi

  printf '%s\n' "${value}"
}

resolve_pd_dir() {
  if [[ "${PD_FORCE_IMAGE:-0}" == "1" ]]; then
    PD_DIR=""
    PD_SOURCE_KIND=""
    PD_SOURCE_DETAIL=""
    return 0
  fi

  if [[ -n "${PD_DIR:-}" ]]; then
    PD_SOURCE_KIND="checkout"
    PD_SOURCE_DETAIL="env:PD_DIR"
    return 0
  fi

  if [[ -d "${ROOT}/../pd" ]]; then
    PD_DIR="${ROOT}/../pd"
    PD_SOURCE_KIND="checkout"
    PD_SOURCE_DETAIL="default:../pd"
    return 0
  fi

  local dotenv_pd_dir
  if dotenv_pd_dir="$(pd_dir_from_dotenv)"; then
    PD_DIR="${dotenv_pd_dir}"
    PD_SOURCE_KIND="checkout"
    PD_SOURCE_DETAIL="dotenv:${ROOT}/.env"
    return 0
  fi

  PD_DIR=""
  PD_SOURCE_KIND=""
  PD_SOURCE_DETAIL=""
}

prepare_pd_image_dir() {
  local container_id
  local copied=0
  local candidate

  require_command docker

  PD_IMAGE_DIR="${PD_IMAGE_DIR:-${PD_RS_INTEGRATION_DIR}/pd-image}"
  export PD_IMAGE_DIR

  rm -rf "${PD_IMAGE_DIR}"
  mkdir -p "${PD_IMAGE_DIR}/bin"

  docker pull "${PD_IMAGE}"
  container_id="$(docker create "${PD_IMAGE}")"
  for candidate in \
    /pd-server \
    /pd/bin/pd-server \
    /bin/pd-server \
    /usr/local/bin/pd-server
  do
    if docker cp "${container_id}:${candidate}" "${PD_IMAGE_DIR}/bin/pd-server" >/dev/null 2>&1; then
      copied=1
      break
    fi
  done
  docker rm -f "${container_id}" >/dev/null 2>&1 || true

  if [[ "${copied}" -ne 1 ]]; then
    echo "failed to find pd-server inside ${PD_IMAGE}" >&2
    exit 1
  fi

  chmod +x "${PD_IMAGE_DIR}/bin/pd-server"
  PD_DIR="${PD_IMAGE_DIR}"
  PD_SOURCE_KIND="image"
  PD_SOURCE_DETAIL="${PD_IMAGE}"
}

ensure_pd_dir() {
  resolve_pd_dir

  if [[ -n "${PD_DIR:-}" ]]; then
    if [[ ! -d "${PD_DIR}" ]]; then
      echo "PD checkout not found at ${PD_DIR} (${PD_SOURCE_DETAIL})" >&2
      exit 1
    fi
    export PD_DIR
    return 0
  fi

  if [[ -z "${PD_IMAGE:-}" ]]; then
    echo "PD checkout not found at ../pd or ${ROOT}/.env, and no PD_IMAGE fallback was configured" >&2
    exit 1
  fi

  prepare_pd_image_dir
  export PD_DIR
}

ensure_pd_server() {
  if [[ "${PD_SOURCE_KIND}" == "checkout" ]]; then
    require_command make
    make -C "${PD_DIR}"
  fi

  if [[ ! -x "${PD_DIR}/bin/pd-server" ]]; then
    echo "pd-server not found at ${PD_DIR}/bin/pd-server" >&2
    exit 1
  fi
}

start_pd_rs_integration_env() {
  require_command python3

  CLIENT_PORT="$(pick_port "${CLIENT_PORT}")"
  PEER_PORT="$(pick_port "${PEER_PORT}" "${CLIENT_PORT}")"
  TSO_PORT="$(pick_port "${TSO_PORT}" "${CLIENT_PORT}" "${PEER_PORT}")"
  export CLIENT_PORT PEER_PORT TSO_PORT

  PD_RS_INTEGRATION_DIR="${PD_RS_INTEGRATION_DIR:-$(pick_tmp_dir)}"
  export PD_RS_INTEGRATION_DIR

  LOG_DIR="${PD_RS_INTEGRATION_DIR}/logs"
  export LOG_DIR

  rm -rf "${PD_RS_INTEGRATION_DIR}/pd" "${LOG_DIR}"
  mkdir -p "${PD_RS_INTEGRATION_DIR}/pd" "${LOG_DIR}"

  ensure_pd_dir
  ensure_pd_server
  export PD_SOURCE_KIND PD_SOURCE_DETAIL

  "${PD_DIR}/bin/pd-server" \
    --name=pd-rs-integration \
    --data-dir="${PD_RS_INTEGRATION_DIR}/pd" \
    --client-urls="http://127.0.0.1:${CLIENT_PORT}" \
    --advertise-client-urls="http://127.0.0.1:${CLIENT_PORT}" \
    --peer-urls="http://127.0.0.1:${PEER_PORT}" \
    --advertise-peer-urls="http://127.0.0.1:${PEER_PORT}" \
    >"${LOG_DIR}/pd.log" 2>&1 &
  PD_PID=$!
  export PD_PID

  wait_for_port "${CLIENT_PORT}"

  "${PD_DIR}/bin/pd-server" services tso \
    --backend-endpoints="http://127.0.0.1:${CLIENT_PORT}" \
    --listen-addr="http://127.0.0.1:${TSO_PORT}" \
    --advertise-listen-addr="http://127.0.0.1:${TSO_PORT}" \
    >"${LOG_DIR}/tso.log" 2>&1 &
  TSO_PID=$!
  export TSO_PID

  wait_for_port "${TSO_PORT}"
}

stop_pd_rs_integration_env() {
  set +e
  if [[ -n "${TSO_PID:-}" ]]; then
    kill "${TSO_PID}" 2>/dev/null || true
    wait "${TSO_PID}" 2>/dev/null || true
  fi
  if [[ -n "${PD_PID:-}" ]]; then
    kill "${PD_PID}" 2>/dev/null || true
    wait "${PD_PID}" 2>/dev/null || true
  fi
}