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

CODEX_DIR="${CODEX_HOME:-$HOME/.codex}"
CONFIG_FILE="$CODEX_DIR/config.toml"
MODEL="${CODEX_MODEL:-gpt-5.6-sol}"

if [[ ! "$MODEL" =~ ^[A-Za-z0-9._-]+$ ]]; then
  echo "错误：CODEX_MODEL 包含不支持的字符。" >&2
  exit 1
fi

echo "Codex GPT-5.6 200K 自动压缩配置"
echo
echo "将更新：$CONFIG_FILE"
echo "  model = \"$MODEL\""
echo "  model_context_window = 272000"
echo "  model_auto_compact_token_limit = 200000"
echo "  model_auto_compact_token_limit_scope = \"total\""
echo
echo "现有 provider、API Key 引用和其他配置会保留，并自动创建备份。"
echo "如作业需要超过 272K 的长上下文，请不要运行本脚本。"

if [[ "${CODEX_GPT56_200K_YES:-0}" != "1" ]]; then
  printf '\n继续应用？[Y/n] '
  read -r answer
  case "${answer:-Y}" in
    y|Y|yes|YES|Yes) ;;
    *) echo "已取消。"; exit 0 ;;
  esac
fi

umask 077
mkdir -p "$CODEX_DIR"
touch "$CONFIG_FILE"

timestamp="$(date +%Y%m%d%H%M%S)"
backup_base="$CONFIG_FILE.bak.$timestamp"
backup_file="$backup_base"
backup_index=1
while [[ -e "$backup_file" ]]; do
  backup_file="$backup_base.$backup_index"
  backup_index=$((backup_index + 1))
done
cp "$CONFIG_FILE" "$CONFIG_FILE.bak"
cp "$CONFIG_FILE" "$backup_file"

has_tomllib=0
if command -v python3 >/dev/null 2>&1 && python3 -c 'import tomllib' >/dev/null 2>&1; then
  has_tomllib=1
fi

if [[ "$has_tomllib" == "1" && -s "$CONFIG_FILE" ]]; then
  if ! python3 - "$CONFIG_FILE" <<'PY'
import sys
import tomllib

with open(sys.argv[1], "rb") as config_file:
    tomllib.load(config_file)
PY
  then
    echo "错误：现有 config.toml 不是有效 TOML，未做修改。" >&2
    exit 1
  fi
fi

tmp_file="$CODEX_DIR/.config.toml.gpt56-200k.$$"
trap 'rm -f "$tmp_file"' EXIT

awk -v model="$MODEL" '
BEGIN {
  in_root = 1
}

function emit_missing(added) {
  added = 0
  if (!seen_model) {
    print "model = \"" model "\""
    seen_model = 1
    added = 1
  }
  if (!seen_context) {
    print "model_context_window = 272000"
    seen_context = 1
    added = 1
  }
  if (!seen_compact) {
    print "model_auto_compact_token_limit = 200000"
    seen_compact = 1
    added = 1
  }
  if (!seen_scope) {
    print "model_auto_compact_token_limit_scope = \"total\""
    seen_scope = 1
    added = 1
  }
  if (added) print ""
}

{
  if (in_root && $0 ~ /^[[:space:]]*\[/) {
    emit_missing()
    in_root = 0
  }

  if (in_root && $0 ~ /^[[:space:]]*model[[:space:]]*=/) {
    if (!seen_model) print "model = \"" model "\""
    seen_model = 1
    next
  }
  if (in_root && $0 ~ /^[[:space:]]*model_context_window[[:space:]]*=/) {
    if (!seen_context) print "model_context_window = 272000"
    seen_context = 1
    next
  }
  if (in_root && $0 ~ /^[[:space:]]*model_auto_compact_token_limit[[:space:]]*=/) {
    if (!seen_compact) print "model_auto_compact_token_limit = 200000"
    seen_compact = 1
    next
  }
  if (in_root && $0 ~ /^[[:space:]]*model_auto_compact_token_limit_scope[[:space:]]*=/) {
    if (!seen_scope) print "model_auto_compact_token_limit_scope = \"total\""
    seen_scope = 1
    next
  }

  print
}

END {
  if (in_root) emit_missing()
}
' "$CONFIG_FILE" > "$tmp_file"

if [[ "$has_tomllib" == "1" ]]; then
  if ! python3 - "$tmp_file" <<'PY'
import sys
import tomllib

with open(sys.argv[1], "rb") as config_file:
    tomllib.load(config_file)
PY
  then
    cp "$backup_file" "$CONFIG_FILE"
    echo "错误：生成的 TOML 校验失败，已自动恢复原配置。" >&2
    exit 1
  fi
fi

mv "$tmp_file" "$CONFIG_FILE"
chmod 600 "$CONFIG_FILE"
trap - EXIT

echo
echo "配置完成：$CONFIG_FILE"
echo "备份文件：$backup_file"
echo "请完全退出并重新打开 Codex，使新配置生效。"