#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""自适应分段拉取（客户 / 活动 3b）— 先查 count，>900 行自动对半切，避免 7 天碎片化。

背景：MCP search_by_mql 的 LIMIT OFFSET 硬上限 1000（单查询最多 1050 行），
而 _pull_segmented.py 用固定 7 天分段，客户表 2.5 年跨度会切出 130+ 段 → 单次运行 >7 分钟被超时杀掉。
本脚本先用 LIMIT 1 拿 count，只对 >900 行的区间对半递归，段数从 130 降到 ~10。

用法:
    python3 pull_adaptive.py customers
    python3 pull_adaptive.py activities
安全闸：写 .tmp → 校验条数 ≥ min_ok → os.replace 覆盖。
"""
import json, os, sys, time
from datetime import datetime, date, timedelta

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import _recover_pull_offset as R  # noqa: E402

DATA_DIR = R.DATA_DIR
PK = R.SALES_PK
TODAY = date.today()
NOW = time.strftime("%Y-%m-%d %H:%M:%S")

SPECS = {
    "customers": {
        "type_key": R.TYPE_CUSTOMER,
        "cols": "name, field_c3224b, field_5ed7ab, field_c8e80d, field_17186c, field_8c870a, "
                "field_79aa0b, field_6415cf, work_item_id, field_e62869, owner, current_status_operator",
        "out": "mcp_customers.json",
        "min_ok": 1600,
        "chunks": [(None, "2024-06-01"), ("2024-06-01", "2025-01-01"),
                   ("2025-01-01", "2025-07-01"), ("2025-07-01", "2026-01-01"),
                   ("2026-01-01", "2026-07-01"), ("2026-07-01", None)],
        "wrap": lambda items: {"updated_at": NOW, "customers": items},
        "sort": (lambda x: x.get("name", ""), False),
    },
    "activities": {
        "type_key": R.TYPE_ACTIVITY,
        "cols": "name, start_time, field_b99055, field_5f20fc, field_76654e, field_5ed7ab, owner, work_item_id",
        "out": "daily_activities_mcp.json",
        "min_ok": 4000,
        "chunks": [("2026-07-01", "2026-08-01"), ("2026-08-01", "2026-09-01"),
                   ("2026-09-01", "2026-10-01"), ("2026-10-01", None)],
        "wrap": lambda items: {"updated": NOW, "total": len(items), "items": items},
        "sort": (lambda x: x.get("start_time", ""), True),
    },
}

FAILS = []
STUBS = [0]


def build(type_key, cols, start, end, tail):
    where = []
    if start:
        where.append("start_time >= '%s'" % start)
    if end:
        where.append("start_time < '%s'" % end)
    mql = "SELECT %s FROM `销售管理`.`%s`" % (cols, type_key)
    if where:
        mql += " WHERE " + " AND ".join(where)
    mql += " ORDER BY start_time DESC"
    return mql + " " + tail


def count_range(type_key, cols, start, end):
    res = R.mcp_call("search_by_mql", {"project_key": PK,
                                      "mql": build(type_key, cols, start, end, "LIMIT 1 OFFSET 0"),
                                      "session_id": ""})
    if not res:
        return None
    lst = res.get("list") or []
    if lst and isinstance(lst[0], dict):
        try:
            return int(lst[0].get("count"))
        except Exception:
            return None
    return None


def pull_seg(type_key, cols, start, end):
    items, seen, offset, empty = [], set(), 0, 0
    for _ in range(21):  # OFFSET 0..1000
        res = R.mcp_call("search_by_mql", {"project_key": PK,
                                           "mql": build(type_key, cols, start, end, "LIMIT 50 OFFSET %d" % offset),
                                           "session_id": ""})
        if res is None:
            FAILS.append("%s~%s OFFSET=%d" % (start or "-", end or "now", offset))
            break
        batch = R.parse_items(res)
        if batch and any(not b for b in batch):
            time.sleep(1.5)
            res2 = R.mcp_call("search_by_mql", {"project_key": PK,
                                                "mql": build(type_key, cols, start, end, "LIMIT 50 OFFSET %d" % offset),
                                                "session_id": ""})
            b2 = R.parse_items(res2) if res2 else []
            if b2 and not any(not b for b in b2):
                batch = b2
            else:
                STUBS[0] += sum(1 for b in batch if not b)
        if not batch:
            empty += 1
            if empty >= 2:
                break
        else:
            empty = 0
            for it in batch:
                if not it:
                    continue
                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
        time.sleep(0.2)
    return items


def midpoint(s, e):
    s_d = datetime.strptime(s or "2024-04-01", "%Y-%m-%d")
    e_d = datetime.strptime(e or TODAY.isoformat(), "%Y-%m-%d")
    mid = (s_d + (e_d - s_d) / 2).strftime("%Y-%m-%d")
    return mid


def seg(type_key, cols, s, e, depth=0):
    n = count_range(type_key, cols, s, e)
    if n and n > 900 and depth < 5:
        mid = midpoint(s, e)
        if mid not in (s, e):
            print("    ↩︎ 段 %s~%s 共 %d 行 → 切分 %s" % (s or "-", e or "now", n, mid))
            return seg(type_key, cols, s, mid, depth + 1) + seg(type_key, cols, mid, e, depth + 1)
    got = pull_seg(type_key, cols, s, e)
    print("  段 %s~%s: count=%s 拉到 %d 条" % (s or "-", e or "now", n, len(got)))
    return got


def main():
    which = sys.argv[1] if len(sys.argv) > 1 else ""
    spec = SPECS.get(which)
    if not spec:
        print("用法: python3 pull_adaptive.py customers|activities")
        return 2
    print("=" * 60)
    print("自适应分段拉取 %s @ %s" % (which, NOW))
    print("=" * 60)
    raw = []
    for s, e in spec["chunks"]:
        raw.extend(seg(spec["type_key"], spec["cols"], s, e))
    seen, items = set(), []
    for it in raw:
        if not it:
            continue
        wid = str(it.get("work_item_id", "") or "")
        if wid and wid in seen:
            continue
        if wid:
            seen.add(wid)
        items.append(it)
    k, rev = spec["sort"]
    items.sort(key=k, reverse=rev)
    print("  合计 %d 条（按 work_item_id 去重）" % len(items))

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

    if FAILS:
        print("  🔴 有分段拉取失败: %s → 不覆盖正式文件，保留 %s" % (FAILS, tmp))
        return 1
    if STUBS[0]:
        print("  🔴 含 %d 条空行 → 不覆盖，保留 %s" % (STUBS[0], tmp))
        return 1
    if len(items) < spec["min_ok"]:
        print("  🔴 仅 %d 条（基准 ≥%d）→ 不覆盖，保留 %s" % (len(items), spec["min_ok"], tmp))
        return 1
    os.replace(tmp, out)
    print("  ✅ 写入 %s: %d 条" % (out, len(items)))
    return 0


if __name__ == "__main__":
    sys.exit(main())