"""Step 1: 从钉钉群拉取消息

用法:
    python step1_collect_messages.py              # 增量（最近3天）
    python step1_collect_messages.py --days 1     # 增量（最近1天）
    python step1_collect_messages.py --full       # 全量拉取
    python step1_collect_messages.py --since "2026-06-05 00:00:00"
"""

import argparse
import json
import os
import subprocess
import sys
from datetime import datetime, timedelta

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from paths import DATA_DIR, RAW_DIR, find_dws

DWS = find_dws()

def fetch_messages(group_id, time_str, forward="true", limit=200):
    cmd = [DWS, "chat", "message", "list", "--group", group_id,
           "--time", time_str, "--forward", forward, "--limit", str(limit), "--format", "json"]
    try:
        result = subprocess.run(cmd, capture_output=True, timeout=60)
        output = result.stdout.decode("utf-8", errors="replace")
        json_start = output.find("{")
        if json_start < 0: return []
        data = json.loads(output[json_start:])
        return data.get("result", {}).get("messages", [])
    except Exception as e:
        print(f"  WARN: {e}", file=sys.stderr)
        return []

def collect_group(group_name, group_id, since):
    print(f"\n=== {group_name} (since {since}) ===")
    all_msgs = {}
    msgs = fetch_messages(group_id, since, "true")
    for m in msgs: all_msgs[m["openMessageId"]] = m
    print(f"  正向: {len(msgs)} 条")
    if len(msgs) >= 180:
        times = sorted([m["createTime"] for m in msgs])
        if times:
            msgs2 = fetch_messages(group_id, times[len(times)//2], "false")
            for m in msgs2: all_msgs[m["openMessageId"]] = m
            print(f"  反向补充: {len(msgs2)} 条")
    result = list(all_msgs.values())
    if result:
        times = sorted([m["createTime"] for m in result])
        print(f"  总计: {len(result)} 条 ({times[0]} ~ {times[-1]})")
    return result

GROUPS = {
    "dc战略问题研究院": "cidoUneRB4Db8TAXaTrKxkQAw==",
    "创新组": "cidMuM+itt5PeY7xNSWsv3M0g==",
}

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--days", type=int)
    parser.add_argument("--since", type=str)
    parser.add_argument("--full", action="store_true")
    args = parser.parse_args()

    if args.since:
        since = args.since
    elif args.full:
        since = "2026-05-01 00:00:00"
    elif args.days:
        since = (datetime.now() - timedelta(days=args.days)).strftime("%Y-%m-%d 00:00:00")
    else:
        since = (datetime.now() - timedelta(days=3)).strftime("%Y-%m-%d 00:00:00")

    print(f"dws: {DWS}\n时间: {since} ~ now")

    os.makedirs(RAW_DIR, exist_ok=True)
    combined_path = os.path.join(RAW_DIR, "all_messages_combined.json")
    existing = {}
    if os.path.isfile(combined_path):
        with open(combined_path, "r", encoding="utf-8") as f:
            existing = json.load(f)

    new_count = 0
    for name, gid in GROUPS.items():
        new_msgs = collect_group(name, gid, since)
        if name not in existing: existing[name] = []
        existing_ids = {m["openMessageId"] for m in existing[name]}
        added = sum(1 for m in new_msgs if m["openMessageId"] not in existing_ids)
        for m in new_msgs:
            if m["openMessageId"] not in existing_ids:
                existing[name].append(m)
                existing_ids.add(m["openMessageId"])
        print(f"  新增: {added} 条 (总计: {len(existing[name])})")
        new_count += added

    with open(combined_path, "w", encoding="utf-8") as f:
        json.dump(existing, f, ensure_ascii=False, indent=2)
    total = sum(len(v) for v in existing.values())
    print(f"\n完成! 新增 {new_count} 条，总计 {total} 条 -> {combined_path}")

if __name__ == "__main__":
    main()