#!/usr/bin/env bash

set -euo pipefail

BASE_URL=""
API_KEY=""
MIRROR=""
NPM_REGISTRY="auto"
SHOW_ONLY=0
TEST_ONLY=0
SKIP_MIRROR=0
SKIP_INSTALL=0
HELP=0

DEFAULT_BASE_URL="https://api.jucode.cn"
CODEX_DIR="$HOME/.codex"
CONFIG_FILE="$CODEX_DIR/config.toml"
AUTH_FILE="$CODEX_DIR/auth.json"

print_info() {
  echo "[INFO] $*"
}

print_success() {
  echo "[SUCCESS] $*"
}

print_warn() {
  echo "[WARNING] $*"
}

print_error() {
  echo "[ERROR] $*" >&2
}

print_info_err() {
  echo "[INFO] $*" >&2
}

print_warn_err() {
  echo "[WARNING] $*" >&2
}

print_success_err() {
  echo "[SUCCESS] $*" >&2
}

show_help() {
  cat <<'EOF'
Codex Ubuntu Deployment Script

Usage:
  bash setup-codex.sh [OPTIONS]

Options:
  --base-url <url>       API base URL
  --api-key <key>        API key
  --mirror <name>        apt mirror: auto|aliyun|tencent|huawei|ustc|tsinghua
  --npm-registry <name>  npm registry: auto|npmjs|npmmirror|tencent|huawei
  --skip-mirror          do not modify apt sources
  --skip-install         do not install nodejs/codex
  --test                 test API only (requires --base-url and --api-key)
  --show                 show current codex config/auth
  --help                 show this help

Examples:
  bash setup-codex.sh --base-url https://api.jucode.cn --api-key sk-xxx --mirror auto --npm-registry auto
  bash setup-codex.sh --show
EOF
}

parse_args() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --base-url)
        BASE_URL="${2:-}"
        shift 2
        ;;
      --api-key)
        API_KEY="${2:-}"
        shift 2
        ;;
      --mirror)
        MIRROR="${2:-}"
        shift 2
        ;;
      --npm-registry)
        NPM_REGISTRY="${2:-}"
        shift 2
        ;;
      --skip-mirror)
        SKIP_MIRROR=1
        shift
        ;;
      --skip-install)
        SKIP_INSTALL=1
        shift
        ;;
      --show)
        SHOW_ONLY=1
        shift
        ;;
      --test)
        TEST_ONLY=1
        shift
        ;;
      --help|-h)
        HELP=1
        shift
        ;;
      *)
        print_error "Unknown argument: $1"
        exit 1
        ;;
    esac
  done
}

require_cmd() {
  command -v "$1" >/dev/null 2>&1
}

need_sudo() {
  [[ "${EUID}" -ne 0 ]]
}

sudo_run() {
  if need_sudo; then
    sudo "$@"
  else
    "$@"
  fi
}

get_codename() {
  if [[ -r /etc/os-release ]]; then
    . /etc/os-release
    if [[ -n "${VERSION_CODENAME:-}" ]]; then
      echo "$VERSION_CODENAME"
      return
    fi
    if [[ -n "${UBUNTU_CODENAME:-}" ]]; then
      echo "$UBUNTU_CODENAME"
      return
    fi
  fi

  if require_cmd lsb_release; then
    lsb_release -cs
    return
  fi

  print_error "Cannot detect Ubuntu codename"
  exit 1
}

get_mirror_host() {
  case "$1" in
    aliyun) echo "mirrors.aliyun.com" ;;
    tencent) echo "mirrors.tencent.com" ;;
    huawei) echo "repo.huaweicloud.com" ;;
    ustc) echo "mirrors.ustc.edu.cn" ;;
    tsinghua) echo "mirrors.tuna.tsinghua.edu.cn" ;;
    *)
      print_error "Unsupported mirror: $1"
      print_error "Use one of: aliyun|tencent|huawei|ustc|tsinghua"
      exit 1
      ;;
  esac
}

