"""Hermes LanceDB → 织忆 MemoryWeave 迁移脚本 (v2)
用单条 commit，走服务端 BGE 编码，确保向量对齐。
"""
import json, time, requests
from concurrent.futures import ThreadPoolExecutor, as_completed
import lancedb
import pyarrow as pa

HERMES_DB = "/home/muc/.hermes/memory_db/lancedb"
ZHIYI_URL = "http://localhost:7821"
ZHIYI_KEY = "zhiyi-dev-key-2026"
WORKERS = 4

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

def arrow_to_dicts(table: pa.Table) -> list[dict]:
    rows = []
    cols = table.column_names
    for i in range(table.num_rows):
        row = {}
        for c in cols:
            row[c] = table.column(c)[i].as_py()
        rows.append(row)
    return rows

def commit_one(r: dict) -> bool:
    """提交单条记忆到织忆"""
    content = str(r.get("content", r.get("text", "")))
    if not content.strip():
        return False
    try:
        resp = requests.post(
            f"{ZHIYI_URL}/api/v1/commit",
            headers={"X-API-Key": ZHIYI_KEY, "Content-Type": "application/json"},
            json={
                "content": content,
                "category": str(r.get("category", "general")),
                "namespace": "hermes-main",
                "agent_id": "hermes-a06",
            },
            timeout=30,
        )
        return resp.status_code in (200, 201)
    except:
        return False

def main():
    log(f"打开 Hermes LanceDB: {HERMES_DB}")
    db = lancedb.connect(HERMES_DB)
    tables = db.table_names()
    log(f"发现 {len(tables)} 个表: {tables}")

    all_records = []
    for tbl_name in tables:
        tbl = db.open_table(tbl_name)
        rows = arrow_to_dicts(tbl.to_arrow())
        log(f"  {tbl_name}: {len(rows)} 条")
        all_records.extend(rows)

    total = len(all_records)
    log(f"总计 {total} 条，准备迁移到织忆（namespace=hermes-main）")

    if total == 0:
        log("无可迁移数据，退出")
        return

    # 并发迁移
    migrated = 0
    failed = 0
    start = time.time()

    with ThreadPoolExecutor(max_workers=WORKERS) as pool:
        futures = {pool.submit(commit_one, r): i for i, r in enumerate(all_records)}
        for f in as_completed(futures):
            idx = futures[f]
            try:
                if f.result():
                    migrated += 1
                else:
                    failed += 1
            except:
                failed += 1

            if (migrated + failed) % 50 == 0:
                elapsed = time.time() - start
                rate = (migrated + failed) / elapsed
                log(f"  [{migrated+failed}/{total}] ✅{migrated} ❌{failed} ({rate:.1f}/s)")

    # 验证
    log(f"\n迁移完成: ✅{migrated} ❌{failed} / {total} ({time.time()-start:.0f}s)")
    resp = requests.get(f"{ZHIYI_URL}/api/v1/stats", headers={"X-API-Key": ZHIYI_KEY})
    stats = resp.json()
    log(f"织忆当前: {stats['total_memories']} 条记忆")

    # 召回验证
    test = requests.post(
        f"{ZHIYI_URL}/api/v1/recall",
        headers={"X-API-Key": ZHIYI_KEY, "Content-Type": "application/json"},
        json={"query": "牧尘用什么系统", "limit": 3},
    ).json()
    log(f"召回测试: {test['count']} 条, top={test['results'][0]['content'][:60] if test['results'] else 'N/A'}")
    log("✅ 迁移完成")

if __name__ == "__main__":
    main()