#!/usr/bin/env python3
"""
setup_codex.py — Codex CLI 代理配置生成器
==========================================

在 Windows 本机运行，生成两个文件到桌面：

  1. connect_XXX.bat  — 每天双击，SSH 连服务器并转发代理端口
  2. setup_proxy.sh   — 一次性脚本，用 MobaXterm 上传到服务器运行

使用流程：
  1. python setup_codex.py          ← 本机运行，生成两个文件
  2. MobaXterm 上传 setup_proxy.sh  ← 拖到服务器家目录
  3. bash setup_proxy.sh            ← 服务器上运行一次
  4. 以后每次：双击 bat → 服务器上 proxyon → codex

前置条件：
  - 本机已运行网络代理程序，能正常访问外网
  - 有目标服务器的 SSH 登录权限
"""

import os
import sys
import socket
import hashlib
from pathlib import Path

# ================================================================
# 配置
# ================================================================

KNOWN_PROXY_PORTS = [
    (7890,  "Proxy-7890"),
    (7897,  "Proxy-7897"),
    (10808, "Proxy-10808"),
    (10809, "Proxy-10809"),
    (1080,  "Proxy-1080"),
    (8080,  "Proxy-8080"),
]

TEST_URL = "http://www.google.com"

LAB_SERVERS = {
    "1": ("10.177.16.130", "2080Ti"),
    "2": ("10.177.19.58",  "5070Ti-001"),
    "3": ("10.177.16.232", "5070Ti-002"),
    "4": ("10.177.11.62",  "3090"),
    "5": ("10.177.16.81",  "4070Ti-001"),
    "6": ("10.177.18.132", "4070Ti-002"),
}


# ================================================================
# 第一步：检测本地代理
# ================================================================

def is_port_open(port, host="127.0.0.1", timeout=1):
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return True
    except (ConnectionRefusedError, TimeoutError, OSError):
        return False


def validate_proxy(port):
    """通过代理 HEAD google.com 验证外网连通性"""
    import urllib.request
    import urllib.error

    proxy = urllib.request.ProxyHandler({"http": f"http://127.0.0.1:{port}"})
    opener = urllib.request.build_opener(proxy)
    try:
        req = urllib.request.Request(TEST_URL, method="HEAD")
        opener.open(req, timeout=10)
        return True
    except urllib.error.HTTPError as e:
        return e.code < 500
    except Exception:
        return False


def detect_proxy():
    """扫描端口 → 验证连通 → 返回 (port, name)"""
    print("[1/3] 检测本地代理...")

    open_ports = [(p, n) for p, n in KNOWN_PROXY_PORTS if is_port_open(p)]

    if not open_ports:
        print(f"  ✗ 未检测到代理端口。请确认代理程序已启动。")
        manual = input("  手动输入端口 [回车退出]: ").strip()
        if manual and manual.isdigit() and is_port_open(int(manual)):
            open_ports = [(int(manual), f"Proxy-{manual}")]
        else:
            sys.exit(1)

    for port, name in open_ports:
        print(f"  验证 {name} → {TEST_URL} ...", end=" ", flush=True)
        if validate_proxy(port):
            print("✓")
            return port, name
        print("✗")

    print("  ✗ 无法通过代理访问外网。")
    sys.exit(1)


# ================================================================
# 第二步：收集信息
# ================================================================

def auto_port(user):
    """用户名 → 17001-17999 的固定端口"""
    return 17001 + int(hashlib.md5(user.encode()).hexdigest()[:4], 16) % 999