pick_fastest_mirror() {
  local codename="$1"
  local mirrors=(aliyun tencent huawei ustc tsinghua)
  local best_name=""
  local best_ms=999999

  print_info_err "Probing mirror latency and selecting the fastest"

  for name in "${mirrors[@]}"; do
    local host
    host="$(get_mirror_host "$name")"
    local test_url="https://$host/ubuntu/dists/$codename/Release"

    local sample_count=2
    local total=0
    local ok=0
    local i

    for (( i=1; i<=sample_count; i++ )); do
      local sec
      sec="$(curl -o /dev/null -sS -m 6 -w '%{time_total}' "$test_url" 2>/dev/null || true)"
      if [[ -n "$sec" ]]; then
        local ms
        ms="$(awk -v v="$sec" 'BEGIN { printf "%d", v * 1000 }')"
        total=$((total + ms))
        ok=$((ok + 1))
      fi
    done

    if (( ok > 0 )); then
      local avg
      avg=$((total / ok))
      print_info_err "Mirror $name avg latency: ${avg}ms"
      if (( avg < best_ms )); then
        best_ms=$avg
        best_name="$name"
      fi
    else
      print_warn_err "Mirror $name probe failed"
    fi
  done

  if [[ -z "$best_name" ]]; then
    print_warn_err "All mirror probes failed, fallback to aliyun"
    best_name="aliyun"
  fi

  print_success_err "Auto selected mirror: $best_name"
  echo "$best_name"
}

resolve_mirror_name() {
  local requested="$1"
  local codename="$2"

  if [[ -z "$requested" ]]; then
    echo ""
    return
  fi

  if [[ "$requested" == "auto" ]]; then
    pick_fastest_mirror "$codename"
    return
  fi

  get_mirror_host "$requested" >/dev/null
  echo "$requested"
}

apply_apt_mirror() {
  local mirror_name="$1"

  if [[ -z "$mirror_name" ]]; then
    print_info "No mirror specified, skip apt source replacement"
    return
  fi

  local codename
  codename="$(get_codename)"
  local host
  host="$(get_mirror_host "$mirror_name")"

  local backup_file="/etc/apt/sources.list.backup.$(date +%Y%m%d_%H%M%S)"
  local source_file="/etc/apt/sources.list"

  if [[ ! -f "$source_file" ]]; then
    print_warn "$source_file not found, trying /etc/apt/sources.list.d/ubuntu.sources only"
  else
    sudo_run cp "$source_file" "$backup_file"
    print_info "Backed up sources.list to: $backup_file"

    cat <<EOF | sudo_run tee "$source_file" >/dev/null
deb https://$host/ubuntu/ $codename main restricted universe multiverse
deb https://$host/ubuntu/ $codename-updates main restricted universe multiverse
deb https://$host/ubuntu/ $codename-backports main restricted universe multiverse
deb https://$host/ubuntu/ $codename-security main restricted universe multiverse
EOF
    print_success "Apt source switched to $mirror_name ($host)"
  fi

  if [[ -f /etc/apt/sources.list.d/ubuntu.sources ]]; then
    local modern_backup="/etc/apt/sources.list.d/ubuntu.sources.backup.$(date +%Y%m%d_%H%M%S)"
    sudo_run cp /etc/apt/sources.list.d/ubuntu.sources "$modern_backup"
    sudo_run sed -i -E "s#https?://(archive|security)\.ubuntu\.com/ubuntu/?#https://$host/ubuntu/#g" /etc/apt/sources.list.d/ubuntu.sources
    print_success "Updated /etc/apt/sources.list.d/ubuntu.sources"
  fi

  sudo_run apt-get update -y
  print_success "apt-get update finished"
}

ensure_base_tools() {
  print_info "Installing base dependencies"
  sudo_run apt-get install -y curl ca-certificates gnupg lsb-release jq
}

ensure_nodejs() {
  if require_cmd node; then
    local version
    version="$(node -v | sed 's/^v//')"
    local major
    major="${version%%.*}"
    if [[ "$major" =~ ^[0-9]+$ ]] && (( major >= 18 )); then
      print_success "Node.js already installed: v$version"
      return
    fi
    print_warn "Node.js version too old: v$version, upgrading"
  else
    print_info "Node.js not installed, installing"
  fi

  if ! require_cmd curl; then
    ensure_base_tools
  fi

  curl -fsSL https://deb.nodesource.com/setup_22.x | sudo_run bash -
  sudo_run apt-get install -y nodejs
  print_success "Node.js installed: $(node -v)"
}

get_npm_registry_url() {
  case "$1" in
    npmjs) echo "https://registry.npmjs.org/" ;;
    npmmirror) echo "https://registry.npmmirror.com/" ;;
    tencent) echo "https://mirrors.tencent.com/npm/" ;;
    huawei) echo "https://repo.huaweicloud.com/repository/npm/" ;;
    *)
      print_error "Unsupported npm registry: $1"
      print_error "Use one of: auto|npmjs|npmmirror|tencent|huawei"
      exit 1
      ;;
  esac
}

