InboundVerify/inbound_verify/store.py

# -*- coding: utf-8 -*-
"""
store.py — 到货核销数据持久化（PostgreSQL）

职责：把 downloads/ 下各站点下载的应到 / 实到 / 未到 Excel 解析后，幂等写入 PostgreSQL。
与下载流程解耦：本模块只读 downloads/ 现有文件入库，不关心谁触发下载、下载了几次。

设计要点：
  - 三张表：expected_record（运单级）/ actual_record（扫描件级）
            / undelivered_record（百世站点直供未到明细，子单级）
  - 每行 = 统一核心列 + raw JSONB（站点原始全列，key=原列名，一字段不丢）
  - 幂等：业务唯一键 UPSERT，重复下载天然合并、零冗余
  - 单号一律按文本读写（dtype=str），防长数字被科学计数 / 精度丢失

命令行：
  python -m inbound_verify.store createdb        创建数据库（幂等）
  python -m inbound_verify.store init            建表（幂等 CREATE TABLE IF NOT EXISTS）
  python -m inbound_verify.store ingest [site]   入库全站或单站（幂等 UPSERT）
  python -m inbound_verify.store ingest-one <site> <kind>  仅入库指定站/类（钩子同款路由）
  python -m inbound_verify.store all             createdb → init → 全站 ingest 一条龙
"""

import os
import sys
from datetime import date, datetime

import pandas as pd
import psycopg
import yaml
from psycopg.types.json import Jsonb

from inbound_verify.paths import BASE_DIR, CONFIG_PATH, DOWNLOAD_DIR
from inbound_verify.domain import (
    BAISHI_FILE,
    _site_cfg,
)  # 站点 / 文件名配置（单一来源）

from inbound_verify import compare  # _read_business_dates（比对侧业务日期读取）

SCHEMA_PATH = os.path.join(BASE_DIR, "schema.sql")

# 有应到 / 实到的 4 站（百世只有站点直供的未到明细，单独处理）
ALL_SITES = ["顺心", "中通", "韵达", "安能"]


# ============================== 配置 / 连接 ==============================


def _load_pg_config():
    """从 config.yaml 读 postgres 段；缺失项给默认。"""
    if not os.path.exists(CONFIG_PATH):
        raise FileNotFoundError(
            f"未找到配置文件 {CONFIG_PATH}（请参考 config.example.yaml 创建 config.yaml）"
        )
    with open(CONFIG_PATH, "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f) or {}
    pg = cfg.get("postgres") or {}
    return {
        "host": pg.get("host", "127.0.0.1"),
        "port": int(pg.get("port", 5432)),
        "user": pg.get("user", "postgres"),
        "password": pg.get("password", ""),
        "dbname": pg.get("dbname", "CQHXDB"),
        "schema": pg.get("schema", "inbound_verify"),
        "auto_ingest": bool(pg.get("auto_ingest", True)),
        "connect_timeout_seconds": int(pg.get("connect_timeout_seconds", 5)),
    }


def _connect(dbname):
    """用关键字参数连接（避开 conninfo 对密码特殊字符的解析）。
    options 设 search_path 到专用 schema + 会话级 statement_timeout=30s
    （cpolar 隧道上防失控查询；connect_timeout 守连接阶段）。"""
    c = _load_pg_config()
    return psycopg.connect(
        host=c["host"],
        port=c["port"],
        dbname=dbname,
        user=c["user"],
        password=c["password"],
        options=f"-c search_path={c['schema']} -c statement_timeout=30s",
        connect_timeout=c["connect_timeout_seconds"],
    )


def ingest_enabled():
    """是否启用下载后自动入库（config.yaml postgres.auto_ingest，默认 True）。
    供 runtime 钩子判定开关，避免它伸手进 _load_pg_config。"""
    return _load_pg_config()["auto_ingest"]


# ============================== 建库 / 建表 ==============================


def create_database():
    """连接维护库 postgres，创建目标数据库（幂等）。"""
    c = _load_pg_config()
    target = c["dbname"]
    with _connect("postgres") as conn:  # autocommit：CREATE DATABASE 不能在事务里
        conn.autocommit = True
        with conn.cursor() as cur:
            cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (target,))
            if cur.fetchone():
                print(f">> [db] 数据库 {target} 已存在，跳过创建")
                return
            cur.execute(f'CREATE DATABASE "{target}"')
            print(f">> [db] 已创建数据库 {target}")


