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

# ---------- styling ----------
if [[ -t 1 ]] && command -v tput >/dev/null 2>&1; then
  BOLD="$(tput bold || true)"
  DIM="$(tput dim || true)"
  UNDERLINE="$(tput smul || true)"
  RESET="$(tput sgr0 || true)"
  RED="$(tput setaf 1 || true)"
  GREEN="$(tput setaf 2 || true)"
  YELLOW="$(tput setaf 3 || true)"
  BLUE="$(tput setaf 4 || true)"
  MAGENTA="$(tput setaf 5 || true)"
  CYAN="$(tput setaf 6 || true)"
else
  BOLD="" DIM="" UNDERLINE="" RESET=""
  RED="" GREEN="" YELLOW="" BLUE="" MAGENTA="" CYAN=""
fi

hr() { printf '%s\n' "${DIM}────────────────────────────────────────────────────────${RESET}"; }
title() { printf '%s%s%s\n' "${BOLD}${MAGENTA}" "$1" "${RESET}"; }
info() { printf '%s%s%s %s\n' "${CYAN}" "info" "${RESET}" "$1"; }
warn() { printf '%s%s%s %s\n' "${YELLOW}" "warn" "${RESET}" "$1"; }
ok()   { printf '%s%s%s %s\n' "${GREEN}" "ok"   "${RESET}" "$1"; }
err()  { printf '%s%s%s %s\n' "${RED}"   "err"  "${RESET}" "$1" >&2; }

# ---------- config ----------
CONFIG_DIR="${HOME}/.codex"
CONFIG_FILE="${CONFIG_DIR}/config.toml"

BLOCK=$(
  cat <<'EOF'
[mcp_servers.numerai]
url = "https://api-tournament.numer.ai/mcp/sse"

[mcp_servers.numerai.env_http_headers]
Authorization = "NUMERAI_MCP_AUTH"
EOF
)

# ---------- helpers ----------
open_url() {
  local url="$1"

  if command -v open >/dev/null 2>&1; then
    open "$url" >/dev/null 2>&1 && return 0
  fi

  if command -v xdg-open >/dev/null 2>&1; then
    xdg-open "$url" >/dev/null 2>&1 && return 0
  fi

  if command -v cmd.exe >/dev/null 2>&1; then
    cmd.exe /c start "$url" >/dev/null 2>&1 && return 0
  fi

  return 1
}

expand_tilde() {
  local p="$1"
  if [[ "${p:0:1}" == "~" ]]; then
    printf '%s' "$HOME${p:1}"
  else
    printf '%s' "$p"
  fi
}

gen_device_token() {
  if command -v openssl >/dev/null 2>&1; then
    openssl rand -hex 16
    return 0
  fi
  LC_ALL=C tr -dc 'a-f0-9' < /dev/urandom | head -c 32
}

require_cmd() {
  local c="$1"
  if ! command -v "$c" >/dev/null 2>&1; then
    err "Missing required command: $c"
    exit 1
  fi
}

extract_api_key_from_json() {
  local json="$1"
  printf '%s' "$json" | sed -nE 's/.*"apiKey"[[:space:]]*:[[:space:]]*"([^"]*)".*/\1/p' | head -n 1
}

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

