#!/usr/bin/env python3
"""
Migrate embeddings from standalone QMD DB → OpenClaw QMD DB.
No re-embedding. Pure SQLite BLOB copy.

Rollback (if needed):
  sqlite3 <openclaw_db> "DROP TABLE IF EXISTS vectors_vec; DELETE FROM content_vectors;"
"""

import sqlite3
import os
import time

VEC0_EXT   = "/home/linuxbrew/.linuxbrew/lib/node_modules/qmd/node_modules/sqlite-vec-linux-x64/vec0.so"
SRC_DB     = os.path.expanduser("~/.cache/qmd/index.sqlite")
DST_DB     = "/home/aria/.openclaw/agents/main/qmd/xdg-cache/qmd/index.sqlite"

def open_db(path):
    conn = sqlite3.connect(path)
    conn.enable_load_extension(True)
    conn.load_extension(VEC0_EXT)
    conn.enable_load_extension(False)  # lock down after load
    return conn

def fmt(n):
    return f"{n:,}"

print("=" * 60)
print("QMD Embedding Migration")
print("=" * 60)

# ── Step 1: Connect ──────────────────────────────────────────
print("\n[1/6] Connecting to databases...")
src = open_db(SRC_DB)
dst = open_db(DST_DB)
print(f"  SRC: {SRC_DB}")
print(f"  DST: {DST_DB}")

# ── Step 2: Pre-flight checks ────────────────────────────────
print("\n[2/6] Pre-flight checks...")
src_vec_count = src.execute("SELECT COUNT(*) FROM content_vectors").fetchone()[0]
dst_content   = dst.execute("SELECT COUNT(*) FROM content").fetchone()[0]
dst_vec_count = dst.execute("SELECT COUNT(*) FROM content_vectors").fetchone()[0]

print(f"  SRC vectors:       {fmt(src_vec_count)}")
print(f"  DST content rows:  {fmt(dst_content)}")
print(f"  DST existing vecs: {fmt(dst_vec_count)}")

if dst_vec_count > 0:
    print(f"\n  ⚠️  DST already has {fmt(dst_vec_count)} vectors. Proceeding with INSERT OR IGNORE.")

# ── Step 3: Create vectors_vec in DST ───────────────────────
print("\n[3/6] Creating vectors_vec virtual table in DST...")
dst.execute("""
    CREATE VIRTUAL TABLE IF NOT EXISTS vectors_vec
    USING vec0(hash_seq TEXT PRIMARY KEY, embedding float[768] distance_metric=cosine)
""")
dst.commit()
print("  ✓ vectors_vec ready")

# ── Step 4: Find eligible hashes ────────────────────────────
print("\n[4/6] Finding eligible hashes (must exist in DST content)...")
dst_hashes = set(r[0] for r in dst.execute("SELECT hash FROM content"))
print(f"  DST content hashes: {fmt(len(dst_hashes))}")

src_cvecs = src.execute(
    "SELECT hash, seq, pos, model, embedded_at FROM content_vectors"
).fetchall()
eligible_cvecs = [r for r in src_cvecs if r[0] in dst_hashes]
print(f"  SRC content_vectors total:    {fmt(len(src_cvecs))}")
print(f"  Eligible (hash overlap):      {fmt(len(eligible_cvecs))}")
print(f"  Skipped (no match in DST):    {fmt(len(src_cvecs) - len(eligible_cvecs))}")

# ── Step 5: Insert content_vectors metadata ──────────────────
print("\n[5/6] Inserting content_vectors metadata...")
inserted_cv = 0
skipped_cv  = 0
BATCH = 500
for i in range(0, len(eligible_cvecs), BATCH):
    batch = eligible_cvecs[i:i+BATCH]
    cur = dst.executemany(
        "INSERT OR IGNORE INTO content_vectors (hash, seq, pos, model, embedded_at) VALUES (?,?,?,?,?)",
        [(r[0], r[1], r[2], r[3], r[4]) for r in batch]
    )
    inserted_cv += cur.rowcount
    dst.commit()
    pct = min(100, int((i + BATCH) / len(eligible_cvecs) * 100))
    print(f"  {pct:3d}% — {fmt(i+len(batch))}/{fmt(len(eligible_cvecs))} rows", end="\r")

print(f"\n  ✓ content_vectors: {fmt(inserted_cv)} inserted, {fmt(len(eligible_cvecs)-inserted_cv)} already existed")

# ── Step 6: Copy vectors ──────────────────────────────────────
print("\n[6/6] Copying vectors (BLOBs) from vectors_vec...")
eligible_hash_seqs = set(f"{r[0]}_{r[1]}" for r in eligible_cvecs)

print("  Reading source vectors (may take a moment)...")
t0 = time.time()
src_vecs = src.execute("SELECT hash_seq, embedding FROM vectors_vec").fetchall()
eligible_vecs = [(r[0], r[1]) for r in src_vecs if r[0] in eligible_hash_seqs]
print(f"  Read {fmt(len(src_vecs))} total, {fmt(len(eligible_vecs))} eligible  ({time.time()-t0:.1f}s)")

inserted_vv = 0
BATCH = 200
t0 = time.time()
for i in range(0, len(eligible_vecs), BATCH):
    batch = eligible_vecs[i:i+BATCH]
    cur = dst.executemany(
        "INSERT OR IGNORE INTO vectors_vec (hash_seq, embedding) VALUES (?,?)",
        batch
    )
    inserted_vv += cur.rowcount
    dst.commit()
    elapsed = time.time() - t0
    rate    = (i + len(batch)) / elapsed if elapsed > 0 else 0
    eta     = (len(eligible_vecs) - i - len(batch)) / rate if rate > 0 else 0
    pct     = int((i + len(batch)) / len(eligible_vecs) * 100)
    print(f"  {pct:3d}% — {fmt(i+len(batch))}/{fmt(len(eligible_vecs))}  ({rate:.0f}/s  ETA {eta:.0f}s)   ", end="\r")

print(f"\n  ✓ vectors_vec: {fmt(inserted_vv)} inserted")

# ── Final verification ───────────────────────────────────────
print("\n" + "=" * 60)
print("Verification")
print("=" * 60)
final_cv  = dst.execute("SELECT COUNT(*) FROM content_vectors").fetchone()[0]
final_vv  = dst.execute("SELECT COUNT(*) FROM vectors_vec").fetchone()[0]
model_cnt = dst.execute("SELECT model, COUNT(*) FROM content_vectors GROUP BY model").fetchall()

print(f"  content_vectors rows: {fmt(final_cv)}")
print(f"  vectors_vec entries:  {fmt(final_vv)}")
for m, c in model_cnt:
    print(f"  model '{m}': {fmt(c)} entries")

src.close()
dst.close()

print("\n✅ Migration complete!")
print(f"\nRollback if needed:")
print(f'  sqlite3 "{DST_DB}" "DROP TABLE IF EXISTS vectors_vec; DELETE FROM content_vectors;"')
