#!/usr/bin/env bash
# ============================================================================
#  Grok 注册机 - 一键启动脚本 (Linux / macOS)
#  自动检测并安装缺失环境，然后启动 register_cli.py
#
#  用法:
#    ./start.sh                          # 默认不限数量（持续注册）
#    ./start.sh --count 10               # 跑 10 个
#    ./start.sh --count 10 --threads 3   # 3 线程跑 10 个
#    ./start.sh --extra 5                # 在已有基础上再注册 5 个
#
#  首次运行会自动:
#    1. 安装 uv（如缺失）
#    2. 安装 Python 3.13（如缺失，经 uv 管理）
#    3. 同步 Python 依赖
#    4. 尝试安装 Chrome/Chromium（如缺失且有权限）
#    5. 从模板复制 .env / config.json（如缺失）
#    6. 无 DISPLAY 时自动装/启 Xvfb（虚拟显示，过 CF 比纯 headless 稳）
#
#  环境变量:
#    SKIP_XVFB=1          # 强制不用 Xvfb（走 Python 侧自动 headless）
#    FORCE_XVFB=1         # 即使已有 DISPLAY 也用 xvfb-run 包一层
#    XVFB_DISPLAY=:99     # 常驻 Xvfb 使用的 display（默认 :99）
# ============================================================================
set -euo pipefail

# ── 颜色 ──
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m'

PROJECT_DIR="$(cd "$(dirname "$0")" && pwd)"
cd "$PROJECT_DIR"

# uv 常见安装路径，确保后续命令能找到
export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH"

# 启动包装：默认直接 exec；无桌面时改为 xvfb-run / 自启 Xvfb
RUN_PREFIX=()
USE_XVFB=0
XVFB_MODE=""   # xvfb-run | persistent | none

echo -e "${CYAN}${BOLD}============================================${NC}"
echo -e "${CYAN}${BOLD}   Grok 注册机 - 一键启动 (Linux/macOS)${NC}"
echo -e "${CYAN}${BOLD}============================================${NC}"
echo ""

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

# 是否已有可用图形显示（含用户已 export 的 DISPLAY / Wayland）
has_display() {
    [ -n "${DISPLAY:-}" ] || [ -n "${WAYLAND_DISPLAY:-}" ]
}

# 尝试用包管理器安装 xvfb（仅免密 sudo 时自动装）
try_install_xvfb() {
    if have_cmd xvfb-run || have_cmd Xvfb; then
        return 0
    fi

    echo -e "  ${YELLOW}⚠${NC} 未检测到 Xvfb，尝试自动安装..."

    if have_cmd apt-get && have_cmd sudo && sudo -n true 2>/dev/null; then
        echo "  → apt 安装 xvfb..."
        if sudo DEBIAN_FRONTEND=noninteractive apt-get update -qq \
            && sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -qq xvfb 2>/dev/null; then
            echo -e "  ${GREEN}✓${NC} xvfb 安装成功"
            return 0
        fi
    elif have_cmd dnf && have_cmd sudo && sudo -n true 2>/dev/null; then
        echo "  → dnf 安装 xorg-x11-server-Xvfb..."
        if sudo dnf install -y xorg-x11-server-Xvfb 2>/dev/null; then
            echo -e "  ${GREEN}✓${NC} Xvfb 安装成功"
            return 0
        fi
    elif have_cmd pacman && have_cmd sudo && sudo -n true 2>/dev/null; then
        echo "  → pacman 安装 xorg-server-xvfb..."
        if sudo pacman -Sy --noconfirm xorg-server-xvfb 2>/dev/null; then
            echo -e "  ${GREEN}✓${NC} Xvfb 安装成功"
            return 0
        fi
    elif have_cmd yum && have_cmd sudo && sudo -n true 2>/dev/null; then
        echo "  → yum 安装 xorg-x11-server-Xvfb..."
        if sudo yum install -y xorg-x11-server-Xvfb 2>/dev/null; then
            echo -e "  ${GREEN}✓${NC} Xvfb 安装成功"
            return 0
        fi
    else
        echo -e "  ${YELLOW}⚠${NC} 无法自动安装（需要免密 sudo 或手动装）"
        echo "    Ubuntu/Debian: sudo apt install -y xvfb"
        echo "    Fedora:        sudo dnf install -y xorg-x11-server-Xvfb"
        echo "    Arch:          sudo pacman -S xorg-server-xvfb"
    fi

    if have_cmd xvfb-run || have_cmd Xvfb; then
        return 0
    fi
    return 1
}