def init_schema():
    """在目标库执行 schema.sql（幂等）。"""
    with open(SCHEMA_PATH, "r", encoding="utf-8") as f:
        sql = f.read()
    with _connect(_load_pg_config()["dbname"]) as conn:
        with conn.cursor() as cur:
            cur.execute(sql)
        conn.commit()
    print(">> [db] 表结构已就绪")


# ============================== 解析辅助 ==============================

# 实到单号列 / 基号列映射（口径取自 domain，与 STATIONS 对齐）：
#   piece     = 实到表里「每件」的单号列（扫描单号 / 子单号 / 复合串）
#   waybill   = 与应到运单号对齐的干净列（中通无干净列，由 piece 复合串 v[:-8] 推导）
#   scan_time = 扫描时间列（缺失则不填，原始值仍在 raw）
#   scan_site = 扫描网点列
ACTUAL_COLMAP = {
    "中通": {
        "piece": "运单号",
        "waybill": None,
        "scan_time": "扫描时间",
        "scan_site": "扫描网点",
    },
    "顺心": {
        "piece": "子单号",
        "waybill": "运单号",
        "scan_time": "操作时间",
        "scan_site": "操作网点",
    },
    "韵达": {
        "piece": "子单号",
        "waybill": "主单号",
        "scan_time": "扫描时间",
        "scan_site": "扫描站点",
    },
    "安能": {
        "piece": "扫描单号",
        "waybill": "所属单号",
        "scan_time": "扫描时间",
        "scan_site": "扫描网点",
    },
}

_TIME_FMTS = (
    "%Y-%m-%d %H:%M:%S",
    "%Y-%m-%d %H:%M",
    "%Y/%m/%d %H:%M:%S",
    "%Y/%m/%d %H:%M",
    "%Y-%m-%d",
    "%Y/%m/%d",
)


def _to_int(v):
    """尽力把单元格转 int（件数）；空 / 非数返回 None。"""
    s = str(v).strip().replace(",", "")
    if s in ("", "-", "nan", "None"):
        return None
    try:
        return int(float(s))
    except (TypeError, ValueError):
        return None


def _parse_time(v):
    """尽力解析多种时间格式为 datetime；失败返回 None（原始值在 raw 里）。"""
    if v is None:
        return None
    if isinstance(v, datetime):
        return v
    s = str(v).strip()
    if not s or s in ("nan", "NaT"):
        return None
    for fmt in _TIME_FMTS:
        try:
            return datetime.strptime(s, fmt)
        except ValueError:
            continue
    try:  # 兜底：交给 pandas 推断
        return pd.to_datetime(s).to_pydatetime()
    except Exception:
        return None


def _parse_date(v):
    if not v:
        return None
    try:
        return datetime.strptime(str(v).strip(), "%Y-%m-%d").date()
    except (ValueError, TypeError):
        return None


def _raw_row(row):
    """把一行原始记录转成 JSONB 兼容 dict（key=原列名；丢空值，保留全部有值字段）。"""
    out = {}
    for k, v in row.items():
        if v is None:
            continue
        if isinstance(v, float) and pd.isna(v):
            continue
        s = str(v).strip()
        if s == "":
            continue
        if isinstance(v, (datetime, date)):
            out[str(k)] = v.isoformat()
        else:
            out[str(k)] = v
    return out


def _read_business_dates():
    """从状态库读各站本次业务日期（与报告口径一致；读不到返回空 dict）。"""
    try:
        return compare._read_business_dates(ALL_SITES + ["百世"]) or {}
    except Exception as e:
        print(f">> [warn] 读取业务日期失败（不影响入库）: {e}")
        return {}


# ============================== UPSERT SQL ==============================

_SQL_EXPECTED = """
    INSERT INTO expected_record
        (site, waybill_no, handover_no, handover_pieces, order_pieces, business_date, raw)
    VALUES (%s,%s,%s,%s,%s,%s,%s)
    ON CONFLICT (site, waybill_no) DO UPDATE SET
        handover_no     = EXCLUDED.handover_no,
        handover_pieces = EXCLUDED.handover_pieces,
        order_pieces    = EXCLUDED.order_pieces,
        business_date   = COALESCE(EXCLUDED.business_date, expected_record.business_date),
        raw             = EXCLUDED.raw,
        ingested_at     = now()
"""

