#!/usr/bin/env bash

set -euo pipefail

OPENAI_BASE_URL_VALUE='https://api.omnirouter.ru/v1'

usage() {
  cat <<'EOF'
OmniRouter Codex installer.

Usage:
  bash -s -- [--launch] [omni_api_key] [-- codex_args...]

Behavior:
  - Prompts for Omni API key if not provided
  - Writes OMNIROUTER_API_KEY to a user env file
  - Leaves any real OpenAI API key untouched
  - Removes legacy omni_ keys from ~/.codex/auth.json when present
  - Updates ~/.codex/config.toml and ~/.codex/omnirouter.config.toml to use OmniRouter
  - Adds source lines to common shell rc files
  - Optionally launches `codex`

Examples:
  curl -fsSL "https://api.omnirouter.ru/install/codex.sh" | bash
  curl -fsSL "https://api.omnirouter.ru/install/codex.sh" | bash -s -- omni_xxx
  curl -fsSL "https://api.omnirouter.ru/install/codex.sh" | bash -s -- --launch -- --yolo -m gpt-5.4

Security:
  - Prefer interactive prompt mode. Passing the key as an argument may leak into shell history.
EOF
}

ensure_source_line() {
  local rc_file source_line
  source_line='[ -f "${XDG_CONFIG_HOME:-$HOME/.config}/omnirouter/codex-env.sh" ] && . "${XDG_CONFIG_HOME:-$HOME/.config}/omnirouter/codex-env.sh"'

  for rc_file in "$HOME/.zshrc" "$HOME/.bashrc" "$HOME/.bash_profile" "$HOME/.profile"; do
    touch "$rc_file"
    if ! grep -Fqx "$source_line" "$rc_file"; then
      printf '\n%s\n' "$source_line" >> "$rc_file"
    fi
  done
}

write_env_file() {
  local env_dir env_file
  env_dir="${XDG_CONFIG_HOME:-$HOME/.config}/omnirouter"
  env_file="$env_dir/codex-env.sh"
  mkdir -p "$env_dir"

  {
    printf 'export OMNIROUTER_API_KEY=%q\n' "$omni_api_key"
  } > "$env_file"

  chmod 600 "$env_file"
  ensure_source_line
}

cleanup_auth_file() {
  local codex_dir auth_file
  codex_dir="${CODEX_HOME:-$HOME/.codex}"
  auth_file="$codex_dir/auth.json"
  mkdir -p "$codex_dir"

  OMNI_API_KEY="$omni_api_key" CODEX_AUTH_PATH="$auth_file" python3 <<'PY'
import json
import os
import pathlib
import sys

auth_path = pathlib.Path(os.environ["CODEX_AUTH_PATH"])

if not auth_path.exists():
    sys.exit(0)

try:
    data = json.loads(auth_path.read_text())
except json.JSONDecodeError as exc:
    print(f"Failed to parse Codex auth file {auth_path}: {exc}", file=sys.stderr)
    sys.exit(1)

openai_key = data.get("OPENAI_API_KEY")
if isinstance(openai_key, str) and openai_key.startswith("omni_"):
    data.pop("OPENAI_API_KEY", None)
    if data.get("auth_mode") == "apikey":
        data.pop("auth_mode", None)
    auth_path.write_text(json.dumps(data, indent=2) + "\n")
PY

  if [[ -f "$auth_file" ]]; then
    chmod 600 "$auth_file"
  fi
}

