Neo-ZQYY/scripts/server/init-server-env.py

"""
服务器环境初始化脚本：删除 skip-worktree 排除的文件/目录 + 创建 export 目录树。

在服务器上 git clone + setup-server-git.py 之后运行。
1. 删除已被 skip-worktree 标记的目录和散文件（释放磁盘空间）
2. 创建完整的 export/ 目录树（ETL/SYSTEM/BACKEND 三大类）

用法:
    python scripts/server/init-server-env.py                # 初始化 test + prod（默认）
    python scripts/server/init-server-env.py --envs test     # 只初始化 test
    python scripts/server/init-server-env.py --envs prod     # 只初始化 prod
    python scripts/server/init-server-env.py --envs test prod
"""

import argparse
import os
import shutil
import sys
from pathlib import Path

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

SERVER_ROOT = Path(r"D:\NeoZQYY")

ENV_PATHS: dict[str, Path] = {
    "test": SERVER_ROOT / "test" / "repo",
    "prod": SERVER_ROOT / "prod" / "repo",
}

# skip-worktree 后可安全删除的目录（与 setup-server-git.py DELETABLE_DIRS 一致）
DELETABLE_DIRS = [
    "export",
    "docs",
    "tests",
    "samples",
    "infra",
    ".kiro",
    ".hypothesis",
    ".pytest_cache",
    r"apps\miniprogram",
    r"scripts\ops",
    r"scripts\audit",
    r"scripts\migrate",
]

# 可安全删除的根目录散文件
DELETABLE_FILES = [
    "coach-detail-full.png",
    "customer-detail-full.png",
    "perf-records-current.png",
    "white-screen-debug.png",
    "NeoZQYY.code-workspace",
    "start-admin.bat",
    ".kiroignore",
]

# export 目录树（所有环境通用）
EXPORT_DIRS = [
    r"export\ETL-Connectors\feiqiu\JSON",
    r"export\ETL-Connectors\feiqiu\LOGS",
    r"export\ETL-Connectors\feiqiu\REPORTS",
    r"export\SYSTEM\LOGS",
    r"export\SYSTEM\REPORTS\dataflow_analysis",
    r"export\SYSTEM\REPORTS\field_audit",
    r"export\SYSTEM\REPORTS\full_dataflow_doc",
    r"export\SYSTEM\CACHE\api_samples",
    r"export\BACKEND\LOGS",
]


# ============================================================================
# 颜色输出（Windows 终端支持 ANSI）
# ============================================================================

class C:
    """ANSI 颜色码"""
    RED = "\033[91m"
    GREEN = "\033[92m"
    YELLOW = "\033[93m"
    CYAN = "\033[96m"
    MAGENTA = "\033[95m"
    WHITE = "\033[97m"
    GRAY = "\033[90m"
    RESET = "\033[0m"


def cprint(msg: str, color: str = C.WHITE) -> None:
    print(f"{color}{msg}{C.RESET}")


# ============================================================================
# 函数
# ============================================================================

def get_dir_size_mb(path: Path) -> float:
    """递归计算目录大小（MB）"""
    total = 0
    try:
        for f in path.rglob("*"):
            if f.is_file():
                total += f.stat().st_size
    except OSError:
        pass
    return round(total / (1024 * 1024), 1)


def get_file_size_mb(path: Path) -> float:
    try:
        return round(path.stat().st_size / (1024 * 1024), 2)
    except OSError:
        return 0.0


def remove_skip_worktree_items(repo_path: Path) -> None:
    """删除指定 repo 下已被 skip-worktree 标记的目录和散文件。"""
    cprint("\n  [删除] 清理 skip-worktree 排除的目录和文件...", C.YELLOW)
    freed_mb = 0.0

    # 删除目录
    for d in DELETABLE_DIRS:
        full = repo_path / d
        if full.exists() and full.is_dir():
            size = get_dir_size_mb(full)
            shutil.rmtree(full, ignore_errors=True)
            cprint(f"    已删除: {d}/ ({size} MB)", C.RED)
            freed_mb += size
        else:
            cprint(f"    跳过:   {d}/ (不存在)", C.GRAY)

    # 删除散文件
    for f in DELETABLE_FILES:
        full = repo_path / f
        if full.exists() and full.is_file():
            size = get_file_size_mb(full)
            full.unlink()
            cprint(f"    已删除: {f} ({size} MB)", C.RED)
            freed_mb += size

    # 删除根目录下所有 .png 文件
    for png in repo_path.glob("*.png"):
        if png.is_file():
            size = get_file_size_mb(png)
            png.unlink()
            cprint(f"    已删除: {png.name} ({size} MB)", C.RED)
            freed_mb += size

    cprint(f"    共释放: {freed_mb:.1f} MB", C.GREEN)


