#!/usr/bin/env bash
set -euo pipefail
export PATH="/usr/bin:/bin:/usr/sbin:/sbin:/usr/local/bin:/opt/homebrew/bin"

CODEX_MODEL="${CODEX_MODEL:-gpt-5.6-sol}"
case "$CODEX_MODEL" in
  *[!A-Za-z0-9._-]*|"") echo "invalid CODEX_MODEL: $CODEX_MODEL" >&2; exit 1 ;;
esac
BAIJIMU_WORKSPACE_ID="${CODEX_WORKSPACE_ID:-${BAIJIMU_WORKSPACE_ID:-${WORKSPACE_ID:-}}}"
BAIJIMU_PROJECT_ID="${CODEX_PROJECT_ID:-${BAIJIMU_PROJECT_ID:-${PROJECT_ID:-}}}"
BAIJIMU_AGENT_CONFIG_ID="${CODEX_AGENT_CONFIG_ID:-${BAIJIMU_AGENT_CONFIG_ID:-}}"
BAIJIMU_AGENT_SESSION_ID="${CODEX_AGENT_SESSION_ID:-${BAIJIMU_AGENT_SESSION_ID:-}}"
BAIJIMU_SESSION_ID="${CODEX_SESSION_ID:-${BAIJIMU_SESSION_ID:-${SESSION_ID:-}}}"
ROUTER_BASE_URL="${CODEX_ROUTER_BASE_URL:-https://router.baijimu.com/api/claudecode/v1}"
case "$BAIJIMU_WORKSPACE_ID" in *[!0-9]*|"") echo "CODEX_WORKSPACE_ID or BAIJIMU_WORKSPACE_ID is required" >&2; exit 1 ;; esac
case "$BAIJIMU_PROJECT_ID" in *[!0-9]*|"") echo "CODEX_PROJECT_ID or BAIJIMU_PROJECT_ID is required" >&2; exit 1 ;; esac

started_at="$(date -u +"%Y-%m-%dT%H:%M:%SZ")"
start_epoch="$(date +%s)"
state_dir="${CODEX_INSTALL_STATE_DIR:-${TMPDIR:-/tmp}/baijimu-codex-install}"
status_path="$state_dir/status.json"
result_path="$state_dir/result.json"
mkdir -p "$state_dir"
llm_credential_created=false

current_step=0
step_count=4
step1_name="Create Baijimu LLM credential and config"; step1_state="pending"; step1_detail=""
step2_name="Verify Baijimu router"; step2_state="pending"; step2_detail=""
step3_name="Verify Codex CLI"; step3_state="pending"; step3_detail=""
step4_name="Finish terminal setup"; step4_state="pending"; step4_detail=""

json_escape() {
  printf '%s' "${1:-}" | awk 'BEGIN { ORS = ""; first = 1 } {
    if (!first) { printf "\\n" }
    first = 0
    gsub(/\\/, "\\\\")
    gsub(/"/, "\\\"")
    gsub(/\r/, "\\r")
    gsub(/\t/, "\\t")
    printf "%s", $0
  }'
}

json_string() { printf '"'; json_escape "${1:-}"; printf '"'; }
write_install_console() { [ "${CODEX_INSTALL_QUIET:-}" = "1" ] && return 0; printf '%s\n' "$*" >&2; }

write_status() {
  {
    printf '{\n'
    printf '  "title": "Baijimu is configuring Codex Terminal",\n'
    printf '  "platform": "macos",\n'
    printf '  "startedAt": '; json_string "$started_at"; printf ',\n'
    printf '  "updatedAt": '; json_string "$(date -u +"%Y-%m-%dT%H:%M:%SZ")"; printf ',\n'
    printf '  "currentStep": %s,\n' "$current_step"
    printf '  "statusPath": '; json_string "$status_path"; printf ',\n'
    printf '  "resultPath": '; json_string "$result_path"; printf ',\n'
    printf '  "steps": [\n'
    for index in 1 2 3 4; do
      eval "name=\${step${index}_name}"
      eval "state=\${step${index}_state}"
      eval "detail=\${step${index}_detail}"
      printf '    { "index": %s, "name": ' "$index"; json_string "$name"; printf ', "state": '; json_string "$state"
      printf ', "detail": '; json_string "$detail"; printf ', "downloadedBytes": null, "totalBytes": null }'
      [ "$index" -lt "$step_count" ] && printf ','
      printf '\n'
    done
    printf '  ]\n'
    printf '}\n'
  } > "$status_path"
}