write_config_file() {
  local codex_dir config_file profile_file
  codex_dir="${CODEX_HOME:-$HOME/.codex}"
  config_file="$codex_dir/config.toml"
  profile_file="$codex_dir/omnirouter.config.toml"
  mkdir -p "$codex_dir"
  touch "$config_file"

  OPENAI_BASE_URL_VALUE="$OPENAI_BASE_URL_VALUE" CODEX_CONFIG_PATH="$config_file" CODEX_PROFILE_CONFIG_PATH="$profile_file" python3 <<'PY'
import os
import pathlib
import re

config_path = pathlib.Path(os.environ["CODEX_CONFIG_PATH"])
profile_config_path = pathlib.Path(os.environ["CODEX_PROFILE_CONFIG_PATH"])
base_url = os.environ["OPENAI_BASE_URL_VALUE"]
start_marker = "# >>> omnirouter codex >>>"
end_marker = "# <<< omnirouter codex <<<"
managed_block = f"""{start_marker}
[model_providers.omnirouter]
name = "OmniRouter"
base_url = "{base_url}"
wire_api = "responses"
env_key = "OMNIROUTER_API_KEY"
{end_marker}
"""
profile_block = 'model_provider = "omnirouter"\n'

text = config_path.read_text() if config_path.exists() else ""

managed_pattern = re.compile(
    rf"(?ms)^\s*{re.escape(start_marker)}\n.*?^\s*{re.escape(end_marker)}\n?"
)
text = managed_pattern.sub("", text).strip()

def strip_table(text: str, table_name: str) -> str:
    pattern = re.compile(
        rf"(?ms)^[ \t]*\[{re.escape(table_name)}\][ \t]*\n.*?(?=^[ \t]*\[|\Z)"
    )
    return pattern.sub("", text)

text = re.sub(r'(?m)^[ \t]*profile[ \t]*=[ \t]*"omnirouter"[ \t]*\n?', "", text)
text = strip_table(text, "model_providers.omnirouter").strip()
text = strip_table(text, "profiles.omnirouter").strip()

if text:
    text += "\n"
text = text.rstrip() + "\n\n" + managed_block

config_path.write_text(text.rstrip() + "\n")

profile_text = profile_config_path.read_text() if profile_config_path.exists() else ""
profile_text = re.sub(r'(?m)^[ \t]*model_provider[ \t]*=.*\n?', "", profile_text).strip()
if profile_text:
    profile_text = profile_block + profile_text.rstrip() + "\n"
else:
    profile_text = profile_block
profile_config_path.write_text(profile_text)
PY
}

prompt_for_key() {
  if [[ -t 0 ]]; then
    read -r -s -p "Enter Omni API key: " omni_api_key
    printf '\n'
    return
  fi

  if [[ -r /dev/tty ]]; then
    read -r -s -p "Enter Omni API key: " omni_api_key < /dev/tty
    printf '\n' > /dev/tty
    return
  fi

  printf 'Omni API key is required.\n' >&2
  printf 'Pass it directly, for example:\n' >&2
  printf '  curl -fsSL "%s/install/codex.sh" | bash -s -- omni_xxx\n' "https://api.omnirouter.ru" >&2
  exit 1
}

launch=0
key_arg=''
declare -a codex_args=()

while (($# > 0)); do
  case "$1" in
    --help|-h)
      usage
      exit 0
      ;;
    --launch)
      launch=1
      shift
      ;;
    --)
      shift
      codex_args=("$@")
      break
      ;;
    -*)
      printf 'Unknown option: %s\n\n' "$1" >&2
      usage >&2
      exit 1
      ;;
    *)
      if [[ -n "$key_arg" ]]; then
        printf 'Unexpected extra argument: %s\n\n' "$1" >&2
        usage >&2
        exit 1
      fi
      key_arg="$1"
      shift
      ;;
  esac
done

if [[ -n "$key_arg" ]]; then
  omni_api_key="$key_arg"
else
  prompt_for_key
fi

if [[ -z "$omni_api_key" ]]; then
  printf 'Omni API key is required.\n' >&2
  exit 1
fi

write_env_file
cleanup_auth_file
write_config_file

export OMNIROUTER_API_KEY="$omni_api_key"

printf 'Codex configured for OmniRouter.\n'
printf 'Saved env vars to %s\n' "${XDG_CONFIG_HOME:-$HOME/.config}/omnirouter/codex-env.sh"
printf 'Updated Codex config: %s\n' "${CODEX_HOME:-$HOME/.codex}/config.toml"
printf 'Updated Codex profile: %s\n' "${CODEX_HOME:-$HOME/.codex}/omnirouter.config.toml"
printf 'Checked Codex auth cache: %s\n' "${CODEX_HOME:-$HOME/.codex}/auth.json"
printf 'Updated shell startup files: ~/.zshrc, ~/.bashrc, ~/.bash_profile, ~/.profile\n'

if ((launch)); then
  exec codex -c 'model_provider="omnirouter"' "${codex_args[@]}"
fi

printf 'Run Codex with: codex -c '\''model_provider="omnirouter"'\''\n'
printf 'Alternative: codex --profile omnirouter\n'
printf 'Open a new shell or run: source ~/.zshrc\n'