#!/usr/bin/env python3
"""
QMD Batch Embedder — bypasses qmd embed's memory-hungry chunking phase.
Reads un-embedded chunks directly from SQLite, embeds via Ollama in small
batches, writes results back. Skips raw session JSONL (too large).
"""

import sqlite3, struct, datetime, urllib.request, json, time, sys, os

DB_PATH      = os.path.expanduser("~/.cache/qmd/index.sqlite")
VEC_EXT      = "/home/linuxbrew/.linuxbrew/lib/node_modules/qmd/node_modules/sqlite-vec-linux-x64/vec0.so"
OLLAMA_URL   = "http://localhost:11434/api/embeddings"
OLLAMA_MODEL = "nomic-embed-text"
MODEL_NAME   = "embeddinggemma"   # match existing content_vectors entries
BATCH_SIZE   = 10                 # conservative — keep memory low
MAX_CHARS    = 4000               # ~1k tokens — safe limit for nomic-embed-text (dense markdown tokenizes heavily)

# Skip raw session JSONL — too large (avg 1.57MB), not useful for memory search
SKIP_COLLECTIONS = {"sessions", "sessions-main"}  # too large/dense for Ollama context

def ollama_embed(text: str) -> list[float]:
    payload = json.dumps({"model": OLLAMA_MODEL, "prompt": text[:MAX_CHARS]}).encode()
    req = urllib.request.Request(OLLAMA_URL, data=payload,
                                  headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=60) as r:
        return json.loads(r.read())["embedding"]

def floats_to_blob(floats: list[float]) -> bytes:
    return struct.pack(f"{len(floats)}f", *floats)

def get_pending(conn) -> list[tuple]:
    rows = conn.execute("""
        SELECT c.hash, c.doc, d.collection
        FROM content c
        JOIN documents d ON c.hash = d.hash
        WHERE c.hash NOT IN (SELECT DISTINCT hash FROM content_vectors)
        GROUP BY c.hash
        ORDER BY
          CASE d.collection
            WHEN 'memory-root' THEN 1
            WHEN 'memory-dir'  THEN 2
            WHEN 'workspace'   THEN 3
            WHEN 'sessions-main' THEN 4
            ELSE 5
          END, c.rowid
    """).fetchall()
    # Filter out skipped collections
    return [(h, doc, col) for h, doc, col in rows if col not in SKIP_COLLECTIONS]

def store(conn, hash_val: str, embedding: list[float]):
    now = datetime.datetime.now(datetime.timezone.utc).isoformat()
    conn.execute("""
        INSERT OR IGNORE INTO content_vectors (hash, seq, pos, model, embedded_at)
        VALUES (?, 0, 0, ?, ?)
    """, (hash_val, MODEL_NAME, now))
    conn.execute("""
        INSERT OR IGNORE INTO vectors_vec (hash_seq, embedding)
        VALUES (?, ?)
    """, (f"{hash_val}_0", floats_to_blob(embedding)))

def main():
    conn = sqlite3.connect(DB_PATH)
    conn.enable_load_extension(True)
    conn.load_extension(VEC_EXT)
    conn.enable_load_extension(False)

    pending = get_pending(conn)
    total   = len(pending)
    done    = 0
    errors  = 0
    skipped = 0

    # Show breakdown
    by_col = {}
    for _, _, col in pending:
        by_col[col] = by_col.get(col, 0) + 1
    print(f"[qmd-batch-embed] {total} chunks to embed:")
    for col, cnt in sorted(by_col.items()):
        print(f"  {col}: {cnt}")
    sys.stdout.flush()

    for i in range(0, total, BATCH_SIZE):
        batch = pending[i : i + BATCH_SIZE]
        for hash_val, doc, col in batch:
            try:
                embedding = ollama_embed(doc)
                store(conn, hash_val, embedding)
                done += 1
            except Exception as e:
                errors += 1
                print(f"  WARN [{col}] {hash_val[:12]}... — {e}")

        conn.commit()
        pct = (done / total * 100) if total else 0
        print(f"  {done}/{total} ({pct:.1f}%) — errors: {errors}", flush=True)
        time.sleep(0.1)

    conn.close()
    print(f"\n[qmd-batch-embed] Done. {done} embedded, {errors} errors, {skipped} skipped.")

if __name__ == "__main__":
    main()