set_step() {
  index="$1"
  state="$2"
  detail="${3:-}"
  current_step="$index"
  eval "step${index}_state=\$state"
  eval "step${index}_detail=\$detail"
  write_status
  eval "name=\${step${index}_name}"
  if [ -n "$detail" ]; then
    write_install_console "[$index/$step_count] $name  $state  $detail"
  else
    write_install_console "[$index/$step_count] $name  $state"
  fi
}

finish_result() {
  ok="$1"
  router_status="${2:-}"
  cli_version="${3:-}"
  cli_smoke="$4"
  error_message="${5:-}"
  elapsed_ms="$(( ($(date +%s) - start_epoch) * 1000 ))"
  {
    printf '{\n'
    printf '  "ok": %s,\n' "$ok"
    printf '  "platform": "macos",\n'
    printf '  "startedAt": '; json_string "$started_at"; printf ',\n'
    printf '  "codexHome": '; json_string "$HOME/.codex"; printf ',\n'
    printf '  "configWritten": true,\n'
    printf '  "authWritten": true,\n'
    printf '  "workspaceId": %s,\n' "$BAIJIMU_WORKSPACE_ID"
    printf '  "projectId": %s,\n' "$BAIJIMU_PROJECT_ID"
    printf '  "llmCredentialCreated": %s,\n' "$llm_credential_created"
    printf '  "routerHttpStatus": '; [ -n "$router_status" ] && printf '%s' "$router_status" || printf 'null'; printf ',\n'
    printf '  "cliVersion": '; json_string "$cli_version"; printf ',\n'
    printf '  "cliSmoke": %s,\n' "$cli_smoke"
    printf '  "model": '; json_string "$CODEX_MODEL"; printf ',\n'
    printf '  "elapsedMs": %s,\n' "$elapsed_ms"
    if [ -n "$error_message" ]; then
      printf '  "warnings": [],\n'
      printf '  "errors": [ '; json_string "$error_message"; printf ' ]\n'
    else
      printf '  "warnings": [],\n'
      printf '  "errors": []\n'
    fi
    printf '}\n'
  } > "$result_path"
  cat "$result_path"
}

write_install_console ""
write_install_console "Baijimu is configuring Codex Terminal"
write_install_console "Please keep this window open."
write_install_console ""
write_status

shared_auth_path() {
  if [ -n "${BAIJIMU_CONFIG_HOME:-}" ]; then
    printf '%s\n' "$BAIJIMU_CONFIG_HOME/baijimu/auth.json"
  else
    printf '%s\n' "$HOME/.config/baijimu/auth.json"
  fi
}

extract_llm_credential_from_json_file() {
  output_file="$1"
  osascript -l JavaScript - "$output_file" <<'JS'
ObjC.import("Foundation");

function run(argv) {
  const path = argv[0];
  const data = $.NSData.dataWithContentsOfFile(path);
  if (!data) {
    return "";
  }
  const document = ObjC.deepUnwrap($.NSJSONSerialization.JSONObjectWithDataOptionsError(data, 0, null));
  const payload = document && document.data ? document.data : document;
  const credential = payload.llmCredential || payload.credential || payload.apiKey || "";
  if (typeof credential === "string") {
    return credential;
  }
  return "";
}
JS
}

