#!/usr/bin/env python3
"""恢复拉取：客户 + 活动，按 start_time 日期分段（每段 <1000 行）避开 MCP OFFSET≤1000 上限。
2026-09-09 发现：MCP search_by_mql 的 LIMIT offset 上限为 1000，OFFSET 翻页无法拉 >1000 行的表。
修复：按 start_time 日期范围切片，每片 OFFSET < 1000，逐片拉取后按 work_item_id 全局去重。
铁律：写临时文件 → 校验条数 vs 首屏总数 → 完整才替换。
"""
import json, os, re, sys, time, urllib.request

MCP_URL = "https://project.feishu.cn/mcp_server/v1"
MCP_TOKEN = "m-abfb29e8-3104-434f-9944-8d0bb592f8cd"
DATA_DIR = "/Users/liuxinyuan/Desktop/Hermes输出-工作类/数据"
SALES_PK = "6593cd71471290e3cc6be6e6"
TYPE_CUSTOMER = "65ae1e403c87b152f3365ca6"
TYPE_ACTIVITY = "65ae1e5d44338dbe7c39a29a"

# 客户分段（2024-05 起，最大片 740 行）
CUSTOMER_CHUNKS = [
    (None, "2024-06-01"),            # < 2024-06-01  → 740
    ("2024-06-01", "2024-07-01"),    # 299
    ("2024-07-01", "2025-01-01"),    # 168
    ("2025-01-01", "2026-01-01"),    # 301
    ("2026-01-01", None),            # >= 2026-01-01 → 219
]
# 活动分段（>= 2026-07-01，最大片 935 行）
ACTIVITY_CHUNKS = [
    ("2026-07-01", "2026-07-11"),
    ("2026-07-11", "2026-07-21"),
    ("2026-07-21", "2026-08-01"),
    ("2026-08-01", "2026-08-11"),
    ("2026-08-11", "2026-08-21"),
    ("2026-08-21", "2026-09-01"),
    ("2026-09-01", None),
]


def mcp_call(method, args, timeout=120):
    payload = {"jsonrpc": "2.0", "id": 1, "method": "tools/call",
               "params": {"name": method, "arguments": args}}
    for attempt in range(5):
        try:
            req = urllib.request.Request(
                MCP_URL, data=json.dumps(payload).encode("utf-8"),
                headers={"X-Mcp-Token": MCP_TOKEN, "Content-Type": "application/json"})
            with urllib.request.urlopen(req, timeout=timeout) as r:
                raw = r.read().decode("utf-8")
        except Exception as e:
            if attempt == 4:
                print(f"    ❌ HTTP 5次失败: {e}")
                return None
            w = 5 * (2 ** attempt)
            print(f"    ⚠️ 请求异常({attempt+1}/5): {e}，{w}s 后重试")
            time.sleep(w)
            continue
        if raw.startswith("data:"):
            raw = raw.split("\n")[0].replace("data: ", "", 1)
        try:
            d = json.loads(raw)
        except Exception:
            if attempt == 4:
                print(f"    ❌ JSON解析失败: {raw[:200]}")
                return None
            time.sleep(8)
            continue
        if "error" in d:
            print(f"    ❌ MCP error: {json.dumps(d['error'], ensure_ascii=False)[:150]}")
            if attempt == 4:
                return None
            time.sleep(12)
            continue
        if "result" not in d or "content" not in d["result"]:
            print(f"    ❌ 意外响应: {json.dumps(d, ensure_ascii=False)[:200]}")
            if attempt == 4:
                return None
            time.sleep(12)
            continue
        for c in d["result"]["content"]:
            t = c.get("text", "")
            if not t or "log_id" in t:
                continue
            cleaned = re.sub(r"\nlog_id:.*$", "", t.strip())
            try:
                return json.loads(cleaned)
            except Exception:
                print(f"    ⚠️ content text 非JSON: {t[:120]}")
                break
        if attempt == 4:
            return None
        time.sleep(15)
    return None