remove_numerai_block() {
  local file="$1"
  local tmp="$(mktemp)"
  awk '
    /^\[mcp_servers\.numerai/ { skip=1; next }
    /^\[/ && !/^\[mcp_servers\.numerai/ { skip=0 }
    !skip { print }
  ' "$file" > "$tmp"
  cat "$tmp" > "$file"
  rm -f "$tmp"
}

clear 2>/dev/null || true
title "Numerai Codex MCP Installer"
printf "  ${DIM}Docs:${RESET} ${CYAN}${UNDERLINE}https://docs.numer.ai/numerai-tournament/mcp${RESET}\n"
hr

mkdir -p "$CONFIG_DIR"

if [[ -f "$CONFIG_FILE" ]] && grep -q '^\[mcp_servers\.numerai\]' "$CONFIG_FILE"; then
  info "A Numerai MCP entry already exists in ${BOLD}${CONFIG_FILE}${RESET}."
  printf "%sOverwrite existing configuration?%s [%sY%s/%sN%s]: " \
    "${BOLD}${MAGENTA}" "${RESET}" "${BOLD}${GREEN}" "${RESET}" "${BOLD}${RED}" "${RESET}"

  read -r overwrite_reply </dev/tty || overwrite_reply=""

  case "$overwrite_reply" in
    y|Y)
      info "Removing existing Numerai MCP configuration..."
      remove_numerai_block "$CONFIG_FILE"
      printf "\n%s\n" "$BLOCK" >>"$CONFIG_FILE"
      ok "MCP configuration overwritten."
      ;;
    *)
      info "Leaving ${BOLD}${CONFIG_FILE}${RESET} unchanged."
      ;;
  esac
else
  if [[ -f "$CONFIG_FILE" ]]; then
    info "Installing Numerai MCP configuration..."
    printf "\n%s\n" "$BLOCK" >>"$CONFIG_FILE"
  else
    info "Creating config file with Numerai MCP configuration..."
    printf "%s\n" "$BLOCK" >"$CONFIG_FILE"
  fi
  ok "MCP configuration installed."
fi

hr
printf "%sCreate MCP API key now?%s [%sY%s/%sN%s]: " \
  "${BOLD}${MAGENTA}" "${RESET}" "${BOLD}${GREEN}" "${RESET}" "${BOLD}${RED}" "${RESET}"

read -r reply </dev/tty || reply=""

case "$reply" in
  y|Y)
    require_cmd curl

    device_token="$(gen_device_token)"
    auth_url="https://numer.ai/mcp/authorize?device_token=${device_token}"
    check_url="https://api-tournament.numer.ai/mcp/check-auth?device_token=${device_token}"

    info "Device token: ${BOLD}${device_token}${RESET}"
    info "Auth URL: ${BOLD}${auth_url}${RESET}"

    if open_url "$auth_url"; then
      ok "Browser opened."
      info "Waiting for API key to be created..."
    else
      warn "Could not open a browser automatically."
      printf "%sOpen this URL in your browser:%s\n  %s\n" "${BOLD}${YELLOW}" "${RESET}" "$auth_url"
      info "Waiting for API key to be created..."
    fi

    api_key=""

    while :; do
      resp="$(curl -fsS --max-time 2 "$check_url" 2>/dev/null || true)"
      if [[ -n "$resp" ]]; then
        api_key="$(extract_api_key_from_json "$resp")"
        if [[ -n "$api_key" ]]; then
          break
        fi
      fi
      sleep 3
    done

    ok "API key detected."
    ;;
  *)
    hr
    printf "%sDo you already have an API key?%s [%sY%s/%sN%s]: " \
      "${BOLD}${MAGENTA}" "${RESET}" "${BOLD}${GREEN}" "${RESET}" "${BOLD}${RED}" "${RESET}"

    read -r have_key </dev/tty || have_key=""

    case "$have_key" in
      y|Y)
        printf "%sPaste your API key:%s " "${BOLD}${MAGENTA}" "${RESET}"
        read -r api_key </dev/tty || api_key=""

        if [[ -z "$api_key" ]]; then
          err "No API key provided. Exiting installer."
          exit 1
        fi

        ok "API key received."
        ;;
      *)
        ok "No API key provided. Exiting installer."
        exit 0
        ;;
    esac
    ;;
esac

hr
default_profile="${HOME}/.bash_profile"
printf "%sInstall NUMERAI_MCP_AUTH into which profile file?%s [%s%s%s]: " \
  "${BOLD}${MAGENTA}" "${RESET}" "${BOLD}${default_profile}${RESET}" ""

read -r profile_path </dev/tty || profile_path=""
profile_path="${profile_path:-$default_profile}"
profile_path="$(expand_tilde "$profile_path")"

profile_dir="$(dirname "$profile_path")"
mkdir -p "$profile_dir"

if [[ ! -f "$profile_path" ]]; then
  : > "$profile_path"
  ok "Created: ${BOLD}${profile_path}${RESET}"
fi

escaped_for_sq="$(escape_for_single_quotes "$api_key")"
export_line="export NUMERAI_MCP_AUTH='Token ${escaped_for_sq}'"

if grep -qE '^[[:space:]]*export[[:space:]]+NUMERAI_MCP_AUTH=' "$profile_path"; then
  tmp="$(mktemp)"
  sed -E "s|^[[:space:]]*export[[:space:]]+NUMERAI_MCP_AUTH=.*$|${export_line}|" "$profile_path" > "$tmp"
  cat "$tmp" > "$profile_path"
  rm -f "$tmp"
  ok "Updated NUMERAI_MCP_AUTH in: ${BOLD}${profile_path}${RESET}"
else
  {
    echo
    echo "# Numerai MCP (added by installer)"
    echo "$export_line"
  } >> "$profile_path"
  ok "Appended NUMERAI_MCP_AUTH to: ${BOLD}${profile_path}${RESET}"
fi

hr
ok "Done! Run this to load the API key in your current terminal:"
printf "  source %s\n" "$profile_path"
ok "Run this to start building your model:"
printf '  codex exec --yolo "find the best neural network architecture to predict target ender"\n'
exit 0