"""OpenClaw LanceDB → 织忆 MemoryWeave 迁移脚本"""
import json, time, requests
import lancedb

OPENCLAW_DB = "/home/muc/.openclaw/memory-lancedb"
ZHIYI_URL = "http://localhost:7821"
ZHIYI_KEY = "zhiyi-dev-key-2026"

def log(msg):
    print(f"[migrate-openclaw] {msg}", flush=True)

def main():
    log(f"打开 OpenClaw LanceDB: {OPENCLAW_DB}")
    db = lancedb.connect(OPENCLAW_DB)
    table = db.open_table("memories")
    total = table.count_rows()
    log(f"共 {total} 条记忆")

    # 用 to_arrow() 避免 pandas 依赖
    arrow_table = table.to_arrow()
    cols = arrow_table.column_names

    success, fail = 0, 0
    for i in range(arrow_table.num_rows):
        row = {c: arrow_table.column(c)[i].as_py() for c in cols}
        content = str(row.get("text", "")).strip()
        if not content:
            fail += 1
            continue

        category = str(row.get("category", "general"))
        scope = str(row.get("scope", "openclaw"))

        try:
            resp = requests.post(
                f"{ZHIYI_URL}/api/v1/commit",
                headers={"X-API-Key": ZHIYI_KEY, "Content-Type": "application/json"},
                json={
                    "content": content,
                    "category": category,
                    "namespace": "openclaw-main",
                    "agent_id": "openclaw-snowy",
                    "metadata": json.dumps({
                        "scope": scope,
                        "openclaw_id": str(row.get("id", "")),
                        "importance": row.get("importance", 0.5),
                        "timestamp": row.get("timestamp", 0),
                    }),
                },
                timeout=30,
            )
            if resp.status_code in (200, 201):
                success += 1
            else:
                fail += 1
                if resp.status_code == 429:
                    time.sleep(1)
                    # retry once
                    resp = requests.post(
                        f"{ZHIYI_URL}/api/v1/commit",
                        headers={"X-API-Key": ZHIYI_KEY, "Content-Type": "application/json"},
                        json={
                            "content": content,
                            "category": category,
                            "namespace": "openclaw-main",
                            "agent_id": "openclaw-snowy",
                        },
                        timeout=30,
                    )
                    if resp.status_code in (200, 201):
                        success += 1
                        fail -= 1
                    else:
                        print(f"  FAIL[{resp.status_code}]: {content[:80]}...")
                else:
                    print(f"  FAIL[{resp.status_code}]: {content[:80]}...")
        except Exception as e:
            fail += 1
            print(f"  ERR: {content[:80]}... -> {e}")

        if (success + fail) % 10 == 0:
            log(f"进度: {success + fail}/{total} (成功={success}, 失败={fail})")
        time.sleep(0.05)  # 温和限流

    log(f"完成: {total} 条, 成功={success}, 失败={fail}")

if __name__ == "__main__":
    main()