_SQL_ACTUAL = """
    INSERT INTO actual_record
        (site, waybill_no, piece_no, scan_time, scan_site, raw)
    VALUES (%s,%s,%s,%s,%s,%s)
    ON CONFLICT (site, piece_no) DO UPDATE SET
        waybill_no  = COALESCE(EXCLUDED.waybill_no, actual_record.waybill_no),
        scan_time   = EXCLUDED.scan_time,
        scan_site   = EXCLUDED.scan_site,
        raw         = EXCLUDED.raw,
        ingested_at = now()
"""

_SQL_UNDELIVERED = """
    INSERT INTO undelivered_record
        (site, waybill_no, piece_no, biz_type, last_scan, raw)
    VALUES (%s,%s,%s,%s,%s,%s)
    ON CONFLICT (site, piece_no) DO UPDATE SET
        waybill_no  = COALESCE(EXCLUDED.waybill_no, undelivered_record.waybill_no),
        biz_type    = EXCLUDED.biz_type,
        last_scan   = EXCLUDED.last_scan,
        raw         = EXCLUDED.raw,
        ingested_at = now()
"""

_SQL_BAISHI_DAILY_STATS = """
    INSERT INTO baishi_daily_stats
        (site, business_date, expected_pieces, arrived_pieces, undelivered_pieces, raw)
    VALUES (%s,%s,%s,%s,%s,%s)
    ON CONFLICT (site, business_date) DO UPDATE SET
        expected_pieces    = COALESCE(EXCLUDED.expected_pieces, baishi_daily_stats.expected_pieces),
        arrived_pieces     = COALESCE(EXCLUDED.arrived_pieces, baishi_daily_stats.arrived_pieces),
        undelivered_pieces = COALESCE(EXCLUDED.undelivered_pieces, baishi_daily_stats.undelivered_pieces),
        raw                = EXCLUDED.raw,
        ingested_at        = now()
"""


# ============================== 入库 ==============================


def _ingest_expected(cur, site, business_date):
    """入库单站应到（运单级，按 waybill_no 去重 keep-first 后 UPSERT）。"""
    cfg = _site_cfg(site)
    path = os.path.join(DOWNLOAD_DIR, cfg["exp"])
    if not os.path.exists(path):
        print(f"   [跳过] {site} 应到：文件不存在 {cfg['exp']}")
        return 0
    df = pd.read_excel(path, dtype=str).fillna("")
    df = df.drop_duplicates(subset=[cfg["exp_wb"]], keep="first")
    biz = _parse_date(business_date)
    rows = []
    for r in df.to_dict("records"):
        wb = str(r.get(cfg["exp_wb"], "")).strip()
        if not wb:
            continue
        rows.append(
            (
                site,
                wb,
                str(r.get(cfg["exp_jd"], "")).strip() or None,
                _to_int(r.get(cfg["exp_qty"])),
                _to_int(r.get("录单件数")),
                biz,
                Jsonb(_raw_row(r)),
            )
        )
    if rows:
        cur.executemany(_SQL_EXPECTED, rows)
    print(f"   [应到] {site}：{len(rows)} 条运单")
    return len(rows)


def _ingest_actual(cur, site):
    """入库单站实到（扫描件级，按 piece_no UPSERT）。"""
    cfg = _site_cfg(site)
    path = os.path.join(DOWNLOAD_DIR, cfg["act"])
    if not os.path.exists(path):
        print(f"   [跳过] {site} 实到：文件不存在 {cfg['act']}")
        return 0
    cm = ACTUAL_COLMAP[site]
    df = pd.read_excel(path, dtype=str).fillna("")
    if site == "韵达":
        # 韵达业务清洗：保留「交接单号」为空的行（到/接件扫描），
        # 抛弃「交接单号」不为空的行（派件/签收等，属重复数据）。
        # 再按子单号去重（一件多扫只留一条；清洗后子单号已天然唯一，drop 为保险）。
        df = df[df["交接单号"].astype(str).str.strip() == ""]
        df = df.drop_duplicates(subset=[cm["piece"]], keep="last")
    rows = []
    for r in df.to_dict("records"):
        piece = str(r.get(cm["piece"], "")).strip()
        if not piece:
            continue
        if site == "中通":  # 复合串 H+运单号(12)+总数(4)+顺序(4)：基号 = v[:-8]
            waybill = piece[:-8] if (len(piece) > 8 and piece[-4:].isdigit()) else piece
        else:
            waybill = str(r.get(cm["waybill"], "")).strip() or None
        rows.append(
            (
                site,
                waybill,
                piece,
                _parse_time(r.get(cm["scan_time"])),
                str(r.get(cm["scan_site"], "")).strip() or None,
                Jsonb(_raw_row(r)),
            )
        )
    if rows:
        cur.executemany(_SQL_ACTUAL, rows)
    print(f"   [实到] {site}：{len(rows)} 条扫描")
    return len(rows)


