PayrollManagementSystem/backend/app/cli.py

"""命令行管理工具。

用法（在 backend/ 目录下执行）::

    uv run python -m app.cli create-admin --username admin

密码优先**交互式输入**而非命令行参数：命令行参数会进入 shell 历史，
在多用户机器上还能被 `ps` 看到。`--password` 仅为自动化脚本保留。
"""

import argparse
import getpass
import json
import signal
import sys
import threading
from datetime import datetime
from pathlib import Path
from uuid import uuid4

from sqlalchemy import select

from app.core.config import get_settings
from app.core.constants import PASSWORD_MAX_BYTES, PASSWORD_MIN_BYTES, ROLE_ADMIN
from app.core.db import dispose_engine, get_session_factory
from app.core.security import hash_password
from app.core.tplus_db import dispose_tplus_engine, get_tplus_session_factory
from app.models import Role, User, UserRole
from app.repositories.tplus_repository import get_batch
from app.services.tplus.sync_service import TPlusSyncService
from app.services.tplus.worker import run_once


def _prompt_password() -> str:
    """交互式读取并二次确认密码，循环直到符合策略。"""
    while True:
        password = getpass.getpass(
            f"请输入密码（{PASSWORD_MIN_BYTES}-{PASSWORD_MAX_BYTES} 字节）: "
        )
        length = len(password.encode("utf-8"))
        if length < PASSWORD_MIN_BYTES:
            print(f"  密码不足 {PASSWORD_MIN_BYTES} 字节，请重新输入。")
            continue
        if length > PASSWORD_MAX_BYTES:
            print(f"  密码超过 {PASSWORD_MAX_BYTES} 字节上限（中文一字占 3 字节），请重新输入。")
            continue
        if password != getpass.getpass("请再次输入以确认: "):
            print("  两次输入不一致，请重新输入。")
            continue
        return password


def _create_admin(args: argparse.Namespace) -> int:
    """幂等创建管理员账号。

    已存在时**不覆盖密码** —— 意外重跑脚本不应把已有账号的密码改掉。
    """
    factory = get_session_factory()
    with factory() as session:
        existing = session.execute(
            select(User).where(User.username == args.username)
        ).scalar_one_or_none()
        if existing is not None:
            print(f"用户 {args.username!r} 已存在（id={existing.id}），未做任何修改。")
            return 0

        role = session.execute(
            select(Role).where(Role.role_code == ROLE_ADMIN)
        ).scalar_one_or_none()
        if role is None:
            print(
                "错误：未找到 admin 角色，请先执行数据库迁移（uv run alembic upgrade head）。",
                file=sys.stderr,
            )
            return 1

        password = args.password or _prompt_password()
        user = User(
            username=args.username,
            password=hash_password(password),
            real_name=args.real_name,
            status=1,
        )
        session.add(user)
        session.flush()  # 取得自增 id 后才能建关联
        session.add(UserRole(user_id=user.id, role_id=role.id))
        session.commit()

        print(f"已创建管理员 {args.username!r}（id={user.id}），角色：{role.role_name}。")
        print("首次登录后请立即修改密码。")
        return 0


def _stage_tplus(args: argparse.Namespace) -> int:
    """读取 T+ 源库并写入 MySQL staging。"""
    target_factory = get_session_factory()
    source_factory = get_tplus_session_factory()
    watermark = datetime.fromisoformat(args.watermark) if args.watermark else None
    with source_factory() as source_session, target_factory() as target_session:
        batch = TPlusSyncService(source_session, target_session).run(
            entity_types=args.entity,
            watermark=watermark,
            batch_size=args.batch_size,
        )
    print(
        f"T+ staging 完成：批次={batch.batch_no} 状态={batch.status} "
        f"源记录={batch.source_record_count} 新增={batch.staged_record_count} "
        f"失败={batch.failed_record_count}"
    )
    return 0 if batch.status == "SUCCEEDED" else 1


def _tplus_worker(args: argparse.Namespace) -> int:
    """独立消费持久化任务；HTTP 进程不承担长时间读取。"""
    if args.connection_file:
        settings = get_settings()
        settings.t_plus_enabled, settings.t_plus_connection_file = True, args.connection_file
    if not 2 <= args.poll_seconds <= 60:
        print("轮询间隔必须在 2 到 60 秒之间", file=sys.stderr)
        return 1
    factory = get_session_factory()
    settings = get_settings()
    stop_file = getattr(args, "stop_file", None) or settings.t_plus_stop_file
    stopped = threading.Event()
    owner = str(uuid4())
    previous_handlers = {}
    if threading.current_thread() is threading.main_thread():
        for sig in (signal.SIGINT, signal.SIGTERM):
            previous_handlers[sig] = signal.signal(sig, lambda *_: stopped.set())
    try:
        while not stopped.is_set() and not (stop_file and Path(stop_file).exists()):
            with factory() as session:
                batch_no = run_once(session, worker_owner=owner)
                batch = get_batch(session, batch_no) if batch_no else None
                if batch_no and batch is None:
                    from app.models import TPlusWorkRun
                    batch = session.scalar(select(TPlusWorkRun).where(TPlusWorkRun.run_no == batch_no))
                status = batch.status if batch else None
            if batch_no:
                print(f"任务完成：{batch_no} 状态={status}", flush=True)
            if not args.watch:
                return 0 if status in (None, "SUCCEEDED") else 1
            if batch_no is None or status != "SUCCEEDED":
                stopped.wait(args.poll_seconds)
        print("T+读取服务已停止领取新任务", flush=True)
        return 0
    except KeyboardInterrupt:
        print("T+ 读取服务已停止")
        return 130
    except Exception as exc:
        print(f"T+ 读取服务停止（{type(exc).__name__}），请检查运行环境", file=sys.stderr)
        return 1
    finally:
        for sig, handler in previous_handlers.items():
            signal.signal(sig, handler)
        try:
            from app.services.tplus.automation import heartbeat

            heartbeat(factory, settings, owner, "STOPPED")
        except Exception as exc:
            print(f"T+停止状态未能写入（{type(exc).__name__}），请检查迁移与运行配置", file=sys.stderr)