# 启动常驻 Xvfb 并 export DISPLAY
start_persistent_xvfb() {
    local disp="${XVFB_DISPLAY:-:99}"
    local sock="/tmp/.X11-unix/X${disp#:}"

    # 已有该 display 的 X socket 则直接复用
    if [ -S "$sock" ] 2>/dev/null; then
        export DISPLAY="$disp"
        echo -e "  ${GREEN}✓${NC} 复用已有 Xvfb DISPLAY=$DISPLAY"
        return 0
    fi
    if ! have_cmd Xvfb; then
        return 1
    fi
    echo "  → 启动常驻 Xvfb $disp ..."
    # 后台起；失败不阻断（上层会回退 headless）
    Xvfb "$disp" -screen 0 1920x1080x24 -ac +extension GLX +render -noreset \
        >/tmp/grok-xvfb.log 2>&1 &
    local xvfb_pid=$!
    sleep 0.5
    if [ -S "$sock" ] 2>/dev/null || kill -0 "$xvfb_pid" 2>/dev/null; then
        export DISPLAY="$disp"
        echo -e "  ${GREEN}✓${NC} Xvfb 已启动 DISPLAY=$DISPLAY pid=$xvfb_pid (log: /tmp/grok-xvfb.log)"
        return 0
    fi
    echo -e "  ${YELLOW}⚠${NC} 常驻 Xvfb 启动失败，见 /tmp/grok-xvfb.log"
    return 1
}

# 无桌面时准备虚拟显示；成功则 RUN_PREFIX / DISPLAY 已设好
step6_ensure_display() {
    echo -e "${CYAN}[6/6] 检测图形显示 / 自动 Xvfb...${NC}"

    # macOS 一般有 Aqua，不强制 Xvfb
    if [ "$(uname -s 2>/dev/null || echo unknown)" = "Darwin" ]; then
        if has_display; then
            echo -e "  ${GREEN}✓${NC} macOS 已有显示 DISPLAY=${DISPLAY:-} WAYLAND=${WAYLAND_DISPLAY:-}"
        else
            echo -e "  ${GREEN}✓${NC} macOS：不强制 Xvfb（Python 侧可 headless）"
        fi
        USE_XVFB=0
        XVFB_MODE="none"
        return 0
    fi

    if [ "${SKIP_XVFB:-0}" = "1" ]; then
        echo -e "  ${YELLOW}⚠${NC} SKIP_XVFB=1：跳过虚拟显示，走 Python 自动 headless"
        USE_XVFB=0
        XVFB_MODE="none"
        return 0
    fi

    local need_xvfb=0
    if [ "${FORCE_XVFB:-0}" = "1" ]; then
        need_xvfb=1
        echo -e "  ${CYAN}→${NC} FORCE_XVFB=1，将使用虚拟显示"
    elif ! has_display; then
        need_xvfb=1
        echo -e "  ${YELLOW}⚠${NC} 无 DISPLAY/WAYLAND（无头服务器），准备自动启用 Xvfb"
    else
        echo -e "  ${GREEN}✓${NC} 已有显示 DISPLAY=${DISPLAY:-} WAYLAND=${WAYLAND_DISPLAY:-}"
        USE_XVFB=0
        XVFB_MODE="none"
        return 0
    fi

    if [ "$need_xvfb" -ne 1 ]; then
        return 0
    fi

    try_install_xvfb || true

    # 优先 xvfb-run（自动分配 display，最省心）
    if have_cmd xvfb-run; then
        RUN_PREFIX=(xvfb-run -a -s "-screen 0 1920x1080x24")
        USE_XVFB=1
        XVFB_MODE="xvfb-run"
        echo -e "  ${GREEN}✓${NC} 将用 xvfb-run 启动（虚拟有头，利于过 CF/Turnstile）"
        echo "    若需纯 headless: SKIP_XVFB=1 ./start.sh ..."
        return 0
    fi

    # 其次常驻 Xvfb + export DISPLAY
    if start_persistent_xvfb; then
        RUN_PREFIX=()
        USE_XVFB=1
        XVFB_MODE="persistent"
        echo -e "  ${GREEN}✓${NC} 使用常驻 Xvfb DISPLAY=$DISPLAY"
        return 0
    fi

    # 都失败 → Python 侧 create_browser_options 会自动 headless
    echo -e "  ${YELLOW}⚠${NC} 未能启用 Xvfb，将依赖 Python 自动 headless（可能更易被 CF 拦）"
    echo "    手动: sudo apt install -y xvfb && xvfb-run -a ./start.sh"
    USE_XVFB=0
    XVFB_MODE="none"
    return 0
}