probe_npm_registry() {
  local name="$1"
  local url
  url="$(get_npm_registry_url "$name")"
  local total=0
  local ok=0
  local i

  for (( i=1; i<=2; i++ )); do
    local sec
    sec="$(curl -o /dev/null -sS -m 6 -w '%{time_total}' "$url" 2>/dev/null || true)"
    if [[ -n "$sec" ]]; then
      local ms
      ms="$(awk -v v="$sec" 'BEGIN { printf "%d", v * 1000 }')"
      total=$((total + ms))
      ok=$((ok + 1))
    fi
  done

  if (( ok == 0 )); then
    echo ""
    return
  fi

  echo $((total / ok))
}

resolve_npm_registry() {
  local requested="$1"
  local selected=""

  if [[ -z "$requested" || "$requested" == "auto" ]]; then
    local candidates=(npmjs npmmirror tencent huawei)
    local best_name=""
    local best_ms=999999

    print_info_err "Probing npm registry latency and selecting the fastest"

    for name in "${candidates[@]}"; do
      local avg
      avg="$(probe_npm_registry "$name")"
      if [[ -n "$avg" ]]; then
        print_info_err "NPM registry $name avg latency: ${avg}ms"
        if (( avg < best_ms )); then
          best_ms=$avg
          best_name="$name"
        fi
      else
        print_warn_err "NPM registry $name probe failed"
      fi
    done

    if [[ -z "$best_name" ]]; then
      print_warn_err "All npm registry probes failed, fallback to npmmirror"
      best_name="npmmirror"
    fi

    selected="$best_name"
  else
    get_npm_registry_url "$requested" >/dev/null
    selected="$requested"
  fi

  echo "$selected"
}

set_npm_registry() {
  local requested="$1"
  local selected
  selected="$(resolve_npm_registry "$requested")"
  local registry_url
  registry_url="$(get_npm_registry_url "$selected")"

  npm config set registry "$registry_url" >/dev/null
  print_success "Using npm registry: $selected ($registry_url)"
}

ensure_codex() {
  set_npm_registry "$NPM_REGISTRY"

  if require_cmd codex; then
    print_success "Codex already installed: $(codex --version || true)"
    return
  fi

  print_info "Installing codex CLI via npm"
  npm install -g @openai/codex
  print_success "Codex installed: $(codex --version || true)"
}

backup_file_if_exists() {
  local path="$1"
  if [[ -f "$path" ]]; then
    cp "$path" "$path.backup.$(date +%Y%m%d_%H%M%S)"
    print_info "Backed up $path"
  fi
}

validate_api_key() {
  local key="$1"
  if [[ "$key" =~ ^[A-Za-z0-9._-]+$ ]]; then
    return 0
  fi

  print_error "Invalid API key format"
  return 1
}

write_codex_settings() {
  local base_url="$1"
  local api_key="$2"

  mkdir -p "$CODEX_DIR"

  backup_file_if_exists "$CONFIG_FILE"
  backup_file_if_exists "$AUTH_FILE"

  cat > "$CONFIG_FILE" <<EOF
model_provider = "OpenAI"
model = "gpt-5.4"
model_reasoning_effort = "high"
disable_response_storage = true

[model_providers.OpenAI]
name = "OpenAI"
base_url = "$base_url"
wire_api = "responses"
requires_openai_auth = true
EOF

  cat > "$AUTH_FILE" <<EOF
{
  "auth_mode": "apikey",
  "OPENAI_API_KEY": "$api_key"
}
EOF

  chmod 600 "$AUTH_FILE" || true

  export OPENAI_API_KEY="$api_key"
  export CODEX_API_KEY="$api_key"

  if require_cmd grep; then
    if ! grep -q '^export OPENAI_API_KEY=' "$HOME/.bashrc" 2>/dev/null; then
      echo "export OPENAI_API_KEY=\"$api_key\"" >> "$HOME/.bashrc"
    else
      sed -i -E "s#^export OPENAI_API_KEY=.*#export OPENAI_API_KEY=\"$api_key\"#" "$HOME/.bashrc"
    fi

    if ! grep -q '^export CODEX_API_KEY=' "$HOME/.bashrc" 2>/dev/null; then
      echo "export CODEX_API_KEY=\"$api_key\"" >> "$HOME/.bashrc"
    else
      sed -i -E "s#^export CODEX_API_KEY=.*#export CODEX_API_KEY=\"$api_key\"#" "$HOME/.bashrc"
    fi
  fi

  print_success "Written: $CONFIG_FILE"
  print_success "Written: $AUTH_FILE"
}

