#!/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
elif [[ "${1:-}" == "--authenticate" ]]; then
  mode=authenticate
  shift
elif [[ "${1:-}" == "--check-config" ]]; then
  mode=check-config
  shift
fi
if (( $# != 0 )); then
  echo "Usage: bazel run //mrjunejune:run_inference_stack -- [--mock|--authenticate|--check-config]" >&2
  exit 2
fi

find_config() {
  local candidate
  for candidate in \
    "${BUILD_WORKSPACE_DIRECTORY:-}/mrjunejune/.config" \
    "$PWD/mrjunejune/.config" \
    "/etc/mrjunejune/.config" \
    "${HOME:-}/.config/mrjunejune/.config"; do
    [[ "$candidate" != /mrjunejune/.config ]] || continue
    [[ "$candidate" != /.config/mrjunejune/.config ]] || continue
    if [[ -f "$candidate" ]]; then
      realpath "$candidate"
      return 0
    fi
  done
  return 1
}

is_supported_config_key() {
  case "$1" in
    UPLOAD_AUTH_TOKEN|S3_REGION|S3_BUCKET|S3_URL_EXPIRES|S3_CLOUDFRONT_URL|\
    DB_PATH|MRJUNEJUNE_DB_PATH|AWS_MRJUNEJUNE_ACCESS_KEY|\
    AWS_MRJUNEJUNE_SECRET_ACCESS_KEY|AUTH_COOKIE_SECRET|\
    AUTH_BOOTSTRAP_USERNAME|AUTH_BOOTSTRAP_PASSWORD_HASH|AUTH_TRUSTED_PROXY|\
    AUTH_SESSION_IDLE_TTL|AUTH_SESSION_ABS_TTL|AUTH_GUEST_TTL|\
    AUTH_DEV_INSECURE_COOKIE|AUTH_GUEST_DAILY_TURNS|\
    AUTH_GUEST_DAILY_OUTPUT_TOKENS|AUTH_GUEST_REQUEST_OUTPUT_TOKENS|\
    SERVER_HOST|MRJUNEJUNE_PORT|MRJUNEJUNE_ALLOW_GUEST_INFERENCE|\
    MRJUNEJUNE_INFERENCE_STATE|MRJUNEJUNE_MOCK_STATE|\
    GITHUB_COPILOT_TOKEN_DIR|LITELLM_MASTER_KEY|LITELLM_HOST|LITELLM_PORT|\
    LITELLM_MODEL|LITELLM_WIRE_API|COPILOT_SESSION_IDLE_SECONDS|\
    COPILOT_MAX_SESSIONS)
      return 0
      ;;
  esac
  return 1
}

load_config() {
  local config_path="$1"
  local line key value
  while IFS= read -r line || [[ -n "$line" ]]; do
    line="${line%$'\r'}"
    [[ "$line" =~ ^[[:space:]]*$ ]] && continue
    [[ "$line" =~ ^[[:space:]]*# ]] && continue
    if [[ "$line" != *=* ]]; then
      echo "Invalid config line in $config_path" >&2
      exit 1
    fi
    key="${line%%=*}"
    value="${line#*=}"
    key="${key#"${key%%[![:space:]]*}"}"
    key="${key%"${key##*[![:space:]]}"}"
    if ! [[ "$key" =~ ^[A-Z][A-Z0-9_]*$ ]] ||
        ! is_supported_config_key "$key"; then
      echo "Unsupported config key in $config_path: $key" >&2
      exit 1
    fi
    printf -v "$key" '%s' "$value"
    export -n "$key" 2>/dev/null || true
  done < "$config_path"
}

config_file="$(find_config || true)"
if [[ -z "$config_file" ]]; then
  echo "Missing mrjunejune/.config." >&2
  echo "Copy mrjunejune/.config.development to mrjunejune/.config." >&2
  exit 1
fi
load_config "$config_file"
: "${MRJUNEJUNE_ALLOW_GUEST_INFERENCE:?MRJUNEJUNE_ALLOW_GUEST_INFERENCE must be set in $config_file}"

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}"
copilot_cache_root="$state_root/cache"
litellm_pid=
server_pid=

if [[ "$mode" == "check-config" ]]; then
  printf 'config=%s\n' "$config_file"
  if [[ -n "${LITELLM_MASTER_KEY:-}" ]]; then
    echo 'litellm_master_key=set'
  else
    echo 'litellm_master_key=missing (generated on first local live run)'
  fi
  printf 'copilot_token_dir=%s\n' \
    "${GITHUB_COPILOT_TOKEN_DIR:-$state_root/litellm-copilot}"
  printf 'copilot_cache_dir=%s\n' "$copilot_cache_root"
  printf 'guest_inference=%s\n' \
    "${MRJUNEJUNE_ALLOW_GUEST_INFERENCE:-false}"
  exit 0
fi

if [[ "$mode" == "authenticate" ]]; then
  token_dir="${GITHUB_COPILOT_TOKEN_DIR:-$state_root/litellm-copilot}"
  mkdir -p "$token_dir" "$copilot_cache_root"
  chmod 700 "$token_dir" "$copilot_cache_root"
  export XDG_CACHE_HOME="$copilot_cache_root"
  exec "$python" "$litellm_zip" --authenticate --token-dir "$token_dir"
fi

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_GUEST_INFERENCE
  # Bind to loopback for local dev; insecure cookies are only allowed here.
  export SERVER_HOST="${SERVER_HOST:-127.0.0.1}"
  export AUTH_DEV_INSECURE_COOKIE="${AUTH_DEV_INSECURE_COOKIE:-true}"
  echo "Starting mock inference at http://${SERVER_HOST}:${MRJUNEJUNE_PORT:-6969}/jrpg"
else
  if [[ -z "${LITELLM_MASTER_KEY:-}" ]]; then
    if [[ -n "${BUILD_WORKSPACE_DIRECTORY:-}" &&
          "$config_file" == "$BUILD_WORKSPACE_DIRECTORY/"* &&
          -w "$config_file" ]]; then
      umask 077
      LITELLM_MASTER_KEY="sk-$(
        "$python" -c 'import secrets; print(secrets.token_hex(24))'
      )"
      printf '\n# Generated once for the local LiteLLM proxy.\nLITELLM_MASTER_KEY=%s\n' \
        "$LITELLM_MASTER_KEY" >> "$config_file"
      chmod 600 "$config_file"
      export LITELLM_MASTER_KEY
      echo "Generated LITELLM_MASTER_KEY in $config_file"
    else
      echo "LITELLM_MASTER_KEY is missing from $config_file" >&2
      exit 1
    fi
  fi
  export LITELLM_MASTER_KEY
  token_dir="${GITHUB_COPILOT_TOKEN_DIR:-$state_root/litellm-copilot}"
  mkdir -p "$state_root/copilot" "$token_dir" "$copilot_cache_root"
  chmod 700 "$state_root" "$state_root/copilot" "$token_dir" "$copilot_cache_root"

  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:run_inference_stack -- --authenticate" >&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 XDG_CACHE_HOME="$copilot_cache_root"
  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_GUEST_INFERENCE
  # run_inference_stack is a local dev command; bind to loopback.
  export SERVER_HOST="${SERVER_HOST:-127.0.0.1}"
  export AUTH_DEV_INSECURE_COOKIE="${AUTH_DEV_INSECURE_COOKIE:-true}"
fi

export MRJUNEJUNE_CONFIG_PATH="$config_file"
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