#!/usr/bin/env python3
"""通用分段拉取（客户 / 活动）— 按 start_time 自动切段，段内 OFFSET<1000，超 1000 行自动对半再切。

背景：
- MCP search_by_mql 的 `LIMIT ... OFFSET` 硬上限 1000，单段最多拉 1050 行，超了会静默截断。
- `_recover_pull_offset.py` 用**静态日期分片**，片大小随数据增长会失效（尤其最新那段 "至今"）。
- 本脚本复用 `pull_orders.py::_pull_range` 的自动对半切分思路，对任意含 start_time 的类型通用。

用法:
    python3 _pull_segmented.py customers   [--months N]
    python3 _pull_segmented.py activities  [--months N]

安全闸：写临时文件 → 校验条数 ≥ 基准下限 → 才 os.replace 覆盖正式文件。
"""
import json, os, sys, time
from datetime import datetime, timedelta, date

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import _recover_pull_offset as R  # noqa: E402  复用 mcp_call / parse_items / 常量

DATA_DIR = R.DATA_DIR
SALES_PK = R.SALES_PK
TODAY = date.today()

SPECS = {
    "customers": {
        "type_key": R.TYPE_CUSTOMER,
        "cols": "name, field_c3224b, field_5ed7ab, field_c8e80d, field_17186c, field_6415cf, work_item_id",
        "out": "mcp_customers.json",
        "min_ok": 1600,
        # 客户表最早一条约 2024-05，往前多留一个月
        "start": "2024-04-01",
        "months": 0,  # 0 = 用 start 固定起点
        "wrap": lambda items, now: {"updated_at": now, "customers": items},
        "sort_key": lambda x: x.get("name", ""),
        "sort_reverse": False,
        "digit": 2,
    },
    "activities": {
        "type_key": R.TYPE_ACTIVITY,
        "cols": "name, start_time, field_b99055, field_5f20fc, owner, work_item_id",
        "out": "daily_activities_mcp.json",
        "min_ok": 4000,
        "start": None,
        "days": 70,  # 滚动近 70 天窗口（≈2.3 个月，与历史 3.1MB 规模一致）
        "months": 0,
        "wrap": lambda items, now: {"updated": now, "total": len(items), "items": items},
        "sort_key": lambda x: x.get("start_time", ""),
        "sort_reverse": True,
        "digit": 2,
    },
}


def window_start(spec):
    if spec.get("days"):
        return (TODAY - timedelta(days=spec["days"])).isoformat()
    if spec["months"]:
        y, m = TODAY.year, TODAY.month
        for _ in range(spec["months"] - 1):
            m -= 1
            if m == 0:
                m, y = 12, y - 1
        return date(y, m, 1).isoformat()
    return spec["start"]


def pull_segment(type_key, cols, start, end, depth=0):
    """拉取 [start, end) 段（OFFSET 翻页），返回 (items, total)。"""
    where = []
    if start:
        where.append(f"start_time >= '{start}'")
    if end:
        where.append(f"start_time < '{end}'")
    mql = f"SELECT {cols} FROM `销售管理`.`{type_key}`"
    if where:
        mql += " WHERE " + " AND ".join(where)
    mql += " ORDER BY start_time DESC"

    items, seen = [], set()
    offset, empty, total = 0, 0, None
    for _ in range(40):  # 硬上限 2000 行/段
        result = R.mcp_call("search_by_mql",
                            {"project_key": SALES_PK, "mql": f"{mql} LIMIT 50 OFFSET {offset}",
                             "session_id": ""})
        if result is None:
            print(f"    ❌ {start or '-'}~{end or '-'} OFFSET={offset} 拉取失败")
            break
        if total is None:
            lst = result.get("list") or []
            if lst and isinstance(lst[0], dict):
                total = lst[0].get("count")
        batch = R.parse_items(result)
        if not batch:
            empty += 1
            if empty >= 2:
                break
        else:
            empty = 0
            for it in batch:
                wid = str(it.get("work_item_id", "") or "")
                if wid and wid in seen:
                    continue
                if wid:
                    seen.add(wid)
                items.append(it)
        offset += 50
        if offset > 1000:
            break  # OFFSET 上限
        time.sleep(0.2)
    return items, total


def pull_range(type_key, cols, start):
    """7 天一段；段总量 >1000 时对半递归再切（最多 2 层）。"""
    def seg(s, e, depth=0):
        got, total = pull_segment(type_key, cols, s, e)
        if total and total > 1000 and depth < 2 and s and e:
            d1 = datetime.strptime(s, "%Y-%m-%d")
            d2 = datetime.strptime(e, "%Y-%m-%d")
            mid = (d1 + (d2 - d1) / 2).strftime("%Y-%m-%d")
            if mid not in (s, e):
                print(f"    ↩︎ 段 {s}~{e} 超 1000 行({total})，切成 {s}~{mid} + {mid}~{e}")
                return seg(s, mid, depth + 1) + seg(mid, e, depth + 1)
        return got

    chunks, d = [], datetime.strptime(start, "%Y-%m-%d").date()
    while d < TODAY:
        nxt = min(d + timedelta(days=7), TODAY)
        chunks.append((d.isoformat(), nxt.isoformat() if nxt < TODAY else None))
        d = nxt
    out = []
    for s, e in chunks:
        got = seg(s, e)
        out.extend(got)
        print(f"  段 {s}~{e or 'now'}: {len(got)} 条")
    return out


def main():
    which = sys.argv[1] if len(sys.argv) > 1 else ""
    spec = SPECS.get(which)
    if not spec:
        print("用法: python3 _pull_segmented.py customers|activities [--months N]")
        sys.exit(2)
    if "--months" in sys.argv:
        spec = dict(spec)
        spec["months"] = int(sys.argv[sys.argv.index("--months") + 1])
        spec["days"] = 0
    if "--days" in sys.argv:
        spec = dict(spec)
        spec["days"] = int(sys.argv[sys.argv.index("--days") + 1])
        spec["months"] = 0

    start = window_start(spec)
    now = time.strftime("%Y-%m-%d %H:%M:%S")
    print("=" * 60)
    print(f"分段拉取 {which}  ({start} ~ now)")
    print("=" * 60)

    raw = pull_range(spec["type_key"], spec["cols"], start)
    seen, items = set(), []
    for it in raw:
        wid = str(it.get("work_item_id", "") or "")
        if wid and wid in seen:
            continue
        if wid:
            seen.add(wid)
        items.append(it)
    items.sort(key=spec["sort_key"], reverse=spec["sort_reverse"])
    print(f"  合计 {len(items)} 条（全局按 work_item_id 去重）")

    out_path = os.path.join(DATA_DIR, spec["out"])
    tmp = out_path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(spec["wrap"](items, now), f, ensure_ascii=False, indent=spec["digit"])

    if len(items) < spec["min_ok"]:
        print(f"  🔴 仅 {len(items)} 条（基准 ≥{spec['min_ok']}），不替换正式文件，保留 {tmp}")
        sys.exit(1)
    os.replace(tmp, out_path)
    print(f"  ✅ 写入 {out_path}: {len(items)} 条")


if __name__ == "__main__":
    main()