start.py

#!/usr/bin/env python3
"""新架构启动与任务投递助手。"""

from __future__ import annotations

import os
import shutil
import subprocess
import sys
from pathlib import Path


ROOT_DIR = Path(__file__).resolve().parent
BACKEND_DIR = ROOT_DIR / "backend"
BACKEND_ENV_EXAMPLE = BACKEND_DIR / ".env.example"
BACKEND_ENV = BACKEND_DIR / ".env"
FRONTEND_DIR = ROOT_DIR / "frontend"
DEFAULT_ACCOUNTS_FILE = ROOT_DIR / "tokens" / "accounts.txt"


def print_banner() -> None:
    """打印启动横幅。"""
    print("=" * 58)
    print("   Codex Auto Register")
    print("   新架构启动与任务投递助手")
    print("=" * 58)
    print()


def backend_python_path() -> Path:
    """定位后端虚拟环境里的 Python。"""
    candidates = [
        BACKEND_DIR / ".venv" / "bin" / "python",
        BACKEND_DIR / ".venv" / "Scripts" / "python.exe",
    ]
    for candidate in candidates:
        if candidate.exists():
            return candidate
    raise FileNotFoundError("未找到 backend/.venv 下的 Python，请先在 backend 目录执行 uv sync")


def ensure_backend_env() -> None:
    """确保后端环境文件存在。"""
    if BACKEND_ENV.exists():
        return
    if not BACKEND_ENV_EXAMPLE.exists():
        raise FileNotFoundError("未找到 backend/.env.example")
    shutil.copyfile(BACKEND_ENV_EXAMPLE, BACKEND_ENV)
    print(f"已创建后端环境文件: {BACKEND_ENV}")


def show_menu() -> str:
    """显示主菜单。"""
    print("请选择操作:")
    print("  1. 启动后端 API")
    print("  2. 启动 Worker")
    print("  3. 投递 OpenAI 账号同步任务")
    print("  4. 查看前端启动命令")
    print("  5. 执行 OpenAI 授权库存巡检")
    print("  6. 导出数据库账号到 sub2api bundle")
    print("  7. 扫描兼容 token 文件导出到 sub2api bundle")
    print("  0. 退出")
    print()
    return input("请输入选项: ").strip()


def read_accounts_file() -> str:
    """读取账号文件路径。"""
    default_value = str(DEFAULT_ACCOUNTS_FILE)
    raw = input(f"账号文件路径（默认 {default_value}）: ").strip()
    value = raw or default_value
    path = Path(value).expanduser().resolve()
    if not path.exists() or not path.is_file():
        raise FileNotFoundError(f"账号文件不存在: {path}")
    return str(path)


def read_proxy() -> str:
    """读取可选代理地址。"""
    return input("代理地址（可留空）: ").strip()


def read_inventory_source() -> str:
    """读取库存巡检来源。"""
    raw = input("库存巡检来源：db=数据库授权，files=兼容 auth 文件（默认 db）: ").strip().lower()
    return raw if raw in {"db", "files"} else "db"


def read_yes_no(prompt: str, default: bool = True) -> bool:
    """读取 yes/no 输入。"""
    suffix = "Y/n" if default else "y/N"
    raw = input(f"{prompt} ({suffix}): ").strip().lower()
    if not raw:
        return default
    return raw in {"y", "yes", "1", "true"}


def read_target_ids() -> list[str]:
    """读取同步目标 ID 列表。"""
    raw = input("同步目标 ID，多个用逗号分隔（可留空表示自动匹配已启用目标）: ").strip()
    if not raw:
        return []
    result: list[str] = []
    for item in raw.split(","):
        value = item.strip()
        if value.isdigit() and int(value) > 0:
            result.append(value)
    return result


def exec_backend_api() -> None:
    """切换到后端 API 进程。"""
    ensure_backend_env()
    python_path = backend_python_path()
    os.chdir(BACKEND_DIR)
    os.execv(
        str(python_path),
        [
            str(python_path),
            "-m",
            "uvicorn",
            "app.main:app",
            "--reload",
        ],
    )


def exec_worker() -> None:
    """切换到 Worker 进程。"""
    ensure_backend_env()
    python_path = backend_python_path()
    os.chdir(BACKEND_DIR)
    os.execv(str(python_path), [str(python_path), "-m", "app.worker.main"])