def collect_info():
    print("\n[2/3] 填写信息")

    # 服务器
    print("  服务器列表:")
    for k, (ip, name) in LAB_SERVERS.items():
        print(f"    {k}. {name} ({ip})")
    choice = input("  选择编号或输入 IP: ").strip()
    if choice in LAB_SERVERS:
        host, server_name = LAB_SERVERS[choice]
        print(f"  → {server_name} ({host})")
    elif choice:
        host, server_name = choice, choice
    else:
        print("  ✗ 必须指定服务器"); sys.exit(1)

    # 用户名
    default_user = os.getenv("USERNAME", "")
    user = input(f"  用户名 [{default_user}]: ").strip() or default_user
    if not user:
        print("  ✗ 必须输入用户名"); sys.exit(1)

    # 密钥
    key_path = input("  SSH 密钥路径 [没有直接回车]: ").strip()

    # 远程端口：用户名哈希自动分配，同一用户名始终同一端口
    remote_port = auto_port(user)
    print(f"  远程端口: {remote_port} (根据用户名自动分配)")

    return {
        "host": host,
        "server_name": server_name,
        "user": user,
        "key_path": key_path,
        "remote_port": remote_port,
    }


# ================================================================
# 第三步：生成文件
# ================================================================

def generate_bat(info, local_port, proxy_name, desktop):
    """生成 Windows bat 连接脚本"""
    key_arg = f'-i "{info["key_path"]}" ' if info["key_path"] else ""
    sn = info["server_name"]

    content = f"""@echo off
chcp 65001 >nul
echo.
echo  [proxy-tunnel] {sn} ({info["host"]})
echo  [proxy-tunnel] 本机:{local_port} ---^> 服务器:127.0.0.1:{info["remote_port"]}
echo  [proxy-tunnel] 连接中... (保持此窗口打开)
echo.
ssh -N -o StrictHostKeyChecking=accept-new {key_arg}-R {info["remote_port"]}:127.0.0.1:{local_port} {info["user"]}@{info["host"]}
echo.
echo  连接已断开。
pause
"""
    path = desktop / f"connect_{sn}.bat"
    path.write_text(content, encoding="utf-8")
    return path