# ── 0. 基础工具 ──
step0_basics() {
    echo -e "${CYAN}[0/6] 检查基础工具...${NC}"
    if ! have_cmd curl && ! have_cmd wget; then
        echo -e "  ${RED}✗${NC} 需要 curl 或 wget 才能自动安装环境"
        echo "  请先安装: sudo apt install -y curl   或   sudo yum install -y curl"
        return 1
    fi
    if have_cmd curl; then
        echo -e "  ${GREEN}✓${NC} curl 可用"
    else
        echo -e "  ${GREEN}✓${NC} wget 可用"
    fi
}

# ── 1. 检测/安装 uv ──
step1_install_uv() {
    echo -e "${CYAN}[1/6] 检测/安装 uv 包管理器...${NC}"

    if have_cmd uv; then
        echo -e "  ${GREEN}✓${NC} uv: $(uv --version 2>&1)"
        return 0
    fi

    echo -e "  ${YELLOW}⚠${NC} 未检测到 uv，正在自动安装..."
    if have_cmd curl; then
        curl -LsSf https://astral.sh/uv/install.sh | sh
    else
        wget -qO- https://astral.sh/uv/install.sh | sh
    fi

    export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH"

    # 部分环境写到 ~/.profile，再补一次常见路径
    if [ -f "$HOME/.local/bin/env" ]; then
        # shellcheck disable=SC1091
        source "$HOME/.local/bin/env" 2>/dev/null || true
    fi

    if have_cmd uv; then
        echo -e "  ${GREEN}✓${NC} uv 安装成功: $(uv --version 2>&1)"
        return 0
    fi

    echo -e "  ${RED}✗${NC} uv 安装失败，请手动: curl -LsSf https://astral.sh/uv/install.sh | sh"
    return 1
}

