#!/usr/bin/env bash
# DeepSeek Key Tool (Linux) — key DeepSeek → Claude Code CLI + extension VS Code / Cursor / Codium
# Cách dùng:
#   chmod +x DeepSeekKeyTool.sh
#   ./DeepSeekKeyTool.sh                 # menu tương tác
#   ./DeepSeekKeyTool.sh sk-xxxx         # áp dụng key luôn
#   ./DeepSeekKeyTool.sh --reset         # gỡ cấu hình
set -euo pipefail
export LANG="${LANG:-C.UTF-8}"
export LC_ALL="${LC_ALL:-C.UTF-8}"

VERSION="1.0.0"
GATEWAY="https://api.nghimmo.com"
MODEL="nghi/deepseek-v4.1-flash"
SETTINGS="${HOME}/.claude/settings.json"
CLAUDE_JSON="${HOME}/.claude.json"
LOG="${HOME}/.local/share/deepseek-key-tool/deepseek-key-tool.log"
ENV_KEYS="ANTHROPIC_AUTH_TOKEN ANTHROPIC_API_KEY ANTHROPIC_BASE_URL ANTHROPIC_MODEL ANTHROPIC_DEFAULT_OPUS_MODEL ANTHROPIC_DEFAULT_SONNET_MODEL ANTHROPIC_DEFAULT_HAIKU_MODEL ANTHROPIC_SMALL_FAST_MODEL CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY CLAUDE_CODE_USE_GATEWAY CLAUDE_CODE_DISABLE_REFUSAL_FALLBACK API_TIMEOUT_MS"

log() {
  mkdir -p "$(dirname "$LOG")"
  printf '%s %s\n' "$(date '+%Y-%m-%d %H:%M:%S')" "$1" >>"$LOG"
}

say() {
  echo "[DeepSeek Key Tool] $*" >&2
  log "$*"
}

# Hoi /v1/models bang key. 0 = key DeepSeek (hoac khong kiem tra duoc), 3 = khong phai DeepSeek, 4 = key sai/het han.
check_key() {
  command -v python3 >/dev/null 2>&1 || return 0
  local rc
  set +e
  GATEWAY="$GATEWAY" KEY="$1" python3 - <<'PY'
import json, os, sys, urllib.request, urllib.error
req = urllib.request.Request(
    os.environ["GATEWAY"].rstrip("/") + "/v1/models",
    headers={"x-api-key": os.environ["KEY"], "user-agent": "DeepSeekKeyTool/1.0"},
)
try:
    with urllib.request.urlopen(req, timeout=15) as r:
        data = json.loads(r.read().decode())
except urllib.error.HTTPError as e:
    try:
        body = json.loads(e.read().decode())
    except Exception:
        body = None
    sys.exit(4 if e.code in (401, 403) and isinstance(body, dict) and "error" in body else 0)
except Exception:
    sys.exit(0)
ids = [str(m.get("id", "")) for m in (data.get("data") or []) if isinstance(m, dict)]
sys.exit(0 if any("deepseek" in i for i in ids) else 3)
PY
  rc=$?
  set -e
  return $rc
}

clear_extension_cache() {
  local rel="User/globalStorage/anthropic.claude-code"
  local base
  for base in \
    "${HOME}/.config/Code" \
    "${HOME}/.config/Code - Insiders" \
    "${HOME}/.config/Cursor" \
    "${HOME}/.config/VSCodium" \
    "${HOME}/.config/Antigravity" \
    "${HOME}/.config/Windsurf" \
    "${HOME}/.config/Trae"
  do
    if [ -d "${base}/${rel}" ]; then
      rm -rf "${base}/${rel}"
      log "cleared extension cache: ${base}/${rel}"
    fi
  done
}