show_settings() {
  print_info "Config file: $CONFIG_FILE"
  if [[ -f "$CONFIG_FILE" ]]; then
    cat "$CONFIG_FILE"
  else
    print_warn "Not found"
  fi

  print_info "Auth file: $AUTH_FILE"
  if [[ -f "$AUTH_FILE" ]]; then
    local key
    key="$(jq -r '.OPENAI_API_KEY // empty' "$AUTH_FILE" 2>/dev/null || true)"
    if [[ -n "$key" ]]; then
      local masked="$key"
      if (( ${#key} > 12 )); then
        masked="${key:0:8}...${key: -4}"
      fi
      jq --arg mk "$masked" '.OPENAI_API_KEY = $mk' "$AUTH_FILE" 2>/dev/null || cat "$AUTH_FILE"
    else
      cat "$AUTH_FILE"
    fi
  else
    print_warn "Not found"
  fi
}

test_api() {
  local base_url="$1"
  local api_key="$2"
  local url="${base_url%/}/health"

  print_info "Testing API: $url"

  if curl -fsS -m 15 -H "Authorization: Bearer $api_key" "$url" >/dev/null; then
    print_success "API test success"
    return 0
  fi

  print_error "API test failed"
  return 1
}

main() {
  parse_args "$@"

  if (( HELP == 1 )); then
    show_help
    exit 0
  fi

  if (( SHOW_ONLY == 1 )); then
    show_settings
    exit 0
  fi

  if [[ -z "$BASE_URL" && -z "$API_KEY" ]]; then
    read -r -p "Enter Base URL [$DEFAULT_BASE_URL]: " input_base
    BASE_URL="${input_base:-$DEFAULT_BASE_URL}"

    while [[ -z "$API_KEY" ]]; do
      read -r -p "Enter API key: " API_KEY
      if [[ -z "$API_KEY" ]]; then
        print_warn "API key is required"
      fi
    done

    if (( SKIP_MIRROR == 0 )); then
      read -r -p "Choose mirror [auto/aliyun/tencent/huawei/ustc/tsinghua/skip] (default: auto): " input_mirror
      case "${input_mirror:-auto}" in
        skip) ;;
        auto|aliyun|tencent|huawei|ustc|tsinghua) MIRROR="${input_mirror:-auto}" ;;
        *) print_warn "Unknown mirror, skip" ;;
      esac
    fi

    if (( SKIP_INSTALL == 0 )); then
      read -r -p "Choose npm registry [auto/npmjs/npmmirror/tencent/huawei] (default: auto): " input_npm_registry
      case "${input_npm_registry:-auto}" in
        auto|npmjs|npmmirror|tencent|huawei) NPM_REGISTRY="${input_npm_registry:-auto}" ;;
        *) print_warn "Unknown npm registry, fallback to auto"; NPM_REGISTRY="auto" ;;
      esac
    fi
  fi

  if [[ -z "$BASE_URL" || -z "$API_KEY" ]]; then
    print_error "Both --base-url and --api-key are required"
    show_help
    exit 1
  fi

  BASE_URL="${BASE_URL%/}"

  if ! validate_api_key "$API_KEY"; then
    exit 1
  fi

  case "${NPM_REGISTRY:-auto}" in
    auto|npmjs|npmmirror|tencent|huawei) ;;
    *) print_error "Unsupported --npm-registry value: $NPM_REGISTRY"; exit 1 ;;
  esac

  if (( TEST_ONLY == 1 )); then
    test_api "$BASE_URL" "$API_KEY"
    exit $?
  fi

  if (( SKIP_MIRROR == 0 )) && [[ -n "$MIRROR" ]]; then
    local codename
    local mirror_to_use
    codename="$(get_codename)"
    mirror_to_use="$(resolve_mirror_name "$MIRROR" "$codename")"
    apply_apt_mirror "$mirror_to_use"
  fi

  if (( SKIP_INSTALL == 0 )); then
    ensure_base_tools
    ensure_nodejs
    ensure_codex
  fi

  write_codex_settings "$BASE_URL" "$API_KEY"

  print_success "Ubuntu Codex deployment completed"
  print_info "Run: source ~/.bashrc"
  print_info "Run: codex --version"
}

main "$@"