def parse_items(result):
    items = []
    if not result:
        return items
    for gid, gitems in result.get("data", {}).items():
        if not isinstance(gitems, list):
            continue
        for item in gitems:
            fields = {}
            for f in (item.get("moql_field_list") or []):
                k = f["key"]
                v = f.get("value")
                if isinstance(v, list) and len(v) > 0:
                    if isinstance(v[0], dict) and "label" in v[0]:
                        fields[k] = v[0]["label"]
                        continue
                    v = v[0]
                if v is None:
                    fields[k] = ""
                elif isinstance(v, dict):
                    if "string_value" in v:
                        fields[k] = v["string_value"]
                    elif "double_value" in v:
                        fields[k] = v["double_value"]
                    elif "long_value" in v:
                        fields[k] = v["long_value"]
                    elif "key_label_value" in v:
                        fields[k] = v["key_label_value"]
                    elif "user_value" in v:
                        fields[k] = v["user_value"]
                    elif "user_value_list" in v:
                        fields[k] = v["user_value_list"]
                    elif "key_label_value_list" in v:
                        fields[k] = v["key_label_value_list"]
                    else:
                        fields[k] = str(v)
                else:
                    fields[k] = str(v)
            items.append(fields)
    return items


def pull_chunk(type_key, cols, start, end):
    """按日期范围拉取一段（<1000 行，OFFSET 安全），返回 items 列表。"""
    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 = 0
    empty_streak = 0
    for page in range(1, 40):  # 40 页硬上限 = 2000 行，远超单段上限
        pmql = f"{mql} LIMIT 50 OFFSET {offset}"
        result = mcp_call("search_by_mql",
                          {"project_key": SALES_PK, "mql": pmql, "session_id": ""})
        if result is None:
            print(f"    ❌ 段 {start or '-'}~{end or '-'} 第{page}页(OFFSET={offset}) 失败")
            break
        batch = parse_items(result)
        if len(batch) == 0:
            empty_streak += 1
            if empty_streak >= 2:
                break
        else:
            empty_streak = 0
            for it in batch:
                wid = it.get("work_item_id", "")
                if wid and wid in seen:
                    continue
                if wid:
                    seen.add(wid)
                items.append(it)
        offset += 50
        time.sleep(0.3)
    return items


def run_customers():
    cols = ("name, field_c3224b, field_5ed7ab, field_c8e80d, "
            "field_17186c, field_6415cf, work_item_id")
    all_items = []
    seen = set()
    for start, end in CUSTOMER_CHUNKS:
        items = pull_chunk(TYPE_CUSTOMER, cols, start, end)
        print(f"  段 {start or '-inf'}~{end or '+inf'}: {len(items)} 条")
        for it in items:
            wid = it.get("work_item_id", "")
            if wid and wid in seen:
                continue
            if wid:
                seen.add(wid)
            all_items.append(it)
    all_items.sort(key=lambda x: x.get("name", ""))
    return all_items


def run_activities():
    cols = "name, start_time, field_b99055, field_5f20fc, owner, work_item_id"
    all_items = []
    seen = set()
    for start, end in ACTIVITY_CHUNKS:
        items = pull_chunk(TYPE_ACTIVITY, cols, start, end)
        print(f"  段 {start or '-inf'}~{end or '+inf'}: {len(items)} 条")
        for it in items:
            wid = it.get("work_item_id", "")
            if wid and wid in seen:
                continue
            if wid:
                seen.add(wid)
            all_items.append(it)
    all_items.sort(key=lambda x: x.get("start_time", ""), reverse=True)
    return all_items


def main():
    which = sys.argv[1] if len(sys.argv) > 1 else "both"
    now = time.strftime("%Y-%m-%d %H:%M:%S")

    if which in ("customers", "both"):
        print("=" * 60)
        print("恢复拉取 客户（按 start_time 分段）")
        print("=" * 60)
        items = run_customers()
        tmp = os.path.join(DATA_DIR, "mcp_customers.json.tmp")
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump({"updated_at": now, "customers": items}, f, ensure_ascii=False, indent=2)
        if len(items) < 1600:
            print(f"  🔴 客户仅 {len(items)} 条（应 ~1727），不替换，保留临时文件！")
        else:
            os.replace(tmp, os.path.join(DATA_DIR, "mcp_customers.json"))
            print(f"  ✅ 客户写入: {len(items)} 条")

    if which in ("activities", "both"):
        print("=" * 60)
        print("恢复拉取 活动（按 start_time 分段）")
        print("=" * 60)
        items = run_activities()
        tmp = os.path.join(DATA_DIR, "daily_activities_mcp.json.tmp")
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump({"updated": now, "total": len(items), "items": items},
                      f, ensure_ascii=False, indent=2)
        if len(items) < 4000:
            print(f"  🔴 活动仅 {len(items)} 条（应 ~4493），不替换，保留临时文件！")
        else:
            os.replace(tmp, os.path.join(DATA_DIR, "daily_activities_mcp.json"))
            print(f"  ✅ 活动写入: {len(items)} 条")

    print("\n恢复拉取结束。")


if __name__ == "__main__":
    main()