def create_export_tree(repo_path: Path) -> None:
    """在指定 repo 下创建完整的 export 目录树。"""
    cprint("\n  [创建] 初始化 export 目录树...", C.YELLOW)

    for d in EXPORT_DIRS:
        full = repo_path / d
        if not full.exists():
            full.mkdir(parents=True, exist_ok=True)
            cprint(f"    已创建: {d}/", C.CYAN)
        else:
            cprint(f"    已存在: {d}/", C.GRAY)


def test_git_setup(repo_path: Path) -> bool:
    """检查 setup-server-git.py 是否已运行。"""
    exclude_file = repo_path / ".git" / "info" / "exclude"
    if not exclude_file.exists():
        return False
    try:
        content = exclude_file.read_text(encoding="utf-8", errors="ignore")
        return "server-exclude" in content
    except OSError:
        return False


# ============================================================================
# 主流程
# ============================================================================

def main() -> None:
    parser = argparse.ArgumentParser(description="NeoZQYY 服务器环境初始化")
    parser.add_argument(
        "--envs",
        nargs="+",
        default=["test", "prod"],
        choices=["test", "prod"],
        help="要初始化的环境列表（默认: test prod）",
    )
    args = parser.parse_args()

    # Windows 终端启用 ANSI 颜色
    if sys.platform == "win32":
        os.system("")  # 触发 ANSI 支持

    cprint("=" * 44, C.WHITE)
    cprint(" NeoZQYY 服务器环境初始化", C.WHITE)
    cprint(f" 目标环境: {', '.join(args.envs)}", C.WHITE)
    cprint("=" * 44, C.WHITE)

    for env in args.envs:
        repo_path = ENV_PATHS.get(env)
        if not repo_path:
            cprint(f"\n[错误] 未知环境: {env}（可选: test, prod）", C.RED)
            continue

        cprint(f"\n{'=' * 10} 环境: {env} {'=' * 10}", C.MAGENTA)
        cprint(f"  路径: {repo_path}")

        # 检查 repo 是否存在
        if not repo_path.exists():
            cprint("  [警告] 目录不存在，跳过。请先 git clone。", C.YELLOW)
            continue

        # 检查 setup-server-git.py 是否已运行
        if not test_git_setup(repo_path):
            cprint("  [警告] 未检测到 setup-server-git.py 的配置。", C.YELLOW)
            cprint("         建议先运行: python scripts/server/setup-server-git.py", C.YELLOW)
            answer = input("  是否继续删除操作？(y/N) ").strip()
            if answer.lower() != "y":
                cprint(f"  已跳过 {env} 环境的删除操作。", C.GRAY)
                # 仍然创建 export 目录
                create_export_tree(repo_path)
                continue

        # 步骤 1：删除排除的文件/目录
        remove_skip_worktree_items(repo_path)

        # 步骤 2：创建 export 目录树
        create_export_tree(repo_path)

        cprint(f"\n  [完成] {env} 环境初始化完毕。", C.GREEN)

    cprint("\n" + "=" * 44, C.WHITE)
    cprint(" 全部完成。", C.GREEN)
    print()
    cprint(" 后续步骤:", C.WHITE)
    cprint("   1. 手动创建各环境的 .env 文件（参考 .env.template）", C.WHITE)
    cprint("   2. 确认 .env 中的 export 路径指向 repo/export/ 下对应子目录", C.WHITE)
    cprint("   3. 运行 uv sync --all-packages 安装依赖", C.WHITE)
    cprint("=" * 44, C.WHITE)


if __name__ == "__main__":
    main()