create_baijimu_llm_credential() {
  command -v "${BAIJIMU_CLI_BIN:-baijimu}" >/dev/null 2>&1 || return 1
  output_file="$(mktemp "${TMPDIR:-/tmp}/baijimu-llm-credential.XXXXXX")"
  error_file="$state_dir/baijimu-llm-credential.err"
  chmod 600 "$output_file"

  cmd=("${BAIJIMU_CLI_BIN:-baijimu}" --json llm-credential create
    --workspace-id "$BAIJIMU_WORKSPACE_ID"
    --project-id "$BAIJIMU_PROJECT_ID"
    --show-secret)
  if [ -n "${BAIJIMU_AGENT_CONFIG_ID:-}" ]; then
    cmd+=(--agent-config-id "$BAIJIMU_AGENT_CONFIG_ID")
  fi
  if [ -n "${BAIJIMU_AGENT_SESSION_ID:-}" ]; then
    cmd+=(--agent-session-id "$BAIJIMU_AGENT_SESSION_ID")
  fi
  if [ -n "${BAIJIMU_SESSION_ID:-}" ]; then
    cmd+=(--session-id "$BAIJIMU_SESSION_ID")
  fi

  if ! "${cmd[@]}" > "$output_file" 2> "$error_file"; then
    rm -f "$output_file"
    return 1
  fi
  credential="$(extract_llm_credential_from_json_file "$output_file")"
  rm -f "$output_file"
  [ -n "$credential" ] || return 1
  printf '%s\n' "$credential"
}

json_escape() {
  printf '%s' "$1" | sed 's/\\/\\\\/g; s/"/\\"/g'
}

toml_escape() {
  json_escape "$1"
}

backup_if_exists() {
  path="$1"
  [ -e "$path" ] || return 0
  cp -p "$path" "$path.bak.$(date +%s)"
}

remove_managed_codex_block() {
  awk '
    $0 == "# >>> baijimu managed codex router" { skipping = 1; next }
    skipping && $0 == "# <<< baijimu managed codex router" { skipping = 0; next }
    !skipping { print }
  ' "$1"
}

remove_toml_table() {
  table_name="$1"
  awk -v table_name="$table_name" '
    function trim(value) {
      sub(/^[[:space:]]+/, "", value)
      sub(/[[:space:]]+$/, "", value)
      return value
    }
    trim($0) == table_name { skipping = 1; next }
    skipping && trim($0) ~ /^\[.+\]$/ { skipping = 0 }
    !skipping { print }
  '
}

remove_root_codex_keys() {
  awk '
    /^[[:space:]]*\[.+\][[:space:]]*$/ { in_table = 1 }
    !in_table && /^[[:space:]]*(model_provider|model|sandbox_mode|approval_policy|cli_auth_credentials_store|forced_login_method)[[:space:]]*=/ { next }
    { print }
  '
}

write_codex_config() {
  api_key="$1"
  codex_dir="$HOME/.codex"
  auth_file="$codex_dir/auth.json"
  config_file="$codex_dir/config.toml"
  work_dir="$(mktemp -d "${TMPDIR:-/tmp}/baijimu-codex-config.XXXXXX")"

  mkdir -p "$codex_dir"
  chmod 700 "$codex_dir"
  backup_if_exists "$auth_file"
  backup_if_exists "$config_file"

  umask 077
  {
    printf '{\n'
    printf '  "OPENAI_API_KEY": "%s",\n' "$(json_escape "$api_key")"
    printf '  "auth_mode": "apikey"\n'
    printf '}\n'
  } > "$work_dir/auth.json"
  mv "$work_dir/auth.json" "$auth_file"
  chmod 600 "$auth_file"

  existing_config="$work_dir/existing-config.toml"
  managed_config="$work_dir/managed-config.toml"
  preserved_config="$work_dir/preserved-config.toml"
  if [ -f "$config_file" ]; then
    cp "$config_file" "$existing_config"
  else
    : > "$existing_config"
  fi
  remove_managed_codex_block "$existing_config" \
    | remove_toml_table "[model_providers.baijimu-router]" \
    | remove_root_codex_keys \
    | sed '/^[[:space:]]*$/N;/^\n$/D' > "$preserved_config"

  {
    printf '# >>> baijimu managed codex router\n'
    printf 'model_provider = "baijimu-router"\n'
    printf 'model = "%s"\n' "$(toml_escape "$CODEX_MODEL")"
    printf 'sandbox_mode = "danger-full-access"\n'
    printf 'approval_policy = "on-request"\n'
    printf 'cli_auth_credentials_store = "file"\n'
    printf 'forced_login_method = "api"\n'
    printf '\n'
    printf '[model_providers.baijimu-router]\n'
    printf 'name = "baijimu-router"\n'
    printf 'base_url = "%s"\n' "$(toml_escape "$ROUTER_BASE_URL")"
    printf 'wire_api = "responses"\n'
    printf 'requires_openai_auth = true\n'
    printf '# <<< baijimu managed codex router\n'
  } > "$managed_config"

  if [ -s "$preserved_config" ]; then
    cat "$managed_config" > "$work_dir/config.toml"
    printf '\n' >> "$work_dir/config.toml"
    cat "$preserved_config" >> "$work_dir/config.toml"
  else
    cat "$managed_config" > "$work_dir/config.toml"
  fi
  mv "$work_dir/config.toml" "$config_file"
  chmod 600 "$config_file"
  rm -rf "$work_dir"
}