def _maintenance(args):
    from app.services.tplus.automation import tick

    if args.connection_file:
        settings = get_settings()
        settings.t_plus_enabled, settings.t_plus_connection_file = True, args.connection_file
    if args.stop_file:
        get_settings().t_plus_stop_file = args.stop_file
    with get_session_factory()() as session:
        result = tick(session)
    print(json.dumps(result, ensure_ascii=False))
    return 0


def _materials(args):
    from app.services.tplus.materials import collect_materials, render_markdown

    if args.connection_file:
        settings = get_settings()
        settings.t_plus_enabled, settings.t_plus_connection_file = True, args.connection_file
    output = Path(args.output_dir).resolve()
    targets = [output / f"materials-{args.customer_id}-{args.month}.{ext}" for ext in ("json", "md")]
    if any(p.exists() for p in targets):
        raise RuntimeError("核对包已存在，请指定新的输出目录，避免覆盖既有验收依据")
    with get_tplus_session_factory()() as source:
        report = collect_materials(source, args.customer_id, args.month)
    output.mkdir(parents=True, exist_ok=True)
    targets[0].write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    targets[1].write_text(render_markdown(report), encoding="utf-8")
    print(f"样本核对包已生成；待补{len(report['missing'])}项；真实金额尚未验收：{output}")
    return 0


def main() -> int:
    """命令行入口。"""
    parser = argparse.ArgumentParser(
        prog="app.cli", description="衍水集团营销人员工资管理系统 - 命令行工具"
    )
    sub = parser.add_subparsers(dest="command", required=True)

    p_admin = sub.add_parser("create-admin", help="创建初始管理员账号（幂等）")
    p_admin.add_argument("--username", default="admin", help="登录用户名，默认 admin")
    p_admin.add_argument("--real-name", default="系统管理员", help="真实姓名")
    p_admin.add_argument(
        "--password",
        default=None,
        help="密码。不传则交互式输入（推荐）——命令行传参会留在 shell 历史中",
    )

    p_tplus = sub.add_parser("tplus-stage", help="读取 T+ 源库并写入原始 staging")
    p_tplus.add_argument(
        "--entity",
        action="append",
        choices=[
            "personnel",
            "products",
            "customers",
            "businesses",
            "shipments",
            "receipts",
            "business_items",
            "shipment_items",
            "receipt_items",
        ],
        required=True,
        help="实体类型，可重复指定",
    )
    p_tplus.add_argument(
        "--watermark",
        default=None,
        help="增量起始时间，ISO 8601 格式；不传表示全量读取",
    )
    p_tplus.add_argument("--batch-size", type=int, default=500, help="每次从源库读取的行数")

    p_worker = sub.add_parser("tplus-worker", help="执行已排队的 T+ 读取或关联核验")
    modes = p_worker.add_mutually_exclusive_group()
    modes.add_argument("--once", action="store_true", help="执行一个任务后退出（默认）")
    modes.add_argument("--watch", action="store_true", help="持续消费手动排队任务")
    p_worker.add_argument("--connection-file", default=None, help="显式指定项目外只读连接 JSON")
    p_worker.add_argument("--poll-seconds", type=int, default=3, help="空闲轮询间隔，2–60 秒")
    p_worker.add_argument("--stop-file", default=None, help="检测到停止标记后不再领取新任务")

    p_maintenance = sub.add_parser("tplus-maintain", help="检查每周原始采集与有限重试，不生成工资")
    p_maintenance.add_argument("--once", action="store_true", help="执行一次检查（默认）")
    p_maintenance.add_argument("--connection-file", default=None)
    p_maintenance.add_argument("--stop-file", default=None)
    p_materials = sub.add_parser("tplus-materials", help="只读整理客户样本核对资料，不生成工资")
    p_materials.add_argument("--customer-id", type=int, required=True)
    p_materials.add_argument("--month", required=True)
    p_materials.add_argument("--connection-file", default=None)
    p_materials.add_argument("--output-dir", required=True)

    args = parser.parse_args()

    try:
        if args.command == "create-admin":
            return _create_admin(args)
        if args.command == "tplus-worker":
            return _tplus_worker(args)
        if args.command == "tplus-maintain":
            return _maintenance(args)
        if args.command == "tplus-materials":
            return _materials(args)
        return _stage_tplus(args)
    except Exception as exc:
        print(f"命令执行失败（{type(exc).__name__}），请检查参数和只读连接配置", file=sys.stderr)
        return 1
    finally:
        dispose_engine()
        dispose_tplus_engine()


if __name__ == "__main__":
    raise SystemExit(main())