def queue_history_sync_task() -> None:
    """调用后端 CLI 投递 OpenAI 账号同步任务。"""
    ensure_backend_env()
    python_path = backend_python_path()
    accounts_file = read_accounts_file()
    proxy = read_proxy()
    sync_to_targets = read_yes_no("导入成功后是否追加同步目标任务", default=True)
    execute_now = read_yes_no("是否立即在当前进程执行任务", default=False)
    target_ids = read_target_ids()

    cmd = [
        str(python_path),
        "-m",
        "app.modules.openai_accounts.cli",
        "--accounts-file",
        accounts_file,
        "--created-by",
        "start-launcher",
    ]
    if proxy:
        cmd.extend(["--proxy", proxy])
    if not sync_to_targets:
        cmd.append("--no-sync-to-targets")
    if execute_now:
        cmd.append("--execute-now")
    for target_id in target_ids:
        cmd.extend(["--target-id", target_id])

    subprocess.run(cmd, cwd=BACKEND_DIR, check=True)


def run_inventory_check() -> None:
    """调用维护 CLI 执行库存巡检。"""
    ensure_backend_env()
    python_path = backend_python_path()
    proxy = read_proxy()
    source = read_inventory_source()
    execute_now = read_yes_no("是否立即在当前进程执行库存巡检任务", default=True)
    cmd = [
        str(python_path),
        "-m",
        "app.modules.openai_accounts.maintenance_cli",
        "inventory-check-files" if source == "files" else "inventory-check",
        "--created-by",
        "start-launcher",
        "--json",
    ]
    if proxy:
        cmd.extend(["--proxy", proxy])
    if execute_now and source == "db":
        cmd.append("--execute-now")
    subprocess.run(cmd, cwd=BACKEND_DIR, check=True)


def export_sub2api_from_database() -> None:
    """调用维护 CLI 导出数据库账号。"""
    ensure_backend_env()
    python_path = backend_python_path()
    import_now = read_yes_no("导出后是否立即导入 sub2api", default=False)
    merge_existing = read_yes_no("是否与现有 bundle 合并写出", default=False)
    cmd = [
        str(python_path),
        "-m",
        "app.modules.openai_accounts.maintenance_cli",
        "export-sub2api",
        "--json",
    ]
    if import_now:
        cmd.append("--import-now")
    if merge_existing:
        cmd.append("--merge-existing")
    subprocess.run(cmd, cwd=BACKEND_DIR, check=True)


def export_sub2api_from_token_files() -> None:
    """调用维护 CLI 扫描兼容 token 文件。"""
    ensure_backend_env()
    python_path = backend_python_path()
    import_now = read_yes_no("导出后是否立即导入 sub2api", default=False)
    merge_existing = read_yes_no("是否与现有 bundle 合并写出", default=False)
    cmd = [
        str(python_path),
        "-m",
        "app.modules.openai_accounts.maintenance_cli",
        "export-sub2api-files",
        "--json",
    ]
    if import_now:
        cmd.append("--import-now")
    if merge_existing:
        cmd.append("--merge-existing")
    subprocess.run(cmd, cwd=BACKEND_DIR, check=True)


def show_frontend_commands() -> None:
    """展示前端启动方式。"""
    pnpm = shutil.which("pnpm") or "pnpm"
    print()
    print("请在新终端执行以下命令启动前端:")
    print(f"cd {FRONTEND_DIR}")
    print("COREPACK_ENABLE_STRICT=0 pnpm install")
    print(f"COREPACK_ENABLE_STRICT=0 {pnpm} -F @vben/web-antd dev")
    print()


def main() -> None:
    """交互式入口。"""
    print_banner()
    choice = show_menu()

    if choice == "1":
        exec_backend_api()
        return
    if choice == "2":
        exec_worker()
        return
    if choice == "3":
        queue_history_sync_task()
        return
    if choice == "4":
        show_frontend_commands()
        return
    if choice == "5":
        run_inventory_check()
        return
    if choice == "6":
        export_sub2api_from_database()
        return
    if choice == "7":
        export_sub2api_from_token_files()
        return
    if choice == "0":
        print("已退出")
        return
    print("无效选项")
    sys.exit(1)


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