def generate_sh(info):
    """生成服务器端一次性配置脚本"""
    rp = info["remote_port"]

    content = f"""#!/bin/bash
# ================================================================
# setup_proxy.sh — 代理开关一键配置
# ================================================================
# 用法: bash setup_proxy.sh
#
# 此脚本在服务器上运行一次，之后就可以用 proxyon/proxyoff 了。
# 配合本机的 bat 文件使用（bat 负责 SSH 隧道转发）。
# ================================================================

set -e

PORT={rp}
PROXY_DIR="$HOME/.proxy.d"
SCRIPT_FILE="$PROXY_DIR/proxy-onoff.sh"

echo "[setup] 远程端口: $PORT"
echo "[setup] 配置目录: $PROXY_DIR"

# 1. 创建代理开关脚本
mkdir -p "$PROXY_DIR"

cat > "$SCRIPT_FILE" << 'PROXYEOF'
# --------------------------------------------------
# 代理开关
# 用法: proxyon  开启 | proxyoff 关闭 | proxystat 状态
# --------------------------------------------------

_PROXY_PORT={rp}
_PROXY_URL="http://127.0.0.1:$_PROXY_PORT"

proxyon() {{
    if command -v ss >/dev/null 2>&1; then
        if ! ss -tln 2>/dev/null | grep -q ":${{_PROXY_PORT}}\\b"; then
            echo "[proxy] ✗ 端口 $_PROXY_PORT 未监听 — 本机 bat 隧道可能没开"
            return 1
        fi
    fi
    export http_proxy="$_PROXY_URL"
    export https_proxy="$_PROXY_URL"
    export HTTP_PROXY="$_PROXY_URL"
    export HTTPS_PROXY="$_PROXY_URL"
    export NO_PROXY="localhost,127.0.0.1,::1,10.177.0.0/16"
    echo "[proxy] ✓ ON → $_PROXY_URL"
}}

proxyoff() {{
    unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY NO_PROXY
    echo "[proxy] OFF"
}}

proxystat() {{
    if [ -n "$http_proxy" ]; then
        echo "[proxy] ON → $http_proxy"
    else
        echo "[proxy] OFF"
    fi
    if command -v ss >/dev/null 2>&1; then
        if ss -tln | grep -q ":${{_PROXY_PORT}}\\b"; then
            echo "[tunnel] ✓ 端口 $_PROXY_PORT 监听中"
        else
            echo "[tunnel] ✗ 端口 $_PROXY_PORT 未监听"
        fi
    fi
}}
PROXYEOF

chmod +x "$SCRIPT_FILE"
echo "[setup] ✓ 代理开关脚本已写入"

# 2. 挂载到 .bashrc（幂等）
if grep -q 'proxy-onoff.sh' "$HOME/.bashrc" 2>/dev/null; then
    echo "[setup] .bashrc 已有挂载，跳过"
else
    printf '\\n# 代理开关 (proxyon / proxyoff / proxystat)\\nsource %s\\n' "$SCRIPT_FILE" >> "$HOME/.bashrc"
    echo "[setup] ✓ 已添加到 .bashrc"
fi

# 3. 检查 Node.js 版本
NODE_OK=0
if command -v node >/dev/null 2>&1; then
    NODE_VER=$(node --version 2>/dev/null | sed 's/^v//')
    NODE_MAJOR=$(echo "$NODE_VER" | cut -d. -f1)
    if [ "$NODE_MAJOR" -ge 20 ] 2>/dev/null; then
        echo "[setup] ✓ Node.js $NODE_VER (满足 >= 20)"
        NODE_OK=1
    else
        echo "[setup] ⚠ Node.js $NODE_VER 版本过低 (需要 >= 20)"
    fi
else
    echo "[setup] ⚠ 未检测到 Node.js"
fi

if [ "$NODE_OK" -eq 0 ]; then
    echo ""
    echo "  请先安装 Node.js (不需要 root):"
    echo ""
    echo "    proxyon"
    echo "    curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.3/install.sh | bash"
    echo "    source ~/.bashrc"
    echo "    nvm install --lts"
    echo ""
    echo "  安装后重新运行本脚本，或直接执行:"
    echo "    npm install -g @openai/codex && codex --login"
    echo ""
fi

echo ""
echo "========================================="
echo "  代理配置完成！"
echo "========================================="
echo ""
echo "  使用方法:"
echo "    1. 本机双击 bat 建立隧道（保持窗口开着）"
echo "    2. 新开终端连服务器，输入 proxyon"
echo "    3. 验证: curl -I http://www.google.com"
if [ "$NODE_OK" -eq 1 ]; then
echo "    4. 安装 Codex CLI:"
echo "         npm install -g @openai/codex"
echo "         codex"
fi
echo ""
"""
    return content


def main():
    print()
    print("=" * 48)
    print("   Codex CLI 代理配置生成器")
    print("=" * 48)
    print()

    # 1. 检测代理
    local_port, proxy_name = detect_proxy()

    # 2. 收集信息
    info = collect_info()

    # 3. 生成文件
    print(f"\n[3/3] 生成文件...")

    desktop = Path(os.environ.get("USERPROFILE", Path.home())) / "Desktop"
    if not desktop.exists():
        desktop = Path.home() / "Desktop"
    desktop.mkdir(exist_ok=True)

    bat_path = generate_bat(info, local_port, proxy_name, desktop)
    print(f"  ✓ {bat_path.name}")

    sh_content = generate_sh(info)
    sh_path = desktop / "setup_proxy.sh"
    sh_path.write_text(sh_content, encoding="utf-8", newline="\n")  # Unix 换行
    print(f"  ✓ {sh_path.name}")

    # 完成
    sn = info["server_name"]
    print(f"""
{"=" * 48}
  桌面已生成两个文件:
{"=" * 48}

  {bat_path.name}
    每天双击，建立代理隧道（窗口保持打开）

  {sh_path.name}
    用 MobaXterm 拖到服务器家目录，运行一次:
      bash setup_proxy.sh

  之后每次使用:
    本机: 双击 {bat_path.name}
    服务器: proxyon → codex / pip / git ...
""")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n已取消。")
        sys.exit(0)