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

server="$(realpath "$1")"
sidecar_launcher="$(realpath "$2")"
sidecar_zip="$(realpath "$3")"
copilot_cli="$(realpath "$4")"
litellm_zip="$(realpath "$5")"
litellm_config="$(realpath "$6")"
python="$(realpath "$7")"

: "${LITELLM_MASTER_KEY:?LITELLM_MASTER_KEY must be set}"

litellm_host="${LITELLM_HOST:-127.0.0.1}"
litellm_port="${LITELLM_PORT:-4000}"
default_state_root="${XDG_STATE_HOME:-$HOME/.local/state}/mrjunejune/inference"
state_root="${MRJUNEJUNE_INFERENCE_STATE:-$default_state_root}"
token_dir="${GITHUB_COPILOT_TOKEN_DIR:-$state_root/litellm-copilot}"
mkdir -p "$state_root/copilot" "$token_dir"
chmod 700 "$state_root" "$state_root/copilot" "$token_dir"

export GITHUB_COPILOT_TOKEN_DIR="$token_dir"
if [[ ! -s "$GITHUB_COPILOT_TOKEN_DIR/access-token" ]]; then
  echo "GitHub Copilot is not authenticated for LiteLLM." >&2
  echo "Run: bazel run //mrjunejune/inference:litellm_proxy -- --authenticate --token-dir \"$GITHUB_COPILOT_TOKEN_DIR\"" >&2
  exit 1
fi

litellm_pid=
server_pid=

cleanup() {
  for pid in "$server_pid" "$litellm_pid"; do
    if [[ -n "$pid" ]]; then
      kill -- "-$pid" 2>/dev/null || true
    fi
  done
  for pid in "$server_pid" "$litellm_pid"; do
    [[ -n "$pid" ]] || continue
    for _ in $(seq 1 30); do
      if ! kill -0 -- "-$pid" 2>/dev/null; then
        break
      fi
      sleep 0.1
    done
    if kill -0 -- "-$pid" 2>/dev/null; then
      kill -KILL -- "-$pid" 2>/dev/null || true
    fi
    wait "$pid" 2>/dev/null || true
  done
}
trap cleanup EXIT INT TERM

setsid "$python" "$litellm_zip" \
  --config "$litellm_config" \
  --host "$litellm_host" \
  --port "$litellm_port" &
litellm_pid=$!

for _ in $(seq 1 150); do
  if ! kill -0 "$litellm_pid" 2>/dev/null; then
    echo "LiteLLM exited before becoming ready" >&2
    exit 1
  fi
  if curl --fail --silent \
      --connect-timeout 1 \
      --max-time 2 \
      "http://${litellm_host}:${litellm_port}/health/readiness" \
      >/dev/null; then
    break
  fi
  sleep 0.1
done

if ! curl --fail --silent \
    --connect-timeout 1 \
    --max-time 2 \
    "http://${litellm_host}:${litellm_port}/health/readiness" \
    >/dev/null; then
  echo "LiteLLM readiness timed out" >&2
  exit 1
fi

export MRJUNEJUNE_PYTHON_PATH="$python"
export MRJUNEJUNE_SIDECAR_ZIP="$sidecar_zip"
export MRJUNEJUNE_INFERENCE_SIDECAR_PATH="$sidecar_launcher"
export MRJUNEJUNE_COPILOT_CLI_PATH="$copilot_cli"
export COPILOT_SIDECAR_HOME="$state_root/copilot"
export LITELLM_BASE_URL="http://${litellm_host}:${litellm_port}/v1"
export LITELLM_MODEL="${LITELLM_MODEL:-jrpg-copilot}"
export LITELLM_WIRE_API="${LITELLM_WIRE_API:-completions}"
export LITELLM_API_KEY="$LITELLM_MASTER_KEY"
export MRJUNEJUNE_ALLOW_ANONYMOUS_INFERENCE="${MRJUNEJUNE_ALLOW_ANONYMOUS_INFERENCE:-0}"

setsid "$server" &
server_pid=$!
while true; do
  if ! kill -0 "$server_pid" 2>/dev/null; then
    set +e
    wait "$server_pid"
    server_status=$?
    set -e
    exit "$server_status"
  fi
  if ! kill -0 "$litellm_pid" 2>/dev/null; then
    wait "$litellm_pid" 2>/dev/null || true
    echo "LiteLLM exited while the public server was running" >&2
    exit 1
  fi
  sleep 0.2
done