# ── 2. 检测/安装 Python 3.13 ──
# 策略: 优先系统/mise 已有 3.13；否则用 uv 下载管理 Python
step2_install_python() {
    echo -e "${CYAN}[2/6] 检测/安装 Python 3.13...${NC}"
    PYTHON_BIN=""

    # mise 托管的 python
    if have_cmd mise && [ -f "$PROJECT_DIR/mise.toml" ]; then
        # 尝试按项目配置安装工具链
        if mise install 2>/dev/null; then
            local py
            py=$(mise which python 2>/dev/null || true)
            if [ -n "$py" ] && "$py" -c "import sys; sys.exit(0 if sys.version_info >= (3,13) else 1)" 2>/dev/null; then
                PYTHON_BIN="$py"
                echo -e "  ${GREEN}✓${NC} Python (mise): $($py --version 2>&1)  at $py"
                return 0
            fi
        fi
    fi

    # 系统 / PATH 上的 3.13+
    local candidate
    for candidate in python3.13 python3.14 python3 python; do
        if have_cmd "$candidate"; then
            if "$candidate" -c "import sys; sys.exit(0 if sys.version_info >= (3,13) else 1)" 2>/dev/null; then
                PYTHON_BIN="$candidate"
                echo -e "  ${GREEN}✓${NC} Python: $($candidate --version 2>&1)"
                return 0
            fi
        fi
    done

    # 项目 .venv
    if [ -x "$PROJECT_DIR/.venv/bin/python" ]; then
        if "$PROJECT_DIR/.venv/bin/python" -c "import sys; sys.exit(0 if sys.version_info >= (3,13) else 1)" 2>/dev/null; then
            PYTHON_BIN="$PROJECT_DIR/.venv/bin/python"
            echo -e "  ${GREEN}✓${NC} Python (.venv): $($PYTHON_BIN --version 2>&1)"
            return 0
        fi
    fi

    # 用 uv 安装托管的 Python 3.13
    echo -e "  ${YELLOW}⚠${NC} 未找到 Python 3.13+，正在通过 uv 安装..."
    if ! uv python install 3.13; then
        echo -e "  ${RED}✗${NC} uv python install 3.13 失败"
        echo "  也可手动: uv python install 3.13"
        return 1
    fi

    # 解析 uv 管理的 python 路径
    local uv_py
    uv_py=$(uv python find 3.13 2>/dev/null || true)
    if [ -n "$uv_py" ] && [ -x "$uv_py" ]; then
        PYTHON_BIN="$uv_py"
        echo -e "  ${GREEN}✓${NC} Python (uv): $($uv_py --version 2>&1)  at $uv_py"
        return 0
    fi

    # uv python find 失败时，后续 uv run 仍可自动解析
    if uv run --python 3.13 python -c "import sys; sys.exit(0 if sys.version_info >= (3,13) else 1)" 2>/dev/null; then
        PYTHON_BIN="uv-managed"
        echo -e "  ${GREEN}✓${NC} Python 3.13 已由 uv 管理（将通过 uv run 调用）"
        return 0
    fi

    echo -e "  ${RED}✗${NC} 无法安装或定位 Python 3.13"
    return 1
}

# ── 3. 同步依赖 ──
step3_sync_deps() {
    echo -e "${CYAN}[3/6] 同步 Python 依赖...${NC}"

    # 固定使用 3.13，避免系统默认 python 过旧
    if ! uv sync --python 3.13; then
        echo -e "  ${YELLOW}⚠${NC} uv sync --python 3.13 失败，尝试默认 uv sync..."
        if ! uv sync; then
            echo -e "  ${RED}✗${NC} 依赖同步失败"
            return 1
        fi
    fi

    echo -e "  ${GREEN}✓${NC} 依赖同步完成"
}

