AutoRpa/选品流程/main.py

from __future__ import annotations

import argparse
import time
from pathlib import Path

import requests

from auth_login import (
    DEFAULT_YAOEX_AUTH_PATH,
    DEFAULT_YBM_COOKIE_PATH,
    PROJECT_ROOT,
    BrowserRuntime,
    load_cookie_file,
    load_dotenv,
    load_yaoex_auth_file,
    login_yaoex,
    login_ybm,
    save_cookie_file,
    save_yaoex_auth,
)
from db import RepoUnavailable
from job import run_batch
from mall_config import YAOEX_MALL_NAME, YBM_MALL_NAME
from repo import OracleSelectRepo
from rules import Candidate
from search_yaoex import YaoexAuth
from search_yaoex import search_candidates as search_yaoex_candidates
from search_ybm import search_candidates as search_ybm_candidates


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="选品班次：捞表1 deal_flag=0，按条形码搜平台，结果写入表2。"
    )
    parser.add_argument("--limit", type=int, default=None, help="本班最多领取 N 条")
    parser.add_argument(
        "--platform",
        choices=["all", "yaoex", "ybm"],
        default="all",
        help="all=两平台都跑",
    )
    parser.add_argument("--secrets", type=Path, default=None, help="secrets.json 路径")
    parser.add_argument(
        "--yaoex-auth",
        type=Path,
        default=None,
        help="已登录的 1号药城认证 JSON（token/user_id/buyer_code），跳过浏览器",
    )
    parser.add_argument(
        "--ybm-cookies",
        type=Path,
        default=None,
        help="已登录的药帮忙 cookie JSON，跳过浏览器",
    )
    parser.add_argument(
        "--keep-browser",
        action="store_true",
        help="登录后不关 Chrome，搜完再留一段时间方便查看",
    )
    parser.add_argument(
        "--hold-seconds",
        type=int,
        default=180,
        help="--keep-browser 时窗口再留多少秒，默认 180",
    )
    parser.add_argument(
        "--fresh-login",
        action="store_true",
        help="忽略本地登录缓存，强制重新打开浏览器",
    )
    return parser.parse_args(argv)


def _platforms(choice: str) -> list[str]:
    if choice == "yaoex":
        return [YAOEX_MALL_NAME]
    if choice == "ybm":
        return [YBM_MALL_NAME]
    return [YAOEX_MALL_NAME, YBM_MALL_NAME]


def main(argv: list[str] | None = None, repo=None) -> None:
    args = parse_args(argv)
    load_dotenv(PROJECT_ROOT / ".env")
    flow_start = time.perf_counter()
    if repo is None:
        try:
            repo = OracleSelectRepo.from_env()
        except RepoUnavailable as exc:
            print(f"数据库不可用：{exc}")
            raise SystemExit(2) from exc

    yaoex_auth: YaoexAuth | None = None
    ybm_cookies: dict[str, str] | None = None
    login_errors: dict[str, str] = {}
    runtime = BrowserRuntime(keep=args.keep_browser, hold_seconds=args.hold_seconds)
    platforms = _platforms(args.platform)

    try:
        if YAOEX_MALL_NAME in platforms:
            try:
                cached = args.yaoex_auth or (
                    None if args.fresh_login else DEFAULT_YAOEX_AUTH_PATH
                )
                if cached and cached.exists():
                    yaoex_auth = load_yaoex_auth_file(cached)
                    print(f"1号药城使用缓存登录：{cached}")
                else:
                    yaoex_auth = login_yaoex(args.secrets, runtime=runtime)
                    save_yaoex_auth(yaoex_auth)
                    print(f"1号药城登录已缓存：{DEFAULT_YAOEX_AUTH_PATH}")
                print("1号药城认证就绪。")
            except SystemExit:
                raise
            except Exception as exc:
                login_errors[YAOEX_MALL_NAME] = str(exc)
                print(f"1号药城登录失败，本平台中止：{exc}")

        if YBM_MALL_NAME in platforms:
            try:
                cached = args.ybm_cookies or (
                    None if args.fresh_login else DEFAULT_YBM_COOKIE_PATH
                )
                if cached and cached.exists():
                    ybm_cookies = load_cookie_file(cached)
                    print(f"药帮忙使用缓存登录：{cached}")
                else:
                    ybm_cookies = login_ybm(args.secrets, runtime=runtime)
                    save_cookie_file(ybm_cookies)
                    print(f"药帮忙登录已缓存：{DEFAULT_YBM_COOKIE_PATH}")
                print("药帮忙认证就绪。")
            except SystemExit:
                raise
            except Exception as exc:
                login_errors[YBM_MALL_NAME] = str(exc)
                print(f"药帮忙登录失败，本平台中止：{exc}")

        yaoex_session = requests.Session()
        ybm_session = requests.Session()
        yaoex_cache: dict[str, list[Candidate]] = {}
        ybm_cache: dict[str, list[Candidate]] = {}
        if ybm_cookies:
            for name, value in ybm_cookies.items():
                ybm_session.cookies.set(name, value)

        searchers = {}
        if yaoex_auth is not None:

            def search_yaoex(query_key: str) -> list[Candidate]:
                assert yaoex_auth is not None
                return search_yaoex_candidates(
                    yaoex_auth, query_key, session=yaoex_session, cache=yaoex_cache
                )

            searchers[YAOEX_MALL_NAME] = search_yaoex

        if ybm_cookies is not None:

            def search_ybm(query_key: str) -> list[Candidate]:
                return search_ybm_candidates(
                    ybm_cookies, query_key, session=ybm_session, cache=ybm_cache
                )

            searchers[YBM_MALL_NAME] = search_ybm

        try:
            summary = run_batch(
                repo,
                platforms=platforms,
                searchers=searchers,
                login_errors=login_errors,
                limit=args.limit,
            )
        finally:
            yaoex_session.close()
            ybm_session.close()

        elapsed = time.perf_counter() - flow_start
        print(
            f"领取={summary.claimed} ADDED={summary.added} "
            f"SKIPPED={summary.skipped} FAILED={summary.failed} "
            f"用时 {elapsed:.2f} 秒"
        )
        print("已写入 rpa_cart_settlement_goods_data，push_status=0，未推订单中心，未加购。")
    finally:
        runtime.close()


if __name__ == "__main__":
    main()