# 选品班次入口：登录 1 号药城 / 药帮忙，按条形码搜，凑齐后按本次 sku 加车。不推接口 2。

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

import requests

from auth_login import (
    DEFAULT_YAOEX_AUTH_PATH,
    DEFAULT_YBM_COOKIE_PATH,
    BrowserRuntime,
    load_cookie_file,
    load_yaoex_auth_file,
    login_yaoex,
    login_ybm,
    save_cookie_file,
    save_yaoex_auth,
)
from rpa_oracle import RepoUnavailable, load_job_env
from job import run_batch
from mall_config import (
    YAOEX_MALL_NAME,
    YAOEX_PLATFORM_TYPE,
    YBM_MALL_NAME,
    YBM_PLATFORM_TYPE,
)
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

DEFAULT_CLAIM_LIMIT = 50
_CART_ROOT = Path(__file__).resolve().parent.parent / "rpa_cart"


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    """解析本班领取条数、跑哪个平台、是否用登录缓存。

    不传 --limit 时用 50，不再把空值当成全部领走。
    """
    parser = argparse.ArgumentParser(
        description="选品班次：捞表1 deal_flag=0，按条形码搜平台，结果写入表2。"
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=DEFAULT_CLAIM_LIMIT,
        help="本班最多领取 N 条，默认 50",
    )
    parser.add_argument(
        "--platform",
        choices=["all", "yaoex", "ybm"],
        default="all",
        help="all=两平台都跑",
    )
    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)


# 把命令行平台收成要登录的商城名。all 是两个都跑。
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]


# 登录两个商城，领表 1 deal_flag=0 选品并写表 2。一边登录失败只停那一边。
def main(argv: list[str] | None = None, repo=None) -> None:
    args = parse_args(argv)
    load_job_env(Path(__file__).resolve().parent)
    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(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(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:

            # 用本班已登录的 1 号药城会话按条形码搜。同一条码走缓存。
            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

        adders, login_by_type = _cart_adders(
            platforms,
            yaoex_auth=yaoex_auth,
            ybm_cookies=ybm_cookies,
            yaoex_session=yaoex_session,
            ybm_session=ybm_session,
            login_errors=login_errors,
        )
        cart_finish = _cart_finish(repo, adders, login_by_type)
        try:
            summary = run_batch(
                repo,
                platforms=platforms,
                searchers=searchers,
                login_errors=login_errors,
                limit=args.limit,
                cart_finish=cart_finish,
            )
        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"加购成功={summary.cart_added} 加购失败={summary.cart_failed} "
            f"用时 {elapsed:.2f} 秒"
        )
        print("已写入 rpa_cart_settlement_goods_data，push_status=0，未推订单中心。")
    finally:
        runtime.close()


# 把加购包放进导入路径。选品目录仍排在前面，避免拿错同名模块。
def _ensure_cart_path() -> None:
    if _CART_ROOT.is_dir() and str(_CART_ROOT) not in sys.path:
        sys.path.append(str(_CART_ROOT))


# 用本班已登录的会话加车。登录失败的平台不配加车函数。
def _cart_adders(
    platforms: list[str],
    *,
    yaoex_auth,
    ybm_cookies,
    yaoex_session: requests.Session,
    ybm_session: requests.Session,
    login_errors: dict[str, str],
) -> tuple[dict, dict[int, str]]:
    _ensure_cart_path()
    from cart_yaoex import add_yaoex_cart
    from cart_ybm import add_ybm_cart

    adders: dict = {}
    if yaoex_auth is not None and YAOEX_MALL_NAME in platforms:

        # 按选品拆定的 sku 加入 1 号药城购物车。
        def add_yaoex(offer, qty: int):
            return add_yaoex_cart(yaoex_auth, offer, qty, session=yaoex_session)

        adders[YAOEX_PLATFORM_TYPE] = add_yaoex
    if ybm_cookies is not None and YBM_MALL_NAME in platforms:

        # 按选品拆定的 sku 加入药帮忙购物车。
        def add_ybm(offer, qty: int):
            return add_ybm_cart(ybm_cookies, offer, qty, session=ybm_session)

        adders[YBM_PLATFORM_TYPE] = add_ybm
    login_by_type: dict[int, str] = {}
    if YAOEX_MALL_NAME in login_errors:
        login_by_type[YAOEX_PLATFORM_TYPE] = login_errors[YAOEX_MALL_NAME]
    if YBM_MALL_NAME in login_errors:
        login_by_type[YBM_PLATFORM_TYPE] = login_errors[YBM_MALL_NAME]
    return adders, login_by_type


# 内存选品仓库上的盖章。字段和加购班的成功、失败、放回一致。
class _MemoryCartStamp:
    # 盖章写回选品仓库里刚插入的表 2 行。
    def __init__(self, select_repo):
        self._repo = select_repo
        self.shift_errors: list = []

    # 成功只改戳，status 保持 1。
    def mark_success(self, fid: int) -> None:
        row = self._row(fid)
        row["update_user_name"] = "加购班"
        row["fail_reason"] = None

    # 失败改 status=2，戳改成加购班。
    def mark_fail(self, fid: int, reason: str) -> None:
        row = self._row(fid)
        row["status"] = 2
        row["fail_reason"] = reason
        row["update_user_name"] = "加购班"

    # 登录或查询失败放回待领。status 不动。
    def mark_release(self, fid: int, reason: str) -> None:
        row = self._row(fid)
        row["update_user_name"] = "系统管理员"
        row["fail_reason"] = reason

    # 记下未捕获异常原文。不改表 2。
    def save_shift_errors(self, items) -> None:
        self.shift_errors.extend(items)

    # 按刚插入的 fid 找表 2 行。找不到则 KeyError。
    def _row(self, fid: int) -> dict:
        for row in self._repo.cart:
            if int(row.get("fid") or 0) == int(fid):
                return row
        raise KeyError(fid)


# 真库和内存选品仓库都能盖加购章。真库复用选品这条池，避免再开一条连接池。
def _stamp_repo(select_repo):
    if isinstance(select_repo, OracleSelectRepo):
        _ensure_cart_path()
        from cart_repo import OracleCartRepo

        return OracleCartRepo(select_repo._pool)
    return _MemoryCartStamp(select_repo)


# 落库之后按拆定的行加车。返回加购成功行数和失败行数。
def _cart_finish(select_repo, adders: dict, login_by_type: dict[int, str]):
    _ensure_cart_path()
    from cart_job import add_known_lines

    # 同一单的拆量行一次加完。
    def finish(lines):
        summary = add_known_lines(
            _stamp_repo(select_repo),
            lines,
            adders=adders,
            login_errors=login_by_type,
        )
        return summary.added, summary.failed

    return finish


if __name__ == "__main__":
    main()