set_step 1 "running" "Creating Baijimu LLM credential and writing Codex config"
if ! local_api_key="$(create_baijimu_llm_credential)"; then
  set_step 1 "failed" "Baijimu LLM credential create failed"
  error_detail="$(tail -n 5 "$state_dir/baijimu-llm-credential.err" 2>/dev/null | tr '\n' ' ')"
  finish_result false "" "" false "failed to create Baijimu LLM credential for workspace $BAIJIMU_WORKSPACE_ID: $error_detail"
  exit 1
fi
if [ -z "$local_api_key" ]; then
  set_step 1 "failed" "Baijimu LLM credential create failed"
  finish_result false "" "" false "Baijimu CLI did not return an LLM credential for workspace $BAIJIMU_WORKSPACE_ID"
  exit 1
fi
llm_credential_created=true
if ! write_codex_config "$local_api_key"; then
  unset local_api_key
  set_step 1 "failed" "Codex config write failed"
  finish_result false "" "" false "failed to write Codex config from Baijimu LLM credential"
  exit 1
fi
if [ -z "$local_api_key" ]; then
  set_step 1 "failed" "Codex auth key was not written"
  finish_result false "" "" false "$HOME/.codex/auth.json does not contain OPENAI_API_KEY"
  exit 1
fi
test "$(stat -f '%Lp' "$HOME/.codex/auth.json")" = "600"
set_step 1 "completed" "Codex config written from Baijimu LLM credential"

set_step 2 "running" "Verifying Baijimu router"
router_err="$state_dir/codex-router.err"
if ! http_code="$(
  curl -sS -m 60 -o /tmp/codex-router-responses.json -w '%{http_code}' \
    -H "Authorization: Bearer $local_api_key" \
    -H 'Content-Type: application/json' \
    -d "{\"model\":\"$CODEX_MODEL\",\"input\":\"Reply with exactly OK\"}" \
    "$ROUTER_BASE_URL/responses" 2> "$router_err"
)"; then
  error_detail="$(tail -n 5 "$router_err" 2>/dev/null | tr '\n' ' ')"
  unset local_api_key
  rm -f /tmp/codex-router-responses.json "$router_err"
  set_step 2 "failed" "Router health check failed"
  finish_result false "" "" false "router /responses health check failed: $error_detail"
  exit 1
fi
unset local_api_key
rm -f /tmp/codex-router-responses.json "$router_err"
if [ "$http_code" != "200" ]; then
  set_step 2 "failed" "Router health check failed"
  finish_result false "$http_code" "" false "router /responses health check failed: HTTP $http_code"
  exit 1
fi
set_step 2 "completed" "Baijimu router verified"

set_step 3 "running" "Checking Codex CLI"
if ! command -v codex >/dev/null 2>&1; then
  set_step 3 "failed" "Codex CLI not found in PATH"
  finish_result false "$http_code" "" false "codex CLI not found in PATH after terminal config"
  exit 1
fi
cli_version="$(codex --version 2>&1)"
if ! codex exec --skip-git-repo-check "Reply exactly OK" 2>&1 | grep -q 'OK'; then
  set_step 3 "failed" "Codex CLI smoke test failed"
  finish_result false "$http_code" "$cli_version" false "codex CLI smoke test failed"
  exit 1
fi
set_step 3 "completed" "Codex CLI verified"
set_step 4 "completed" "Codex Terminal configured"
write_install_console ""
write_install_console "Codex Terminal setup completed. You can close this window."
finish_result true "$http_code" "$cli_version" true ""
