#!/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")"
mock_sidecar="${8:-}"
if (( $# >= 8 )); then
  shift 8
else
  shift 7
fi

mode=live
if [[ "${1:-}" == "--mock" ]]; then
  mode=mock
  shift
fi
if (( $# != 0 )); then
  echo "Usage: bazel run //mrjunejune:run_inference_stack -- [--mock]" >&2
  exit 2
fi

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}"
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

if [[ "$mode" == "mock" ]]; then
  if [[ -z "$mock_sidecar" ]]; then
    echo "Mock sidecar is unavailable in this runtime" >&2
    exit 1
  fi
  mock_state="${MRJUNEJUNE_MOCK_STATE:-$state_root/mock}"
  mkdir -p "$mock_state"
  chmod 700 "$mock_state"
  export MRJUNEJUNE_DB_PATH="${MRJUNEJUNE_DB_PATH:-$mock_state/conversations.db}"
  if [[ "$mock_sidecar" == /* ]]; then
    export MRJUNEJUNE_INFERENCE_SIDECAR_PATH="$mock_sidecar"
  else
    export MRJUNEJUNE_INFERENCE_SIDECAR_PATH="$PWD/$mock_sidecar"
  fi
  export MRJUNEJUNE_COPILOT_CLI_PATH="$MRJUNEJUNE_INFERENCE_SIDECAR_PATH"
  export MRJUNEJUNE_ALLOW_ANONYMOUS_INFERENCE=1
  echo "Starting mock inference at http://127.0.0.1:${MRJUNEJUNE_PORT:-6969}/jrpg"
else
  : "${LITELLM_MASTER_KEY:?LITELLM_MASTER_KEY must be set}"
  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

  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}"
fi

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 [[ -n "$litellm_pid" ]] && ! 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