# ── 4. 检测/尝试安装浏览器 ──
step4_check_browser() {
    echo -e "${CYAN}[4/6] 检测/安装浏览器...${NC}"

    # Prefer Chromium for turnstilePatch (--load-extension works).
    # Google Chrome alone often leaves CF token=0 on Xvfb servers.
    local candidate
    local found_chromium=""
    local found_chrome=""
    for candidate in \
        chromium chromium-browser \
        /usr/bin/chromium /usr/bin/chromium-browser /snap/bin/chromium; do
        if have_cmd "$candidate" || [ -x "$candidate" ] 2>/dev/null; then
            found_chromium="$(command -v "$candidate" 2>/dev/null || echo "$candidate")"
            break
        fi
    done
    for candidate in \
        google-chrome-stable google-chrome \
        /usr/bin/google-chrome-stable /usr/bin/google-chrome \
        "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"; do
        if have_cmd "$candidate" || [ -x "$candidate" ] 2>/dev/null; then
            found_chrome="$(command -v "$candidate" 2>/dev/null || echo "$candidate")"
            break
        fi
    done

    if [ -n "$found_chromium" ]; then
        echo -e "  ${GREEN}✓${NC} Chromium: $found_chromium （可加载 turnstilePatch）"
        if [ -n "$found_chrome" ]; then
            echo -e "  ${CYAN}·${NC} 另检测到 Chrome: $found_chrome（注册默认优先 Chromium）"
        fi
        return 0
    fi

    if [ -n "$found_chrome" ]; then
        echo -e "  ${YELLOW}⚠${NC} 仅有 Google Chrome: $found_chrome"
        echo "    Chrome 无法加载 turnstilePatch，Turnstile 容易 token=0。"
        echo "    将尝试安装 Chromium..."
    else
        echo -e "  ${YELLOW}⚠${NC} 未检测到 Chrome/Chromium，尝试自动安装 Chromium..."
    fi

    local installed=false

    # Debian/Ubuntu — prefer chromium (extension support)
    if have_cmd apt-get; then
        if have_cmd sudo; then
            if sudo -n true 2>/dev/null; then
                echo "  → 使用 apt 安装 chromium（免密 sudo）..."
                if sudo apt-get update -qq && sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -qq chromium chromium-browser 2>/dev/null; then
                    installed=true
                fi
            else
                echo -e "  ${YELLOW}⚠${NC} 检测到 apt，但需要交互 sudo，跳过自动安装浏览器"
                echo "    可手动: sudo apt install -y chromium"
            fi
        fi
    # Arch
    elif have_cmd pacman; then
        if have_cmd sudo && sudo -n true 2>/dev/null; then
            echo "  → 使用 pacman 安装 chromium..."
            if sudo pacman -Sy --noconfirm chromium 2>/dev/null; then
                installed=true
            fi
        else
            echo "    可手动: sudo pacman -S chromium"
        fi
    # Fedora/RHEL
    elif have_cmd dnf; then
        if have_cmd sudo && sudo -n true 2>/dev/null; then
            echo "  → 使用 dnf 安装 chromium..."
            if sudo dnf install -y chromium 2>/dev/null; then
                installed=true
            fi
        else
            echo "    可手动: sudo dnf install -y chromium"
        fi
    # macOS Homebrew — chromium cask if available, else chrome
    elif have_cmd brew; then
        echo "  → 使用 Homebrew 安装 chromium / chrome..."
        if brew install --cask chromium 2>/dev/null || brew install --cask google-chrome 2>/dev/null; then
            installed=true
        fi
    fi

    # 再检测一次（Chromium 优先）
    for candidate in \
        chromium chromium-browser \
        /usr/bin/chromium /usr/bin/chromium-browser /snap/bin/chromium \
        google-chrome-stable google-chrome \
        /usr/bin/google-chrome-stable /usr/bin/google-chrome \
        "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome"; do
        if have_cmd "$candidate" || [ -x "$candidate" ] 2>/dev/null; then
            local resolved
            resolved="$(command -v "$candidate" 2>/dev/null || echo "$candidate")"
            case "$resolved" in
                *chromium*)
                    echo -e "  ${GREEN}✓${NC} 浏览器已就绪 (Chromium): $resolved"
                    return 0
                    ;;
                *)
                    echo -e "  ${YELLOW}⚠${NC} 仅有 Chrome: $resolved"
                    echo "    注册机将用 CDP stealth；过 CF 仍建议: sudo apt install -y chromium"
                    return 0
                    ;;
            esac
        fi
    done

    if $installed; then
        echo -e "  ${YELLOW}⚠${NC} 安装命令已执行，但仍未在 PATH 中找到浏览器，请重启终端后重试"
    else
        echo -e "  ${YELLOW}⚠${NC} 未能自动安装浏览器"
        echo "  注册需要 Chromium（推荐，可加载 turnstilePatch）或 Chrome。"
        echo "  手动安装示例:"
        echo "    Ubuntu/Debian: sudo apt install -y chromium"
        echo "    Arch:          sudo pacman -S chromium"
        echo "    macOS:         brew install --cask chromium"
        echo "  ${YELLOW}脚本将继续，注册可能因缺少浏览器 / Turnstile 失败。${NC}"
    fi
}