def _ingest_undelivered_baishi(cur):
    """入库百世应到未到明细（子单级，按 (site, piece_no) UPSERT）。"""
    path = os.path.join(DOWNLOAD_DIR, BAISHI_FILE)
    if not os.path.exists(path):
        print(f"   [跳过] 百世 未到：文件不存在 {BAISHI_FILE}")
        return 0
    df = pd.read_excel(path, dtype=str).fillna("")
    rows = []
    for r in df.to_dict("records"):
        rows.append(
            (
                "百世",
                str(r.get("运单号", "")).strip() or None,
                str(r.get("子单号", "")).strip() or None,
                str(r.get("类型", "")).strip() or None,
                str(r.get("最新扫描记录", "")).strip() or None,
                Jsonb(_raw_row(r)),
            )
        )
    if rows:
        cur.executemany(_SQL_UNDELIVERED, rows)
    print(f"   [未到] 百世：{len(rows)} 条")
    return len(rows)


def upsert_baishi_daily_stats(exp, arr, business_date=None):
    """直接落库百世当日应到/实到基数（应扫/已扫，站级日聚合）。
    供 baishi 下载时抓到基数后直接调用（一步落库，不绕 state_store→store）。
    business_date 默认今天（百世固定当天）。best-effort：失败只告警，不影响下载流程。"""
    biz = business_date or date.today()
    if exp is None and arr is None:
        return
    undel = (exp - arr) if (exp is not None and arr is not None) else None
    try:
        with _connect(_load_pg_config()["dbname"]) as conn:
            with conn.cursor() as cur:
                cur.execute(
                    _SQL_BAISHI_DAILY_STATS,
                    (
                        "百世",
                        biz,
                        exp,
                        arr,
                        undel,
                        Jsonb({"expected": exp, "arrived": arr, "undelivered": undel}),
                    ),
                )
            conn.commit()
        print(f"   [基数] 百世 {biz}: 应扫 {exp} / 已扫 {arr} / 未扫 {undel}")
    except Exception as e:
        print(f"   [基数] 百世 {biz} 入库失败（不影响下载）: {e}")


def ingest(site=None):
    """入库：指定 site 则单站（百世只入未到），否则全站。返回总条数。"""
    dates = _read_business_dates()
    total = 0
    with _connect(_load_pg_config()["dbname"]) as conn:
        with conn.cursor() as cur:
            sites = ALL_SITES if site in (None, "百世") else [site]
            for s in sites:
                if s not in ALL_SITES:
                    print(f"   [跳过] 不支持应到/实到的站点: {s}")
                    continue
                total += _ingest_expected(cur, s, dates.get(s))
                total += _ingest_actual(cur, s)
            if site in (None, "百世"):
                total += _ingest_undelivered_baishi(cur)
        conn.commit()
    print(f">> [ingest] 完成，共 {total} 条")
    return total


