"""Batch-populate embedding BLOBs on all knowledge tables.

Usage:
    python -m pulse.brain.index.embedding_populator          # populate all
    python -m pulse.brain.index.embedding_populator --table facts  # single table
    python -m pulse.brain.index.embedding_populator --dry-run     # count only
"""
from __future__ import annotations
import logging, struct, time
from typing import Optional

import numpy as np

log = logging.getLogger("pulse.embedding_populator")

KNOWLEDGE_TABLES = (
    "lessons", "failures", "patterns", "facts",
    "evidence", "intents", "decisions", "private",
)
BATCH_SIZE = 64
EMBED_DIM = 384  # all-MiniLM-L6-v2


def _embed_batch(texts: list[str]) -> list[bytes]:
    """Embed a batch of texts and return BLOB bytes (float32)."""
    from pulse.brain.index.rule_embeddings.backends import embed_texts
    vectors = embed_texts(texts)
    return [np.array(v, dtype=np.float32).tobytes() for v in vectors]


def _count_missing(conn, table: str) -> int:
    return conn.execute(
        f"SELECT COUNT(*) FROM {table} WHERE embedding IS NULL AND id IS NOT NULL"
    ).fetchone()[0]


def populate_table(table: str, batch_size: int = BATCH_SIZE,
                   dry_run: bool = False) -> dict:
    """Populate embeddings for one table. Returns stats dict."""
    from pulse.core.brain_db import get_conn
    conn = get_conn()

    total = conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]
    missing = _count_missing(conn, table)
    if missing == 0:
        return {"table": table, "total": total, "embedded": 0, "skipped": total}

    if dry_run:
        return {"table": table, "total": total, "missing": missing, "dry_run": True}

    log.info("%s: %d/%d missing embeddings", table, missing, total)
    embedded = 0
    errors = 0
    start = time.time()

    while True:
        rows = conn.execute(
            f"SELECT id, content, title FROM {table} "
            f"WHERE embedding IS NULL AND id IS NOT NULL LIMIT ?",
            (batch_size,)
        ).fetchall()
        if not rows:
            break

        ids = [r[0] for r in rows]
        texts = [(r[1] or r[2] or "").strip() for r in rows]

        # For empty/short texts, use a zero vector
        nonempty_idx = [i for i, t in enumerate(texts) if len(t) >= 5]
        blobs = [b'\x00' * (EMBED_DIM * 4)] * len(texts)  # zero vectors

        if nonempty_idx:
            nonempty_texts = [texts[i] for i in nonempty_idx]
            try:
                nonempty_blobs = _embed_batch(nonempty_texts)
                for j, idx in enumerate(nonempty_idx):
                    blobs[idx] = nonempty_blobs[j]
            except Exception as e:
                log.error("Embedding batch failed (%d items): %s", len(nonempty_texts), e)
                errors += len(nonempty_texts)
                # Mark these with a sentinel so we don't retry forever
                for row_id in ids:
                    conn.execute(
                        f"UPDATE {table} SET embedding = ? WHERE id = ?",
                        (b'\x00' * (EMBED_DIM * 4), row_id))
                conn.commit()
                continue

        with conn:
            for row_id, blob in zip(ids, blobs):
                conn.execute(
                    f"UPDATE {table} SET embedding = ? WHERE id = ?",
                    (blob, row_id))
        embedded += len(ids)

        if embedded % 500 == 0:
            elapsed = time.time() - start
            rate = embedded / elapsed if elapsed > 0 else 0
            remaining = (missing - embedded) / rate if rate > 0 else 0
            log.info("%s: %d/%d (%.0f/s, ~%.0fs remaining)",
                     table, embedded, missing, rate, remaining)

    elapsed = time.time() - start
    stats = {
        "table": table, "total": total, "embedded": embedded,
        "errors": errors, "elapsed_s": round(elapsed, 1),
    }
    log.info("%s: done — %d embedded in %.1fs", table, embedded, elapsed)
    return stats


def populate_all(tables: tuple = KNOWLEDGE_TABLES,
                 batch_size: int = BATCH_SIZE,
                 dry_run: bool = False) -> list[dict]:
    """Populate embeddings across all knowledge tables."""
    results = []
    total_start = time.time()
    for table in tables:
        try:
            stats = populate_table(table, batch_size=batch_size, dry_run=dry_run)
            results.append(stats)
        except Exception as e:
            log.error("Failed to populate %s: %s", table, e)
            results.append({"table": table, "error": str(e)})

    total_elapsed = time.time() - total_start
    total_embedded = sum(r.get("embedded", 0) for r in results)
    log.info("ALL DONE: %d items embedded across %d tables in %.1fs",
             total_embedded, len(tables), total_elapsed)
    return results


def embedding_to_numpy(blob: bytes) -> np.ndarray:
    """Convert a stored BLOB back to numpy array."""
    return np.frombuffer(blob, dtype=np.float32)


def numpy_to_blob(arr: np.ndarray) -> bytes:
    """Convert numpy array to storage BLOB."""
    return arr.astype(np.float32).tobytes()


if __name__ == "__main__":
    import argparse, sys
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")

    parser = argparse.ArgumentParser(description="Populate embedding BLOBs")
    parser.add_argument("--table", choices=KNOWLEDGE_TABLES, help="Single table")
    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    if args.table:
        result = populate_table(args.table, batch_size=args.batch_size, dry_run=args.dry_run)
        print(result)
    else:
        results = populate_all(batch_size=args.batch_size, dry_run=args.dry_run)
        for r in results:
            print(r)