# ── 5. 配置文件模板 ──
step5_config_templates() {
    echo -e "${CYAN}[5/6] 检查配置文件...${NC}"

    if [ ! -f "$PROJECT_DIR/.env" ] && [ -f "$PROJECT_DIR/.env.example" ]; then
        cp "$PROJECT_DIR/.env.example" "$PROJECT_DIR/.env"
        echo -e "  ${YELLOW}⚠${NC} 已从 .env.example 创建 .env，请编辑填入真实配置"
    elif [ -f "$PROJECT_DIR/.env" ]; then
        echo -e "  ${GREEN}✓${NC} .env 已存在"
    else
        echo -e "  ${YELLOW}⚠${NC} 无 .env / .env.example（可直接写 config.json）"
    fi

    if [ ! -f "$PROJECT_DIR/config.json" ] && [ -f "$PROJECT_DIR/config.example.json" ]; then
        cp "$PROJECT_DIR/config.example.json" "$PROJECT_DIR/config.json"
        echo -e "  ${YELLOW}⚠${NC} 已从 config.example.json 创建 config.json，请编辑填入真实配置"
    elif [ -f "$PROJECT_DIR/config.json" ]; then
        echo -e "  ${GREEN}✓${NC} config.json 已存在"
    else
        echo -e "  ${RED}✗${NC} 缺少 config.json 与 config.example.json"
        return 1
    fi
}

# ── 主流程 ──
main() {
    step0_basics || exit 1
    echo ""
    step1_install_uv || exit 1
    echo ""
    step2_install_python || exit 1
    echo ""
    step3_sync_deps || exit 1
    echo ""
    step4_check_browser
    echo ""
    step5_config_templates || exit 1
    echo ""
    step6_ensure_display
    echo ""

    echo -e "${CYAN}${BOLD}============================================${NC}"
    echo -e "${GREEN}${BOLD}  环境就绪，启动 register_cli.py${NC}"
    echo -e "${CYAN}${BOLD}============================================${NC}"
    if [ "$USE_XVFB" = "1" ]; then
        echo -e "  显示模式: ${GREEN}Xvfb (${XVFB_MODE})${NC}  DISPLAY=${DISPLAY:-(xvfb-run 自动分配)}"
    elif has_display; then
        echo -e "  显示模式: ${GREEN}headed${NC}  DISPLAY=${DISPLAY:-} WAYLAND=${WAYLAND_DISPLAY:-}"
    else
        echo -e "  显示模式: ${YELLOW}headless (Python 自动)${NC}"
    fi
    echo ""

    # 默认不限数量（--count 0），除非用户显式传参
    local -a args
    if [ $# -eq 0 ]; then
        echo -e "${YELLOW}[*] 默认模式：不限数量，持续注册（Ctrl+C 停止）${NC}"
        echo ""
        args=(--count 0)
    else
        args=("$@")
    fi

    # 统一走 uv run，确保用项目 venv + 正确 Python
    # mise 仅作可选加速；有 mise.toml 时先确保工具已装，仍用 uv 启动更稳
    if have_cmd mise && [ -f "$PROJECT_DIR/mise.toml" ]; then
        mise install >/dev/null 2>&1 || true
    fi

    if [ ${#RUN_PREFIX[@]} -gt 0 ]; then
        echo -e "${CYAN}[*] 启动命令: ${RUN_PREFIX[*]} uv run --python 3.13 python -u register_cli.py ${args[*]}${NC}"
        echo ""
        exec "${RUN_PREFIX[@]}" uv run --python 3.13 python -u register_cli.py "${args[@]}"
    else
        exec uv run --python 3.13 python -u register_cli.py "${args[@]}"
    fi
}

main "$@"