write_claude_code() {
  local key="$1"
  mkdir -p "$(dirname "$SETTINGS")"
  clear_extension_cache
  if command -v python3 >/dev/null 2>&1; then
    SETTINGS="$SETTINGS" CLAUDE_JSON="$CLAUDE_JSON" GATEWAY="$GATEWAY" KEY="$key" MODEL="$MODEL" python3 <<'PY'
import json, os
from pathlib import Path
sf = Path(os.environ["SETTINGS"])
data = {}
if sf.exists():
    try:
        data = json.loads(sf.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        data = {}
if not isinstance(data, dict):
    data = {}
env = data.get("env")
if not isinstance(env, dict):
    env = {}
model = os.environ["MODEL"]
env["ANTHROPIC_BASE_URL"] = os.environ["GATEWAY"]
env["ANTHROPIC_AUTH_TOKEN"] = os.environ["KEY"]
for k in ("ANTHROPIC_MODEL", "ANTHROPIC_DEFAULT_OPUS_MODEL", "ANTHROPIC_DEFAULT_SONNET_MODEL",
          "ANTHROPIC_DEFAULT_HAIKU_MODEL", "ANTHROPIC_SMALL_FAST_MODEL"):
    env[k] = model
env["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1"
env["CLAUDE_CODE_USE_GATEWAY"] = "1"
env["CLAUDE_CODE_DISABLE_REFUSAL_FALLBACK"] = "1"
env["API_TIMEOUT_MS"] = "600000"
env.pop("ANTHROPIC_API_KEY", None)
data["env"] = env
sf.write_text(json.dumps(data, indent=2) + "\n", encoding="utf-8")
cj = Path(os.environ["CLAUDE_JSON"])
if cj.exists():
    try:
        cdata = json.loads(cj.read_text(encoding="utf-8"))
        if isinstance(cdata, dict):
            for k in ("primaryApiKey", "customApiKeyResponses"):
                cdata.pop(k, None)
            cj.write_text(json.dumps(cdata, indent=2) + "\n", encoding="utf-8")
    except json.JSONDecodeError:
        pass
PY
  else
    cat >"$SETTINGS" <<EOF
{
  "env": {
    "ANTHROPIC_BASE_URL": "$GATEWAY",
    "ANTHROPIC_AUTH_TOKEN": "$key",
    "ANTHROPIC_MODEL": "$MODEL",
    "ANTHROPIC_DEFAULT_OPUS_MODEL": "$MODEL",
    "ANTHROPIC_DEFAULT_SONNET_MODEL": "$MODEL",
    "ANTHROPIC_DEFAULT_HAIKU_MODEL": "$MODEL",
    "ANTHROPIC_SMALL_FAST_MODEL": "$MODEL",
    "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY": "1",
    "CLAUDE_CODE_USE_GATEWAY": "1",
    "CLAUDE_CODE_DISABLE_REFUSAL_FALLBACK": "1",
    "API_TIMEOUT_MS": "600000"
  }
}
EOF
  fi
}

reset_all() {
  clear_extension_cache
  if command -v python3 >/dev/null 2>&1 && [ -f "$SETTINGS" ]; then
    SETTINGS="$SETTINGS" CLAUDE_JSON="$CLAUDE_JSON" ENV_KEYS="$ENV_KEYS" python3 <<'PY'
import json, os
from pathlib import Path
for path_key in ("SETTINGS", "CLAUDE_JSON"):
    p = Path(os.environ[path_key])
    if not p.exists():
        continue
    try:
        data = json.loads(p.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        continue
    if path_key == "SETTINGS" and isinstance(data.get("env"), dict):
        for k in os.environ["ENV_KEYS"].split():
            data["env"].pop(k, None)
    if path_key == "CLAUDE_JSON" and isinstance(data, dict):
        for k in ("primaryApiKey", "customApiKeyResponses"):
            data.pop(k, None)
    p.write_text(json.dumps(data, indent=2) + "\n", encoding="utf-8")
PY
  elif [ -f "$SETTINGS" ]; then
    rm -f "$SETTINGS"
  fi
}

ask_key() {
  local preset="${1:-}"
  if [ -n "$preset" ]; then
    printf '%s' "$preset"
    return
  fi
  local out=""
  if [ -t 0 ] && [ -r /dev/tty ]; then
    read -r -s -p "API key DeepSeek (sk-...): " out </dev/tty
    echo >&2
  fi
  printf '%s' "$out"
}

do_apply() {
  local key="$1" rc=0
  key="$(printf '%s' "$key" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')"
  if [ -z "$key" ] || [ "${#key}" -lt 8 ] || [[ "$key" == *YOUR_API_KEY* ]]; then
    say "Key không hợp lệ."
    exit 1
  fi
  say "Đang kiểm tra key..."
  check_key "$key" || rc=$?
  if [ "$rc" = "3" ]; then
    say "Key này không phải key DeepSeek. Key Claude dùng ClaudeKeyTool.sh."
    exit 1
  fi
  if [ "$rc" = "4" ]; then
    say "Key sai, bị khoá hoặc đã hết hạn. Kiểm tra tại ${GATEWAY}/check"
    exit 1
  fi
  write_claude_code "$key"
  say "Đã ghi ~/.claude/settings.json (model ${MODEL})"
  say "Claude Code CLI: chạy lại 'claude'."
  say "VS Code / Cursor: Ctrl+Shift+P → Reload Window."
  say "Log: $LOG"
}

do_reset() {
  reset_all
  say "Đã reset cấu hình Claude Code trên máy này."
}

# --- main ---
ARG1="${1:-}"
case "$ARG1" in
  --reset|-r)
    do_reset
    exit 0
    ;;
  sk-*)
    do_apply "$ARG1"
    exit 0
    ;;
esac

echo ""
echo "  DeepSeek Key Tool (Linux) v${VERSION}"
echo "  Gateway: ${GATEWAY} · Model: ${MODEL}"
echo ""
echo "  1) Áp dụng key DeepSeek"
echo "  2) Reset cấu hình"
echo "  3) Thoát"
read -r -p "Chọn [1]: " choice </dev/tty 2>/dev/null || choice="1"
case "${choice:-1}" in
  2)
    do_reset
    ;;
  3)
    exit 0
    ;;
  *)
    key="$(ask_key "")"
    do_apply "$key"
    ;;
esac
