"""Compute cross_refs (top-3 similar items) using stored embeddings.

Usage:
    python -m pulse.brain.index.cross_ref_builder          # all tables
    python -m pulse.brain.index.cross_ref_builder --table facts
    python -m pulse.brain.index.cross_ref_builder --top-k 5
"""
from __future__ import annotations
import json, logging, time
from typing import Optional

import numpy as np

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

KNOWLEDGE_TABLES = (
    "lessons", "failures", "patterns", "facts",
    "evidence", "intents", "decisions", "private",
)
TOP_K = 3
CHUNK_SIZE = 500  # rows scored at a time against the full matrix
EMBED_DIM = 384


def _load_all_embeddings(conn) -> tuple[list[str], list[str], np.ndarray]:
    """Load all embeddings across tables. Returns (ids, table_names, matrix)."""
    all_ids, all_tables, all_vecs = [], [], []
    for table in KNOWLEDGE_TABLES:
        try:
            rows = conn.execute(
                f"SELECT id, embedding FROM {table} WHERE embedding IS NOT NULL"
            ).fetchall()
            for row in rows:
                blob = row[1]
                if blob and len(blob) == EMBED_DIM * 4:
                    vec = np.frombuffer(blob, dtype=np.float32)
                    if np.any(vec != 0):  # skip zero vectors
                        all_ids.append(row[0])
                        all_tables.append(table)
                        all_vecs.append(vec)
        except Exception as e:
            log.warning("Failed to load embeddings from %s: %s", table, e)
    if not all_vecs:
        return [], [], np.array([])
    matrix = np.stack(all_vecs)
    # L2 normalize for cosine similarity via dot product
    norms = np.linalg.norm(matrix, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    matrix = matrix / norms
    return all_ids, all_tables, matrix


def build_cross_refs(table: Optional[str] = None, top_k: int = TOP_K) -> list[dict]:
    """Compute cross-refs for all items with embeddings."""
    from pulse.core.brain_db import get_conn
    conn = get_conn()

    log.info("Loading all embeddings into memory...")
    all_ids, all_tables, matrix = _load_all_embeddings(conn)
    n = len(all_ids)
    if n == 0:
        log.warning("No embeddings found. Run embedding_populator first.")
        return []
    log.info("Loaded %d embeddings, matrix shape: %s", n, matrix.shape)

    # Build index: table -> list of (position, id)
    table_positions = {}
    for i, (item_id, tbl) in enumerate(zip(all_ids, all_tables)):
        table_positions.setdefault(tbl, []).append((i, item_id))

    tables_to_process = [table] if table else list(KNOWLEDGE_TABLES)
    results = []
    total_start = time.time()

    for tbl in tables_to_process:
        positions = table_positions.get(tbl, [])
        if not positions:
            results.append({"table": tbl, "updated": 0, "skipped": True})
            continue

        log.info("%s: computing cross-refs for %d items", tbl, len(positions))
        updated = 0
        start = time.time()

        for chunk_start in range(0, len(positions), CHUNK_SIZE):
            chunk = positions[chunk_start:chunk_start + CHUNK_SIZE]
            chunk_indices = [p[0] for p in chunk]
            chunk_ids = [p[1] for p in chunk]

            # Batch cosine similarity: chunk_vecs @ full_matrix.T
            chunk_matrix = matrix[chunk_indices]
            similarities = chunk_matrix @ matrix.T  # (chunk_size, n)

            updates = []
            for local_i, (global_i, item_id) in enumerate(chunk):
                sims = similarities[local_i]
                # Zero out self-similarity
                sims[global_i] = -1.0
                # Get top-k indices
                top_indices = np.argpartition(sims, -top_k)[-top_k:]
                top_indices = top_indices[np.argsort(sims[top_indices])[::-1]]
                # Filter: only keep similarity > 0.3
                refs = []
                for idx in top_indices:
                    if sims[idx] > 0.3:
                        refs.append({
                            "id": all_ids[idx],
                            "table": all_tables[idx],
                            "sim": round(float(sims[idx]), 3),
                        })
                if refs:
                    updates.append((json.dumps(refs), item_id))

            if updates:
                with conn:
                    conn.executemany(
                        f"UPDATE {tbl} SET cross_refs = ? WHERE id = ?",
                        updates)
                updated += len(updates)

            if chunk_start % 2000 == 0 and chunk_start > 0:
                elapsed = time.time() - start
                log.info("%s: %d/%d chunks processed (%.1fs)",
                         tbl, chunk_start, len(positions), elapsed)

        elapsed = time.time() - start
        stats = {"table": tbl, "updated": updated, "total": len(positions),
                 "elapsed_s": round(elapsed, 1)}
        results.append(stats)
        log.info("%s: %d cross-refs written in %.1fs", tbl, updated, elapsed)

    total_elapsed = time.time() - total_start
    log.info("ALL DONE: cross-refs built in %.1fs", total_elapsed)
    return results


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

    parser = argparse.ArgumentParser(description="Build cross-references")
    parser.add_argument("--table", choices=KNOWLEDGE_TABLES)
    parser.add_argument("--top-k", type=int, default=TOP_K)
    args = parser.parse_args()

    results = build_cross_refs(table=args.table, top_k=args.top_k)
    for r in results:
        print(r)