def ingest_task(site, kind):
    """按 (site, kind) 入库本次刚下载的文件（幂等 UPSERT），返回总条数。
    与 ingest(site) 的区别：只入本次刷新的那一类，避免重读写另一类文件（同步钩子里减少阻塞）。
    kind 路由：
      expected/actual        各入其列；
      undelivered 百世        入未到；
      undelivered 4 站        _site_undelivered_handler 内部连带下了 expected+actual，故入两者；
      __compare__ / 其它组合  返回 0。
    """
    if site == "__compare__":
        return 0
    # 百世无应到/实到（只有站点直供的未到）；非 undelivered 直接返回 0，避免
    # _ingest_expected/_ingest_actual 走到 _site_cfg(百世)=None 而 TypeError。
    if site == "百世" and kind != "undelivered":
        return 0
    dates = _read_business_dates()
    total = 0
    with _connect(_load_pg_config()["dbname"]) as conn:
        with conn.cursor() as cur:
            if kind == "expected":
                total += _ingest_expected(cur, site, dates.get(site))
            elif kind == "actual":
                total += _ingest_actual(cur, site)
            elif kind == "undelivered":
                if site == "百世":
                    total += _ingest_undelivered_baishi(cur)
                else:  # 顺心/中通/韵达/安能
                    total += _ingest_expected(cur, site, dates.get(site))
                    total += _ingest_actual(cur, site)
            # 其它组合（如 百世/expected，正常不经钩子触发）：防御性返回 0
        conn.commit()
    return total


def get_existing_handover_nos(site):
    """查该站点已落库的交接单号集合（expected_record.handover_no）。
    供"提交导出任务前"去重：已落库的交接单号不再重复提交导出任务。
    PG 不可用（cpolar 抖动等）时返回空集 + 告警，调用方按"未确认存在"处理
    （继续提交导出，UPSERT 兜底，绝不因去重查询失败而漏数据）。"""
    try:
        with _connect(_load_pg_config()["dbname"]) as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT handover_no FROM expected_record "
                    "WHERE site=%s AND handover_no IS NOT NULL AND handover_no <> ''",
                    (site,),
                )
                return {str(r[0]).strip() for r in cur.fetchall()}
    except Exception as e:
        print(f">> [去重] 查询已落库交接单号失败（{site}），本次不去重: {e}")
        return set()


# ============================== PG 数据存在性查询 ==============================


def has_data(site, kind, target_date):
    """查询 PG：指定站点在 target_date 是否有业务数据。
    target_date: str 'YYYY-MM-DD' 或 date 对象。
    返回 (has_rows: bool, count: int)。
    PG 不可达时返回 (False, 0)，不抛异常——调用方按「未确认存在」处理。

    kind 路由：
      expected     → expected_record (business_date)
      actual       → actual_record (scan_time::date)
      undelivered  → 百世: baishi_daily_stats；4 站: 不单独查（由调用方 expected∧actual 派生）
    """
    if site == "百世" and kind == "undelivered":
        sql = (
            "SELECT COUNT(*) FROM baishi_daily_stats"
            " WHERE site = %s AND business_date = %s"
        )
        params = (site, target_date)
    elif kind == "expected":
        sql = (
            "SELECT COUNT(*) FROM expected_record"
            " WHERE site = %s AND business_date = %s"
        )
        params = (site, target_date)
    elif kind == "actual":
        sql = (
            "SELECT COUNT(*) FROM actual_record"
            " WHERE site = %s AND scan_time::date = %s"
        )
        params = (site, target_date)
    else:
        return (False, 0)
    try:
        with _connect(_load_pg_config()["dbname"]) as conn:
            with conn.cursor() as cur:
                cur.execute(sql, params)
                row = cur.fetchone()
                cnt = int(row[0]) if row else 0
                return (cnt > 0, cnt)
    except Exception as e:
        print(f">> [状态] PG 查询 {site}/{kind}/{target_date} 失败: {e}")
        return (False, 0)


# ============================== 命令行 ==============================


def main():
    cmd = sys.argv[1] if len(sys.argv) > 1 else "all"
    site = sys.argv[2] if len(sys.argv) > 2 else None
    if cmd == "createdb":
        create_database()
    elif cmd == "init":
        init_schema()
    elif cmd == "ingest":
        ingest(site)
    elif cmd == "all":
        create_database()
        init_schema()
        ingest()
    elif cmd == "ingest-one":
        kind = sys.argv[3] if len(sys.argv) > 3 else None
        if not site or kind not in ("expected", "actual", "undelivered"):
            print(
                "用法: python -m inbound_verify.store ingest-one <site> <expected|actual|undelivered>"
            )
            sys.exit(1)
        total = ingest_task(site, kind)
        print(f">> [ingest-one] {site}/{kind} 入库 {total} 条")
    else:
        print(__doc__)
        sys.exit(1)


if __name__ == "__main__":
    main()